DiLoCo: Nesterov on parameter server

crates/messages/src/lib.rs

@@ -370,7 +370,11 @@ pub enum Executor {
         config: DiLoCoConfig,
     },
     #[serde(rename = "parameter-server")]
-    ParameterServer { updates: Receive, results: Send },
+    ParameterServer {
+        updates: Receive,
+        results: Send,
+        optimizer: Optimizer,
+    },
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -381,9 +385,9 @@ pub enum Optimizer {
         betas: Option<[f64; 2]>,
         epsilon: Option<f64>,
     },
-    Sgd {
+    Nesterov {
         learning_rate: f64,
-        momentum: Option<f64>,
+        momentum: f64,
     },
 }
 

crates/scheduler/src/bin/hypha-scheduler.rs

@@ -267,6 +267,10 @@ async fn run(config: ConfigWithMetadata<Config>) -> Result<()> {
                 executor: Executor::ParameterServer {
                     updates: Receive::peers(worker_ids.clone()),
                     results: Send::peers(worker_ids.clone(), SelectionStrategy::All),
+                    optimizer: Optimizer::Nesterov {
+                        learning_rate: 0.7,
+                        momentum: 0.9,
+                    },
                 },
             })
             .await

crates/worker/src/executor/mod.rs

@@ -10,6 +10,8 @@ mod process;
 pub use parameter_server::ParameterServerExecutor;
 pub use process::ProcessExecutor;
 
+use crate::executor::parameter_server::TensorOpError;
+
 #[derive(Error, Debug)]
 pub enum Error {
     #[error("Bridge error")]
@@ -19,6 +21,10 @@ pub enum Error {
     Io(#[from] std::io::Error),
     #[error("Unsupported job spec")]
     UnsupportedJobSpec(),
+    #[error("Unsupported optimizer")]
+    UnsupportedOptimizer(),
+    #[error("Tensor error")]
+    Tensor(#[from] TensorOpError),
 }
 
 pub trait JobExecutor {