Update on "[WIP] Add DiLoCo"

Still WIP but open to feedback on the API

## API Usage

```python
# LocalSGD example
model = SimpleModel()
optimizer = optim.SGD(model.parameters())
manager = create_autospec(Manager)
with LocalSGD(manager, model, optimizer, sync_every=2):
    for inp, label in dataloader:
        loss = model(inp).mean()
        loss.backward()
        optimizer.step()
        
# DiLoCo example
model = SimpleModel()
inner_optimizer = torch.optim.AdamW(
    m.parameters(), lr=4e-4, weight_decay=0.1, betas=(0.9, 0.95)
)
outer_optimizer = torch.optim.SGD(
    m.parameters(), lr=0.7, momentum=0.9, nesterov=True
)
manager = create_autospec(Manager)
with DiLoCo(manager, m, inner_optimizer, outer_optimizer, sync_every=2):
    for inp, label in dataloader:
        loss = model(inp).mean()
        loss.backward()
        inner_optimizer.step()
        # outer_optimizer is used every 'sync_every' steps

```

## Changes
- Updated `LocalSGD` to be a context manager rather than a `nn.Module` wrapper. This required adding a pre_forward_hook to the model start the quorum
- Added DiLoCo. This is a subclass of LocalSGD since a lot of code is shared
- TODO: should be working, but still validating some tests

discussion doc: https://docs.google.com/document/d/11c5JwQpSzilrDvK-vNsgQhpXAihbMn-hTRC8y3LiGqY/edit?tab=t.0#heading=h.izo4yi6jz4mk

[ghstack-poisoned]

Author: H-Huang

.pyre_configuration

@@ -6,5 +6,9 @@
       "import_root": ".",
       "source": "torchft"
     }
+  ],
+  "search_path": [
+    {"site-package": "torchx"},
+    {"site-package": "parameterized"}
   ]
 }

.torchxconfig

@@ -0,0 +1,3 @@
+[cli:run]
+component=torchft/torchx.py:hsdp
+scheduler=local_cwd

CONTRIBUTING.md

@@ -96,6 +96,21 @@ make livehtml
 The docs will be built in the `docs/build/html` directory and served at http://localhost:8000.
 The page will be automatically re-built as long as the process is kept running.
 
+### Running Multiple Replica Local Job
+
+We use torchx to run multiple worker local test jobs. You need to run the
+lighthouse first and then you can use torchx to launch as many replica groups as
+you want. This uses the [train_ddp.py](./train_ddp.py) script.
+
+```sh
+$ torchft_lighthouse --min_replicas 2 --join_timeout_ms 10000 &
+$ torchx run -- --replicas 10
+```
+
+Once the Lighthouse has started you can view the status of all the workers at the Lighthouse dashboard.
+
+Default address is: http://localhost:29510
+
 ## Contributor License Agreement ("CLA")
 
 In order to accept your pull request, we need you to submit a CLA. You only need to do this once to work on any of

Cargo.toml

@@ -12,6 +12,7 @@ log = "0.4.22"
 prost = "0.13.3"
 prost-types = "0.13.3"
 pyo3 = {version="0.22.3", features = ["extension-module"]}
+rand = "0.8.5"
 slog = "2.7.0"
 slog-stdlog = "4.1.1"
 stderrlog = "0.6.0"

README.md

@@ -52,7 +52,7 @@ greatly improve efficiency by avoiding stop the world training on errors.
 
 Before proceeding, ensure you have the following installed:
 
-- Rust (with necessaray dependencies)
+- Rust (with necessary dependencies)
 - `protobuf-compiler` and the corresponding development package for Protobuf.
 
 Note that the Rust versions available in many conda environments may be outdated. To install the latest version of Rust, we recommend downloading it directly from the official website as shown in the below command:

output.txt

