research/design-F4-decoupled-diloco-s3.md

# F4 — Decoupled DiLoCo on Serverless + S3 (AWS-native, concrete)

**Facet:** the headline distributed-training question — how N independent SageMaker Training Jobs (or EKS Indexed-Job pods) each run inner DiLoCo (AdamW H steps) then sync pseudo-gradients ONCE per ~500 steps through S3 via `ObjectStoreAllReduce`, with no cross-job NCCL.

**Environment (LIVE):** account `386931836011`, region `us-west-2`, `Admin`/Isengard. Verified live below:
- Rendezvous-ready bucket exists: `s3://amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d` (matches the `sagemaker`-name pattern that `AmazonSageMakerFullAccess` grants S3 on — load-bearing, see IAM).
- Two ready execution roles: `arn:aws:iam::386931836011:role/service-role/AmazonSageMaker-ExecutionRole-20250725T133247` (use this) and `...-20241223T082691`.
- CPU training quota = **30** instances each for `ml.m5.xlarge` / `ml.m5.large` / `ml.c5.xlarge` in us-west-2 → the 2-replica smoke is trivially in budget.
- Warm-pool quota = **0** for all (default). `KeepAlivePeriodInSeconds` set without a quota increase silently no-ops. This is THE cold-start gotcha.

---

## 1. The decision that already shipped (ADR-005) and why it's right for AWS

ADR-005 chose **object-store rendezvous, not cross-job NCCL**, as the DiLoCo comm primitive. DiLoCo (Douillard et al. 2023, arXiv:2311.08105 §3.2) syncs once per H≈500–1000 inner steps — ~10–30 min wall-clock. The exchange per outer round is one `PutObject` of the pseudo-gradient (~2 GB for a 1B bf16 model) + (N−1) `GetObject`s. At N=8 that's ~128 GB read spread over 30 min ≈ 70 MB/s aggregate, ~$0.05/round on S3. The repo realizes exactly this in `ObjectStoreAllReduce.allreduce()` (`composer_replication/diloco/serverless/allreduce.py:131`): PUT `round_{NNNNNN}/rank_{RRRR}.pt`, poll-until-all-peers-exist, mean, `tensor.copy_(avg)`.

**Why this is correct on AWS specifically (not just plausible):** S3 has been **strongly read-after-write consistent for PUT/GET/LIST in all regions since Dec 2 2020, at no extra cost** (`aws.amazon.com/s3/consistency`; `docs.aws.amazon.com/AmazonS3/latest/userguide/Welcome.html`). The poll loop's correctness depends on exactly this: rank R PUTs `rank_0003.pt`, and the instant any peer's `_exists()` returns True the subsequent `_get()` is guaranteed to return that byte-complete object — no "object not yet visible" race, no eventual-consistency retry shim. The repo's `_put` even writes a `.tmp` then `os.replace` on the local path (atomic on POSIX); on S3 a single `PutObject` is atomic per key by definition, so a peer never sees a half-written `.pt`. **The substrate is consistency-correct by construction on S3.**

The one-line architectural payoff on K8s/serverless (report §8): a straggler replica simply blocks peers at the poll loop (bounded by `timeout_s=1800`) instead of deadlocking an NCCL gang. Inter-replica DiLoCo sync needs **no gang scheduling** — only intra-replica FSDP does.

---

## 2. Exact lifecycle: `SageMakerExecutor.launch_replicas(N)` → collect

`composer_replication/diloco/serverless/sagemaker.py` (`SageMakerExecutor`) already implements the full `ServerlessExecutor` Protocol. The end-to-end flow for one Decoupled DiLoCo run:

```
make_diloco_run(N=4)                          # orchestrator (NEW, ~120 LOC — see §6)
  └─ exec = SageMakerExecutor(
         role_arn="arn:aws:iam::386931836011:role/service-role/AmazonSageMaker-ExecutionRole-20250725T133247",
         image_uri="386931836011.dkr.ecr.us-west-2.amazonaws.com/composer-diloco:latest",
         output_s3_path="s3://amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d/diloco-out/run42/",
         region="us-west-2")
  └─ handles = exec.launch_replicas(
         n_replicas=4, entrypoint=<ignored>,
         entrypoint_args={
           "rendezvous_uri": "s3://amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d/diloco-rdv/run42/",
           "trainer_module": "composer_replication.trainer.composer_trainer",
           "trainer_fn": "diloco_train",            # NEW thin wrapper — see §6
           "trainer_kwargs": {"model_name":"Qwen/Qwen2.5-0.5B","sync_every":500,"total_steps":2000}},
         gpu="A10G", timeout=86400)
```

