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) GetObjects. 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 canceled then it raises (clean abort, no orphan jobs).

Inside each job, replica_entrypoint.main (replica_entrypoint.py:38):

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

Inside diloco_train (the trainer_fn):

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):

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).
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.
diloco_train trainer_fn (~40 LOC): wraps the existing trainer in make_diloco_outer_loop, accepts manager/rank/world_size.
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.