@@ -41,6 +41,7 @@ message QuorumMember {
     string store_address = 3;
     int64 step = 4;
     uint64 world_size = 5;
+    bool shrink_only = 6;
 }
 
 message Quorum {
@@ -72,6 +73,7 @@ message ManagerQuorumRequest {
     int64 rank = 1;
     int64 step = 2;
     string checkpoint_server_addr = 3;
+    bool shrink_only = 4;
 }
 
 message ManagerQuorumResponse {

proto/torchft.proto

@@ -27,7 +27,8 @@ dev = [
     "pyre-check",
     "parameterized",
     "expecttest",
-    "numpy"
+    "numpy",
+    "torchx"
 ]
 
 [tool.maturin]

pyproject.toml

@@ -6,6 +6,9 @@
 
 pub mod lighthouse;
 pub mod manager;
+mod net;
+mod retry;
+mod timeout;
 
 use core::time::Duration;
 use std::env;
@@ -46,6 +49,7 @@ impl Manager {
         store_addr: String,
         world_size: u64,
         heartbeat_interval: Duration,
+        connect_timeout: Duration,
     ) -> PyResult<Self> {
         py.allow_threads(move || {
             let runtime = Runtime::new()?;
@@ -58,6 +62,7 @@ impl Manager {
                     store_addr,
                     world_size,
                     heartbeat_interval,
+                    connect_timeout,
                 ))
                 .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
             let handle = runtime.spawn(manager.clone().run());
@@ -84,47 +89,47 @@ impl Manager {
 struct ManagerClient {
     runtime: Runtime,
     client: ManagerServiceClient<Channel>,
-    timeout: Duration,
 }
 
 #[pymethods]
 impl ManagerClient {
     #[new]
-    fn new(py: Python<'_>, addr: String, timeout: Duration) -> PyResult<Self> {
+    fn new(py: Python<'_>, addr: String, connect_timeout: Duration) -> PyResult<Self> {
         py.allow_threads(move || {
             let runtime = Runtime::new()?;
             let client = runtime
-                .block_on(manager::manager_client_new(addr, timeout))
+                .block_on(manager::manager_client_new(addr, connect_timeout))
                 .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
 
             Ok(Self {
                 runtime: runtime,
                 client: client,
-                timeout: timeout,
             })
         })
     }
 
-    #[pyo3(signature = (rank, step, checkpoint_server_addr, timeout=None))]
     fn quorum(
-        &mut self,
+        &self,
         py: Python<'_>,
         rank: i64,
         step: i64,
         checkpoint_server_addr: String,
-        timeout: Option<Duration>,
+        shrink_only: bool,
+        timeout: Duration,
     ) -> Result<(i64, i64, i64, String, String, i64, Option<i64>, i64, bool), StatusError> {
         py.allow_threads(move || {
             let mut request = tonic::Request::new(ManagerQuorumRequest {
                 rank: rank,
                 step: step,
                 checkpoint_server_addr: checkpoint_server_addr,
+                shrink_only: shrink_only,
             });
-            // This notifies the server about the timeout but doesn't affect the
-            // endpoint timeout which we set on client creation.
-            request.set_timeout(timeout.unwrap_or(self.timeout));
 
-            let response = self.runtime.block_on(self.client.quorum(request))?;
+            // This timeout is processed on the server side so we also enable
+            // keep alives to detect server health.
+            request.set_timeout(timeout);
+
+            let response = self.runtime.block_on(self.client.clone().quorum(request))?;
             let resp = response.into_inner();
             Ok((
                 resp.quorum_id,
@@ -140,47 +145,49 @@ impl ManagerClient {
         })
     }
 
-    #[pyo3(signature = (rank, timeout=None))]
     fn checkpoint_address(
-        &mut self,
+        &self,
         py: Python<'_>,
         rank: i64,
-        timeout: Option<Duration>,
+        timeout: Duration,
     ) -> Result<String, StatusError> {
         py.allow_threads(move || {
             let mut request = tonic::Request::new(CheckpointAddressRequest { rank: rank });
-            // This notifies the server about the timeout but doesn't affect the
-            // endpoint timeout which we set on client creation.
-            request.set_timeout(timeout.unwrap_or(self.timeout));
+
+            // This timeout is processed on the server side so we also enable
+            // keep alives to detect server health.
+            request.set_timeout(timeout);
 
             let response = self
                 .runtime
-                .block_on(self.client.checkpoint_address(request))?;
+                .block_on(self.client.clone().checkpoint_address(request))?;
             let resp = response.into_inner();
             Ok(resp.checkpoint_server_address)
         })
     }
 
-    #[pyo3(signature = (rank, step, should_commit, timeout=None))]
     fn should_commit(
-        &mut self,
+        &self,
         py: Python<'_>,
         rank: i64,
         step: i64,
         should_commit: bool,
-        timeout: Option<Duration>,
+        timeout: Duration,
     ) -> Result<bool, StatusError> {
         py.allow_threads(move || {
             let mut request = tonic::Request::new(ShouldCommitRequest {
                 rank: rank,
                 step: step,
                 should_commit: should_commit,
             });
+
             // This notifies the server about the timeout but doesn't affect the
             // endpoint timeout which we set on client creation.
-            request.set_timeout(timeout.unwrap_or(self.timeout));
+            request.set_timeout(timeout);
 
-            let response = self.runtime.block_on(self.client.should_commit(request))?;
+            let response = self
+                .runtime
+                .block_on(self.client.clone().should_commit(request))?;
             let resp = response.into_inner();
             Ok(resp.should_commit)
         })
@@ -297,11 +304,16 @@ impl From<Status> for StatusError {
 #[pymodule]
 fn torchft(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // setup logging on import
-    stderrlog::new()
-        .verbosity(2)
+    let mut log = stderrlog::new();
+    log.verbosity(2)
         .show_module_names(true)
-        .timestamp(stderrlog::Timestamp::Millisecond)
-        .init()
+        .timestamp(stderrlog::Timestamp::Millisecond);
+
+    if env::var("CLICOLOR_FORCE").is_ok() {
+        log.color(stderrlog::ColorChoice::AlwaysAnsi);
+    }
+
+    log.init()
         .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
 
     m.add_class::<Manager>()?;