**Per replica, `launch_replicas` submits ONE single-instance Training Job** (`sagemaker.py:309-369`):
- `ResourceConfig.InstanceCount == 1` — deliberately NOT one multi-instance job (which would wire SageMaker's intra-job NCCL via `resourceconfig.json` and couple the replicas — the wrong model). N replicas = N separate jobs.
- `AlgorithmSpecification.ContainerEntrypoint = ["python","-m","composer_replication.diloco.serverless.replica_entrypoint"]`, `ContainerArguments = ["--rendezvous", s3uri, "--world-size","4","--trainer-module",...,"--trainer-fn","diloco_train","--trainer-kwargs-json", "{...}"]` (each token a separate list element).
- `Environment = {"REPLICA_RANK":"<rank>","WORLD_SIZE":"4","RENDEZVOUS_URI":s3uri}` — the rank channel.
- `EnableNetworkIsolation=False` — **load-bearing**, pinned, never a knob: True severs the container's outbound S3 and dead-locks the allreduce poll until `timeout_s`. The bucket access is granted on `RoleArn` instead (SageMaker's IRSA analog).
- On any rank's `create_training_job` failure, already-launched siblings are best-effort `cancel`ed then it raises (clean abort, no orphan jobs).

**Inside each job, `replica_entrypoint.main`** (`replica_entrypoint.py:38`):
1. reads `REPLICA_RANK` from env (falls back to argv via the dual contract — argv for SageMaker/Local, env for EKS),
2. builds `store = ObjectStoreAllReduce(uri=s3uri, rank, world_size)`. For an `s3://` URI this hits `_init_fsspec()` → `fsspec.filesystem("s3")` (s3fs; in the `[serverless]` extra),
3. `manager = MockManager(store)`,
4. imports `trainer_module.trainer_fn`, injects `manager=`, `rank=`, `world_size=`, calls it.

