Initial commit

crates/training/rtx-federated/src/aggregation/fednova.rs

@@ -0,0 +1,96 @@
+//! FedNova implementation with normalized averaging for non-IID data
+
+use super::{AggregationAlgorithm, AggregationConfig, ModelUpdate};
+use crate::error::{FederatedError, Result};
+use async_trait::async_trait;
+use serde::{Deserialize, Serialize};
+use tracing::debug;
+use uuid::Uuid;
+
+/// FedNova aggregation algorithm for normalized averaging
+#[derive(Debug)]
+pub struct FedNova {
+    config: FedNovaConfig,
+    tau_effective: f64,
+    round_count: usize,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct FedNovaConfig {
+    pub tau_effective: f64,
+    pub momentum_factor: f64,
+    pub normalize_weights: bool,
+}
+
+impl FedNova {
+    pub async fn new(tau_eff: f64, momentum: f64) -> Result<Self> {
+        let config = FedNovaConfig {
+            tau_effective: tau_eff,
+            momentum_factor: momentum,
+            normalize_weights: true,
+        };
+
+        Ok(Self {
+            config,
+            tau_effective: tau_eff,
+            round_count: 0,
+        })
+    }
+}
+
+#[async_trait]
+impl AggregationAlgorithm for FedNova {
+    async fn aggregate(&self, updates: &[ModelUpdate]) -> Result<ModelUpdate> {
+        if updates.is_empty() {
+            return Err(FederatedError::NoUpdatesCollected);
+        }
+
+        debug!(
+            "🔄 Starting FedNova aggregation with {} updates",
+            updates.len()
+        );
+
+        // Compute normalized weights based on local epochs
+        let mut weights = Vec::new();
+        for update in updates {
+            let local_steps = update.local_epochs as f64;
+            let normalized_weight = local_steps / self.tau_effective;
+            weights.push(normalized_weight * update.sample_count as f64);
+        }
+
+        let aggregated_params = super::WeightedAggregator::weighted_average(updates, &weights)?;
+
+        let mut aggregated_update = ModelUpdate::new(Uuid::new_v4());
+        aggregated_update.parameters = aggregated_params;
+
+        let total_samples: usize = updates.iter().map(|u| u.sample_count).sum();
+        aggregated_update.sample_count = total_samples;
+        aggregated_update.set_metadata(
+            "aggregation_algorithm",
+            serde_json::Value::String("FedNova".to_string()),
+        );
+
+        Ok(aggregated_update)
+    }
+
+    fn get_config(&self) -> AggregationConfig {
+        AggregationConfig::FedNova {
+            tau_effective: self.config.tau_effective,
+            momentum_factor: self.config.momentum_factor,
+            normalize_weights: self.config.normalize_weights,
+        }
+    }
+
+    async fn update_state(&mut self, round: usize, _updates: &[ModelUpdate]) -> Result<()> {
+        self.round_count = round;
+        Ok(())
+    }
+
+    fn supports_async(&self) -> bool {
+        false
+    }
+
+    fn name(&self) -> &'static str {
+        "FedNova"
+    }
+}