|
| 1 | +//! The math1 family's Rust: the `FedAvg` kernel, and the storage protocol its twins and triggers speak. |
| 2 | +//! |
| 3 | +//! A trigger stores the canonical input JSON as [`INPUT_FILENAME`] in its own bucket and relays the frame |
| 4 | +//! [`pointer_frame`] builds, naming it. Each twin waits for that frame, reads the input from storage, runs |
| 5 | +//! [`fed_avg`], and stores what [`output_json`] renders as [`OUTPUT_FILENAME`] in its own bucket. The kernel uses only |
| 6 | +//! `+ - * /` on f64 in a fixed order, so every runtime that follows IEEE 754 computes the same model to the bit, the |
| 7 | +//! twins in other languages included. |
| 8 | +//! |
| 9 | +//! `no_std` with `alloc`, so a twin built for a target with no standard library uses it as a browser or WASI one does. |
| 10 | +
|
| 11 | +#![no_std] |
| 12 | +#![expect(clippy::float_arithmetic, reason = "the FedAvg kernel is float math by design")] |
| 13 | + |
| 14 | +extern crate alloc; |
| 15 | + |
| 16 | +use alloc::string::{String, ToString as _}; |
| 17 | +use alloc::vec::Vec; |
| 18 | + |
| 19 | +use serde::Deserialize; |
| 20 | +use serde_json::{Value, json}; |
| 21 | + |
| 22 | +/// The `type` of the pointer frame a trigger relays. |
| 23 | +pub const INPUT_TYPE: &str = "math1-input"; |
| 24 | +/// Storage object name of the canonical input, inside the trigger's bucket. |
| 25 | +pub const INPUT_FILENAME: &str = "math1-input.json"; |
| 26 | +/// Storage object name each twin stores its global model as, inside its own bucket. |
| 27 | +pub const OUTPUT_FILENAME: &str = "math1-output.json"; |
| 28 | + |
| 29 | +/// The canonical input: per-client (feature, target) samples plus the training hyperparameters. |
| 30 | +#[derive(Debug, Deserialize)] |
| 31 | +#[non_exhaustive] |
| 32 | +pub struct Math1Input { |
| 33 | + /// Each client's (feature, target) samples. |
| 34 | + pub clients: Vec<Vec<(f64, f64)>>, |
| 35 | + /// How many rounds of local training and merging to run. |
| 36 | + pub rounds: u32, |
| 37 | + /// How many full-batch gradient-descent epochs each client runs per round. |
| 38 | + pub epochs: u32, |
| 39 | + /// The gradient-descent step size. |
| 40 | + pub learning_rate: f64, |
| 41 | +} |
| 42 | + |
| 43 | +/// Where a trigger stored the input: the bucket and file a pointer frame names. |
| 44 | +#[derive(Clone, Debug, Deserialize)] |
| 45 | +#[non_exhaustive] |
| 46 | +pub struct InputPointer { |
| 47 | + /// The trigger's bucket. |
| 48 | + pub bucket: String, |
| 49 | + /// The input's name in that bucket. |
| 50 | + pub filename: String, |
| 51 | +} |
| 52 | + |
| 53 | +/// Why a text frame is not the math1 input pointer. |
| 54 | +#[derive(Debug, thiserror::Error)] |
| 55 | +#[non_exhaustive] |
| 56 | +pub enum PointerError { |
| 57 | + /// The frame is not JSON, or claims to be the pointer but does not name a bucket and file. |
| 58 | + #[error(transparent)] |
| 59 | + Json(#[from] serde_json::Error), |
| 60 | + /// The frame is JSON, but a message other than the pointer. |
| 61 | + #[error("not a math1-input frame")] |
| 62 | + OtherType, |
| 63 | +} |
| 64 | + |
| 65 | +/// A global model, fitting `target = weight * feature + bias`. |
| 66 | +#[derive(Clone, Copy, Debug, PartialEq)] |
| 67 | +#[non_exhaustive] |
| 68 | +pub struct Model { |
| 69 | + /// The fitted weight. |
| 70 | + pub weight: f64, |
| 71 | + /// The fitted bias. |
| 72 | + pub bias: f64, |
| 73 | +} |
| 74 | + |
| 75 | +/// The pointer frame a trigger relays, naming [`INPUT_FILENAME`] in `bucket`. |
| 76 | +#[must_use] |
| 77 | +pub fn pointer_frame(bucket: &str) -> String { |
| 78 | + json!({ "type": INPUT_TYPE, "bucket": bucket, "filename": INPUT_FILENAME }).to_string() |
| 79 | +} |
| 80 | + |
| 81 | +/// The input pointer the text frame `text` carries. |
| 82 | +/// |
| 83 | +/// # Errors |
| 84 | +/// |
| 85 | +/// When `text` is not JSON, is a message other than the pointer, or is a pointer naming no bucket and file. |
| 86 | +pub fn parse_pointer(text: &str) -> Result<InputPointer, PointerError> { |
| 87 | + let value: Value = serde_json::from_str(text)?; |
| 88 | + if value.get("type").and_then(Value::as_str) != Some(INPUT_TYPE) { |
| 89 | + return Err(PointerError::OtherType); |
| 90 | + } |
| 91 | + Ok(serde_json::from_value(value)?) |
| 92 | +} |
| 93 | + |
| 94 | +/// Sample count as f64, accumulated additively to avoid an integer-to-float cast. |
| 95 | +fn sample_count(samples: &[(f64, f64)]) -> f64 { |
| 96 | + samples.iter().fold(0.0_f64, |count, _| count + 1.0_f64) |
| 97 | +} |
| 98 | + |
| 99 | +/// Run the `FedAvg` simulation on `input`, returning the final global model. |
| 100 | +/// |
| 101 | +/// Each round, every client starts from the global model, runs its local epochs of full-batch gradient descent, and |
| 102 | +/// the results are merged as an average weighted by each client's sample count. |
| 103 | +#[must_use] |
| 104 | +pub fn fed_avg(input: &Math1Input) -> Model { |
| 105 | + let mut weight = 0.0_f64; |
| 106 | + let mut bias = 0.0_f64; |
| 107 | + let total_samples: f64 = input |
| 108 | + .clients |
| 109 | + .iter() |
| 110 | + .fold(0.0_f64, |acc, samples| acc + sample_count(samples)); |
| 111 | + for _ in 0_u32..input.rounds { |
| 112 | + let mut merged_weight = 0.0_f64; |
| 113 | + let mut merged_bias = 0.0_f64; |
| 114 | + for samples in &input.clients { |
| 115 | + let count = sample_count(samples); |
| 116 | + let mut client_weight = weight; |
| 117 | + let mut client_bias = bias; |
| 118 | + for _ in 0_u32..input.epochs { |
| 119 | + let mut grad_weight = 0.0_f64; |
| 120 | + let mut grad_bias = 0.0_f64; |
| 121 | + for &(feature, target) in samples { |
| 122 | + let residual = client_weight * feature + client_bias - target; |
| 123 | + grad_weight += residual * feature; |
| 124 | + grad_bias += residual; |
| 125 | + } |
| 126 | + client_weight -= input.learning_rate * (2.0_f64 * grad_weight / count); |
| 127 | + client_bias -= input.learning_rate * (2.0_f64 * grad_bias / count); |
| 128 | + } |
| 129 | + merged_weight += client_weight * count; |
| 130 | + merged_bias += client_bias * count; |
| 131 | + } |
| 132 | + weight = merged_weight / total_samples; |
| 133 | + bias = merged_bias / total_samples; |
| 134 | + } |
| 135 | + Model { weight, bias } |
| 136 | +} |
| 137 | + |
| 138 | +/// The JSON a twin stores as [`OUTPUT_FILENAME`]: the name of the `module` that computed `model`, and the model. |
| 139 | +#[must_use] |
| 140 | +pub fn output_json(module: &str, model: Model) -> String { |
| 141 | + json!({ "module": module, "weight": model.weight, "bias": model.bias }).to_string() |
| 142 | +} |
0 commit comments