97 lines
2.7 KiB
Rust
97 lines
2.7 KiB
Rust
//! 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"
|
|
}
|
|
}
|