**Inside `diloco_train` (the trainer_fn):**
```python
diloco = make_diloco_outer_loop(
    manager=manager, model_fragments=[model],
    inner_optimizer=AdamW(model.parameters(), lr=1e-5),
    outer_lr=0.7, outer_momentum=0.9, nesterov=True, sync_every=500)
with diloco:
    for step in range(total_steps):
        inner_optim.zero_grad(); loss = composer_loss(...); loss.backward()
        inner_optim.step()       # at step%500==0 DiLoCo's post-hook fires
```
At every 500th `inner_optim.step()`, torchft's DiLoCo post-hook fires `prepare_sync` → `perform_sync`. `perform_sync` computes the pseudo-gradient `θ_initial − θ_local` (sign convention pinned in `diloco/__init__.py:13-38`), calls `manager.allreduce(pseudograd)` → `MockManager.allreduce` (`allreduce.py:268`) → `ObjectStoreAllReduce.allreduce`: PUT `round_000000/rank_0003.pt`, poll for all 4 ranks' files, mean, `copy_` back. `MockManager.should_commit()` always returns True (no FT failover; replica failure is the orchestrator's job), then the **outer Nesterov SGD step** applies the averaged pseudo-gradient and redistributes the new global weights. `start_quorum()` bumps `_step` so `current_step()` advances exactly once per round (fragment-rotation math). Repeat for the next 500 inner steps → `round_000001/`, etc.

**Collect:** `exec.collect(handles, timeout=...)` polls `describe_training_job` per handle until terminal (`Completed`/`Failed`/`Stopped`), returns rank-ordered result dicts incl. `S3ModelArtifacts`. `poll` maps `TrainingJobStatus` via `_STATUS_MAP` (refining `InProgress`+queued → `pending`). `stream_logs` reads `/aws/sagemaker/TrainingJobs` CloudWatch stream `<job>/algo-1-<epoch>`.

The DiLoCo math, `MockManager`, `ObjectStoreAllReduce`, `make_diloco_outer_loop`, the trainer, and the loss are **all byte-for-byte unchanged** across Local→EKS→SageMaker. That is the whole point of the Protocol.

---

## 3. S3 rendezvous specifics for AWS

- **Bucket/prefix:** `s3://amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d/diloco-rdv/<run_id>/` for the rendezvous; a sibling `…/diloco-out/<run_id>/` for `OutputDataConfig`. Use the sagemaker-named bucket so the stock execution-role policy already grants access (see IAM). Layout written by the substrate: `…/diloco-rdv/<run_id>/round_{NNNNNN}/rank_{RRRR}.pt`.
- **Consistency:** strong RAW + strong LIST, all regions, since 2020 — the poll loop (`_exists` then `_get`) needs no consistency shim. A peer's file becomes visible atomically on PUT completion.
- **IAM (load-bearing):** the jobs need `s3:GetObject` + `s3:PutObject` (and ideally `s3:ListBucket`) on `…/diloco-rdv/<run_id>/*` on the **execution `RoleArn`**, NOT the caller. The stock `AmazonSageMakerFullAccess` on the existing roles grants S3 only on buckets whose name contains `sagemaker`/`Sagemaker`/`SageMaker`/`aws-glue` — which is exactly why the `amazon-sagemaker-…` bucket works out of the box and a custom-named `s3://composer-diloco-rdv` would silently 403 the first PUT and hang every peer until `timeout_s`. Either keep the rendezvous in a sagemaker-named bucket (recommended, zero IAM work) or attach an inline policy scoping `s3:GetObject/PutObject/ListBucket` to the custom prefix.
- **Poll timeout vs stragglers:** `ObjectStoreAllReduce(timeout_s=1800, poll_interval_s=1.0)`. 30 min comfortably covers a SageMaker cold-started replica that joins late (3–5 min provision + first-500-steps lag). For Spot churn at larger N, raise to `timeout_s=3600`. A `TimeoutError` names the exact missing `rank_R` + `round_N` — the orchestrator can then cancel + relaunch that rank (DiLoCo's `should_commit==True` means a stalled round does not silently skip; the run aborts cleanly rather than averaging a partial set).
- **Cost:** ~$0.05/round (ADR-005), negligible vs GPU. For a 2000-step / sync_every=500 run that's 4 rounds ≈ $0.20 of S3 for the whole run.

---

## 4. What's missing to run it for real (the deferred smoke) + the cheapest validating run

**The gap.** ADR-005 §"Open/deferred" flags a "real serverless smoke" as never run; report §10 Phase 0 lists "EKSExecutor + S3 rendezvous + dep bump" as substrate hardening. Concretely, what exists vs what's missing:

| Layer | State |
|---|---|
| `ObjectStoreAllReduce` over `file://` | **Proven** — `test_serverless_diloco_integration.py` runs a 2-process `LocalProcessExecutor` run and asserts cross-rank weight convergence after one outer round. |
| `ObjectStoreAllReduce` over `s3://` | **Code path exists, never exercised against real S3.** `_init_fsspec()` is untested with live s3fs; the `_exists`/`_get`/`_put` S3 branches have only mock coverage. |
| `SageMakerExecutor` against real boto3 | **Never submitted a real job.** Tests inject a `_MockSMClient`. The ECR `image_uri` does not yet exist (no Dockerfile baking `composer_replication`). |
| `diloco_train` trainer_fn | **Missing.** `composer_trainer.py` has the trainer but no thin DiLoCo-wrapped entry that accepts the injected `manager`/`rank`/`world_size` kwargs. |

**The exact remaining gaps to close, in order (each cheap, each decisive):**
1. **s3:// smoke (≈$0, ~10 min):** point the *existing* `test_serverless_diloco_integration` multi-process test at `s3://amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d/diloco-rdv/smoke-$(uuid)/` instead of a tmp dir, with 2 local processes. This exercises the real s3fs PUT/poll/GET/mean path and the strong-consistency assumption with zero GPU spend. Only `s3fs` + `boto3` (already in `[serverless]`) and ambient `Admin` creds are needed. Gate: both ranks converge to identical weights (the existing assertion).
2. **Dockerfile + ECR push (~$0):** ~30-line image `FROM pytorch/pytorch`, `pip install -e .[serverless,diloco,train]`, entrypoint `python -m composer_replication.diloco.serverless.replica_entrypoint`. Push to `386931836011.dkr.ecr.us-west-2.amazonaws.com/composer-diloco:latest`.
3. **`diloco_train` trainer_fn (~40 LOC):** wraps the existing trainer in `make_diloco_outer_loop`, accepts `manager/rank/world_size`.
4. **2 tiny CPU SageMaker jobs (the validating run, ~$0.10–0.30):** `SageMakerExecutor(gpu=None → ml.m5.xlarge, $0.23/hr on-demand)`, `n_replicas=2`, `nn.Linear` or a 0.5B model with `total_steps=4, sync_every=2`, `rendezvous_uri` in the sagemaker bucket. Two jobs × ~5 min cold-start + ~2 min run ≈ 0.24 instance-hours ≈ **$0.06–0.30**. Within the 30-instance CPU quota with massive headroom. Gate: `collect()` returns both `succeeded`, both ranks' `round_000000/rank_000{0,1}.pt` appear in S3, and the two model artifacts in `…/diloco-out/` are byte-identical (proves the cross-job allreduce ran through real S3, not just locally).

That ladder — `file://` (done) → `s3://` local (step 1) → 2 real CPU jobs (step 4) — is the report's prescribed cheapest path and closes the deferred-smoke gap for well under the ADR's $2–5 estimate (the original estimate assumed GPU + Modal).

---

## 5. Streaming DiLoCo (`fragment_sync_delay>0`) on this substrate

Streaming DiLoCo (Liu et al. 2025, "Eager Updates for Overlapped Communication in DiLoCo", arXiv:2501.18512; the DiLoCo line is arXiv:2311.08105) splits the model into fragments synced on staggered schedules so the allreduce of fragment k overlaps inner computation of the next steps. `make_diloco_outer_loop` already exposes the knobs: `fragment_sync_delay>0`, `fragment_update_alpha`, and `model_fragments=[frag_0,...,frag_M-1]` (`diloco/__init__.py:72-93`).

From torchft's `local_sgd.py` (`_StreamingDiLoCoFragment`, verified via DeepWiki against the upstream source):
- **prepare_sync** fires at inner step `sync_every − fragment_sync_delay`: computes pseudo-gradients (`_save_grads`) and **launches the allreduce without waiting**, recording the `Work` in `self._allreduce_work`, on a separate CUDA stream.
- **perform_sync** fires at `sync_every`: **waits** on that `Work`, restores global params, `should_commit()`, applies the outer step, then `_merge_parameters` blends `local*alpha + global*(1−alpha)` (`fragment_update_alpha`; 0.0 = standard full-replacement DiLoCo).
- **`_current_fragment() = current_step() % len(fragments)`** — round-robin; each fragment syncs on its own offset. This is exactly why `MockManager.start_quorum()` bumps `_step` once per round and `current_step()` is faithful: get this wrong and replicas pick different fragments and diverge.

**The load-bearing gotcha for the object-store substrate.** Streaming's overlap depends on `manager.allreduce()` being **asynchronous** — prepare_sync launches it, `fragment_sync_delay` steps of inner compute run, then perform_sync waits. But the repo's `MockManager.allreduce` is **synchronous**: it calls `ObjectStoreAllReduce.allreduce`, which **blocks** on the poll-until-all-peers loop before returning `_ImmediateWork` (whose `.wait()` is a no-op). So on this substrate, **today, prepare_sync blocks for the full S3 rendezvous and `fragment_sync_delay` buys zero overlap** — Streaming degrades to vanilla, correctly but without the comm/compute overlap benefit. This is fine for correctness (and is why the same API "configures Streaming" per the docstring) but defeats the point on a 2 GB-per-fragment, 30-min-cadence S3 sync.

**The fix to make Streaming real here (~60 LOC, deferred):** give `ObjectStoreAllReduce` a non-blocking mode: `allreduce_async(tensor)` PUTs `round_N/rank_R.pt` and returns immediately; the returned `Work.wait()` then runs the poll-GET-mean-copy. `MockManager.allreduce` returns that deferred `Work` instead of `_ImmediateWork`. Now prepare_sync's PUT returns instantly, the `fragment_sync_delay` inner steps run while peers are PUTting concurrently, and perform_sync's `.wait()` does the poll/mean. Because S3 is strongly consistent, by the time perform_sync waits `fragment_sync_delay` steps later, peer files are far more likely already present — the overlap is genuine. Per-fragment streaming further shrinks each PUT to (model_size / M), so the poll is over smaller objects. This is the natural Streaming-DiLoCo realization on object storage and is the right Phase-5 upgrade (report §10) for multi-fragment large-model runs; vanilla (`fragment_sync_delay=0`, single fragment) is correct and sufficient for Phases 0–4.

---

## 6. Repo delta (what to build in `composer_replication/`)

| File | Delta | ~LOC |
|---|---|---|
| `composer_replication/trainer/composer_trainer.py` | NEW `diloco_train(*, manager, rank, world_size, model_name, sync_every=500, total_steps, **kw)` — wraps the existing trainer body in `make_diloco_outer_loop(manager=manager,...)`; this is the `trainer_fn` `replica_entrypoint` calls. | ~40 |
| `composer_replication/diloco/serverless/run.py` | NEW thin orchestrator `make_diloco_run(executor, n, rendezvous_uri, trainer_module, trainer_fn, trainer_kwargs, gpu, timeout)` → `launch_replicas` + `collect`; surfaces straggler `TimeoutError` → relaunch-rank. | ~120 |
| `docker/Dockerfile.diloco` | NEW — `FROM pytorch/pytorch:*-cuda*`, `pip install -e .[serverless,diloco,train]`, ENTRYPOINT to `replica_entrypoint`. Push to ECR `composer-diloco:latest`. | ~30 |
| `composer_replication/diloco/serverless/sagemaker.py` | EXTEND: optional `keep_alive_period_s` → `ResourceConfig.KeepAlivePeriodInSeconds` (warm pool; **document the 0-default quota gotcha**); optional `use_spot` → `EnableManagedSpotTraining=True` + `MaxWaitTimeInSeconds` (> `MaxRuntimeInSeconds`) + `CheckpointConfig`. Both default off. | ~40 |
| `composer_replication/diloco/serverless/allreduce.py` | EXTEND (Phase 5, for real Streaming): `allreduce_async` + a deferred `Work` whose `.wait()` runs the poll/mean; `MockManager` returns it. Vanilla path untouched. | ~60 |
| `composer_replication/diloco/serverless/tests/test_serverless_diloco_integration.py` | EXTEND: parametrize rendezvous over `file://` AND `s3://amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d/diloco-rdv/smoke-<uuid>/` (gated on `AWS_SMOKE=1`) — closes the s3:// path gap with zero GPU. | ~30 |
| `docs/adrs/ADR-005-serverless-diloco.md` | UPDATE the "Open/deferred — Real serverless smoke" clause: replace the Modal $2–5 estimate with the SageMaker 2×CPU-job ($0.06–0.30) plan above. | ~10 |

Untouched: `loss.py`, `teacher_replay.py`, `safety/`, `make_diloco_outer_loop`, `MockManager` core, `ObjectStoreAllReduce` core, `replica_entrypoint` (its dual argv/env contract already supports both SageMaker and EKS).

---

## 7. Open questions / falsifiers

- **Warm-pool quota:** default 0 in us-west-2 (verified). If iterative dev wants warm starts (skip the 3–5 min cold-start per round-of-jobs), request a `ml.<type> for training warm pool usage` quota increase first; otherwise `KeepAlivePeriodInSeconds` no-ops.
- **Pseudo-gradient dtype/size at real model scale:** the smoke uses tiny tensors; a 0.5B–8B bf16 pseudo-gradient is 1–16 GB per PUT — confirm s3fs multipart upload throughput and that `torch.save({"rank","tensor"})` round-trips bf16 on CPU (the code casts peer tensors back to device+dtype on GET).
- **Streaming overlap:** only real after the `allreduce_async` delta (§5); until then `fragment_sync_delay>0` is correct-but-no-overlap on S3. Measure round wall-clock with vs without before claiming the benefit.
- **N>16 straggler fragility under Spot:** the bounded `timeout_s` poll is the mitigation; the HyperPod-attached-node-group path (report §9, same S3 rendezvous) is the resilience escalation for multi-day runs.