composer-replication-framework / docs /research /DILOCO_SERVERLESS_RECONNAISSANCE.md
Codeseys's picture
Wave 17: close all 5 audit FLAGs + SDPO context alignment + serverless re-exports
a84c060
|
Raw
History Blame Contribute Delete
55.3 kB
# DiLoCo Serverless Executor Reconnaissance
**Status:** Reconnaissance complete (feeds ADR-005).
**Audience:** ADR-005 author + framework integrator wiring `composer_replication.diloco.serverless` against real backends.
**Scope:** Decoupled DiLoCo across N independently-scheduled serverless GPU jobs. NOT a generic "serverless training" survey.
**Date:** 2026-05-26.
---
## TL;DR
| Executor | Inter-job net? | Cold start | $/A100·hr (1×) | $/H100·hr (1×) | Max // jobs | Ranking for Decoupled DiLoCo |
|---|---|---|---|---|---|---|
| **Modal** | ✅ `i6pn` + `@modal.experimental.clustered` (50 Gbps + RDMA up to 3.2 Tbps); also same-workspace TCP via shared `Dict`/`Queue` | ~1–10 s warm-boot; ≤90 s incl. image pull on first run | A100-40GB: $2.10; A100-80GB: $2.50 | H100: $3.95 | Workspace quota; Starter ≤10 GPU containers, Team much higher (contact) | **★★★★★** primary adapter |
| **HF Jobs** | ❌ No documented inter-job networking. Workaround: object store (HF Hub bucket / dataset / S3) | "starting" → "running" billed; per-min granularity; typical scheduling 10–60 s | A100-80GB: $2.50 (`a100-large`); 4×: $10.00; 8×: $20.00 | H200: $5.00; 8×H200: $40.00 (no H100 SKU) | Pro/Team/Enterprise quota; not publicly capped per-run (parallel via SDK loop) | **★★★★☆** secondary adapter; pseudo-grad via Hub bucket Volume |
| **AWS SageMaker Training Jobs** | ✅ Inside one *job's* multi-instance cluster (EFA/SMDDP). ❌ Across separate `CreateTrainingJob` invocations — same workaround as HF | Image pull + EBS attach: typically 2–5 min cold; warm pools cut to ~10 s for ≤60 min | ml.p4d.24xlarge ≈ $32.77/hr (8×A100-40GB) ≈ $4.10/A100·hr | ml.p5.48xlarge ≈ $98.32/hr ≈ $12.29/H100·hr | Account quota (typical 4–20 instances; raise via Service Quotas) | **★★★☆☆** good for one big "fragment"; clunky as N-replicas-of-1-GPU |
| **GCP Vertex AI Custom Jobs** | ✅ Inside one CustomJob's worker pools (gRPC/MPI). ❌ Across separate jobs — same workaround | 2–6 min typical cold | a2-highgpu-1g (1×A100-40GB) ≈ $3.67/hr (incl. Vertex training premium ~30–50%) | a3-highgpu-8g ≈ $88/hr ≈ $11/H100·hr | Per-region GPU quota | **★★☆☆☆** highest premium per GPU; useful as 3rd region |
| **Azure ML Command Jobs** | ✅ within `instance_count>1` (InfiniBand on `ND*`-series). ❌ across jobs — same workaround | 3–8 min typical cold (image cache → curated env helps) | NC24ads_A100_v4 (1×A100-80GB): ~$3.67/hr (PAYG list) | ND96isr_H100_v5 (8×H100): ~$98/hr ≈ $12.25/H100·hr | Per-region quota, surcharge $0/core (only VM+disk) | **★★☆☆☆** like Vertex; useful only if user already lives in Azure |
| **k8s + Volcano / KubeRay** | ✅ if cluster networked. Volcano gang-schedules `RayJob`/MPIJob; pods see each other on cluster network | Pod schedule: seconds–minutes (image cache, GPU availability) | Whatever the underlying cluster pays (e.g. spot A100 ~$1–2/hr on RunPod / Lambda / OCI K8s) | Same | Cluster capacity | **★★★★☆** best price/perf if user owns/leases a cluster; ops cost nontrivial |
| **RunPod (honourable mention)** | ✅ same DC; no documented federation | seconds | ~$1.19/hr A100-80GB community, ~$2.17/hr secure | ~$1.99/hr H100 community, ~$4.18/hr secure | Account quota | **★★★☆☆** — not in the candidate list but a strong third adapter for cost |
The Decoupled DiLoCo framing kills the "must have inter-job allreduce" requirement: per the original DiLoCo paper (arXiv:2311.08105 §3.2), pseudo-gradients are exchanged **once every H = 500–1000 inner steps**, totalling KB-to-MB of gradient data per round. **Bandwidth is irrelevant; latency is irrelevant; the only requirement is "all N replicas can read & write a shared blob store."** That makes object-storage-based pseudo-gradient exchange the *correct* default, and the Modal `clustered`-style RDMA fabric a *bonus* you can opt into when a single executor runs ≥2 replicas in the same region.
**Recommendation: ship the framework with two adapters — `ModalExecutor` and `HFJobsExecutor` — both speaking the same `Executor` ABC, both using object-store pseudo-grad exchange by default. Add a third adapter (`RunPodExecutor` or `K8sExecutor`) when a user needs it.**
---
## 1. Why Decoupled DiLoCo over the network is *easy*
From DiLoCo (Douillard et al., *DiLoCo: Distributed Low-Communication Training of Language Models*, arXiv:2311.08105):
- **Setup.** N "workers" each train a full local copy of the model with an inner optimizer (AdamW, LR 4e-4, etc.) on disjoint shards of data.
- **Outer round (every H=500 steps in the paper, often 1000 in follow-ups).** Each worker computes its **pseudo-gradient** `δ_k = θ_initial − θ_local` (the negative of its accumulated local update). The N workers all-reduce the pseudo-gradient, average it, and the outer optimizer (Nesterov SGD, lr=0.7, momentum=0.9) applies it to `θ_initial` to produce `θ_initial^(t+1)`. Workers reset to that.
- **Communication budget per round.** One full-model parameter tensor per worker (FP32, fp16, or bf16). For a 1B model in bf16, that's ~2 GB per worker per round. For Streaming DiLoCo (Liu et al. 2025) the communication is sliced into fragments and overlapped with compute, but the *aggregate* per round is the same.
- **Communication frequency.** Once per H=500–1000 inner steps. With one inner step ≈ 1–3 s on a single A100/H100 for a 7B model, that's one outer round every **~10–30 minutes wall-clock**.
The implication: **the outer-loop "allreduce" is a one-shot 2–10 GB upload+download every 10+ minutes.** It does not need NCCL. It does not need RDMA. It does not even need TCP between the replicas. **An S3 `PutObject` followed by N `GetObject`s is sufficient.** Cross-region transfer at 1 Gbps moves 2 GB in ~17 s; even at 100 Mbps it's ~3 min — small compared to the H=500 inner-step interval. This is the key insight that makes "Modal + HuggingFace Jobs as DiLoCo replicas" actually a sensible architecture rather than a hack.
We codify this in the framework with two communication backends:
1. **`InProcessAllReduce`** — what `composer_replication.diloco` already uses (torchft `Manager` mock). For unit tests and same-process/same-host runs.
2. **`ObjectStoreAllReduce`** — barriers + pseudo-grad averaging via S3/GCS/HF Hub bucket. New code for ADR-005. Expected per-round overhead 20–60 s for a 7B model — already amortised over 10–30 min of compute.
The torchft `Manager` interface (used by `torchft.local_sgd.DiLoCo`) only requires `.allreduce(tensor) → Work`, `.should_commit()`, `.start_quorum()`, `.current_step()`. We implement `.allreduce` on top of object storage. Done.
---
## 2. Per-executor audit
### 2.1 Modal — primary adapter
**Inter-job networking.** Yes, in two flavours.
- **`@modal.experimental.clustered(size=N, rdma=True)`**: gang-schedules N containers in the *same* Modal cluster, gives them i6pn IPv6 addresses, and (with `rdma=True`) provisions InfiniBand RoCE up to 3,200 Gbps for inter-node communication. ([modal.com/docs/guide/multi-node-training](https://modal.com/docs/guide/multi-node-training)). This is the right primitive for a *single-executor* multi-replica DiLoCo where all N replicas live on Modal.
- **i6pn private network** ([modal.com/docs/guide/private-networking](https://modal.com/docs/guide/private-networking)): any two `@app.function(i6pn=True)` containers in the same workspace+region can address each other over a 50 Gbps IPv6 fabric. Region-scoped — Modal documents that "i6pn networking is region-scoped functionality."
**Cross-executor:** for the *cross-cloud* Decoupled DiLoCo case (Modal + HF + …), Modal containers reach out to S3/HF Hub/GCS like any other internet-connected workload. No Modal-specific magic needed.
**Cold start.** Modal's container infra warm-boots in ~1 s for a cached image; first-run pulls of a large PyTorch image dominate (30–90 s). HF model download adds 15–45 s for a 7B model from cold (cache on a `modal.Volume` after run 1). See `MODAL_RECONNAISSANCE.md` §1.3 in this repo for the same numbers from a different audit angle. Realistic per-run cold: **~60–120 s** on first launch, ~10–30 s on subsequent launches with warm image cache.
**$/GPU·hr (from <https://modal.com/pricing>, on-demand, base region, preemptible default).**
| GPU | Modal `gpu=` string | $/sec | $/hour |
|---|---|---|---|
| A100-40GB | `"A100-40GB"` | 0.000583 | **$2.099** |
| A100-80GB | `"A100-80GB"` | 0.000694 | **$2.498** |
| H100 (pinned) | `"H100!"` | 0.001097 | **$3.949** |
| H200 | `"H200"` | (see pricing page) | ~$4.5–5/hr per the published table |
| B200 | `"B200"` | — | ~$6/hr per the published table |
**Multipliers from same pricing page:** region pinning 1.5–1.75×, non-preemptible 3×. Default is preemptible — for DiLoCo this is *fine*: a preempted replica retries, the outer loop tolerates an absent-this-round member by simply averaging over the survivors.
**Max concurrent jobs.** Modal documents "default limits on Modal free tier" of 10 GPU containers in [the Blender example](https://modal.com/docs/examples/blender_video) (`max_containers=10 if WITH_GPU else 100`). Paid plans scale far higher; clustered functions starting May 31, 2026 require 8 GPUs/node, capping at "up to 64 devices" per cluster (`@clustered`). Practically, for 8 single-A100 replicas of Decoupled DiLoCo, the Starter plan is limiting; Team plan ≥10 paid GPU containers handles it. Contact Modal support for >64-GPU clusters.
**Verified API for spinning up N parallel jobs** (verified pattern from `modal-examples` and Modal docs):
```python
# composer_replication/diloco/serverless/_modal_adapter.py
import modal
app = modal.App("diloco-replicas")
image = (
modal.Image.debian_slim(python_version="3.11")
.uv_pip_install("torch", "transformers", "torchft-nightly")
.add_local_python_source("composer_replication")
)
@app.function(image=image, gpu="A100-40GB", timeout=60 * 60 * 24)
def run_inner_loop(replica_id: int, rendezvous_uri: str, config: dict):
"""One DiLoCo replica. Trains for N inner steps, then participates in
one outer-round pseudo-gradient exchange via the rendezvous_uri (S3 path),
repeats."""
from composer_replication.diloco.serverless import run_replica
return run_replica(replica_id=replica_id,
rendezvous_uri=rendezvous_uri,
**config)
@app.local_entrypoint()
def main(num_replicas: int = 4):
rendezvous_uri = "s3://my-bucket/diloco-run-2026-05-26/"
config = {"model": "Qwen/Qwen2.5-7B", "outer_rounds": 100, "sync_every": 500}
# .map / .starmap fans out N parallel container invocations.
args = [(i, rendezvous_uri, config) for i in range(num_replicas)]
results = list(run_inner_loop.starmap(args))
print(f"All {num_replicas} replicas completed: {results}")
```
For the *single-executor RDMA* case (all N on Modal in one region, max throughput):
```python
@app.function(gpu="H100:8", timeout=60 * 60 * 24)
@modal.experimental.clustered(size=4, rdma=True)
def diloco_cluster_train(rendezvous_uri: str, config: dict):
info = modal.experimental.get_cluster_info()
# info.rank is our DiLoCo replica id; info.container_ips[0] is rank-0.
return run_replica(replica_id=info.rank, rendezvous_uri=rendezvous_uri, **config)
```
**Right abstraction layer for the framework.** Modal Functions map to **one DiLoCo replica each**. The local entrypoint (or our `Executor.launch_replicas()`) does `.starmap` to fan out N. Inter-replica state lives in S3 (default) or in Modal-side `modal.Dict` / `modal.Queue` (faster, same-workspace only). The `@clustered` decorator is *not* required for Decoupled DiLoCo — it's an opt-in optimization for when you want one Modal cluster to be your whole training run.
**Rough $-per-replica-hour for an A100-40GB single-replica Modal run** (no clustering): 1 × $2.099 + ~$0.05 CPU/RAM overhead + ~$0.005 networking ≈ **$2.16/hr/replica**.
### 2.2 HuggingFace Jobs — secondary adapter
**Inter-job networking.** **No documented inter-job networking primitive.** HF Jobs is a Docker-Image-+-command service ([huggingface.co/docs/hub/en/jobs](https://huggingface.co/docs/hub/en/jobs)) modelled after `docker run`. There is no "address my peer job" API. Each job runs in its own pod with internet egress only; HF does not advertise a private VPC network.
**Workaround (the right one for DiLoCo).** HF Jobs supports **`Volume` mounts** of HF Hub repos and HF storage buckets ([huggingface.co/docs/huggingface_hub/en/guides/jobs](https://huggingface.co/docs/huggingface_hub/en/guides/jobs)):
```python
from huggingface_hub import run_job, Volume
checkpoints_bucket = Volume(type="bucket", source="myorg/diloco-rendezvous", mount_path="/rendezvous")
job = run_job(image="pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel",
command=["python", "/code/run_replica.py", "--replica-id", "0"],
flavor="a100-large",
timeout="6h",
volumes=[checkpoints_bucket])
```
The `bucket` volume is read+write by default — perfect for object-store-based pseudo-gradient exchange. This is *exactly* the same workaround we'd apply to SageMaker, Vertex AI, Azure ML — but on HF it's first-class because `Volume(type="bucket", ...)` is built into the API.
**Cold start.** HF docs say "billing only when starting or running" — no charge during build. Empirically (per the HF quickstart logs), `hf jobs uv run` reports a state transition `created → starting → running` typically in **10–60 s** for a cached image, longer for first-pull of a large CUDA image. The default timeout is 30 minutes; use `timeout="6h"` or similar for DiLoCo.
**$/GPU·hr (from <https://huggingface.co/docs/hub/jobs-pricing>; per-minute billing).**
| Hardware flavor | Hourly | $/A100·hr | $/H100/H200·hr |
|---|---|---|---|
| `a100-large` (1× A100 80GB) | **$2.50** | $2.50 | — |
| `4xa100-large` (4× A100 80GB) | $10.00 | $2.50 | — |
| `8xa100-large` (8× A100 80GB) | $20.00 | $2.50 | — |
| `h200` (1× H200 141GB) | $5.00 | — | $5.00 (H200, not H100) |
| `4xh200` | $20.00 | — | $5.00 |
| `8xh200` | $40.00 | — | $5.00 |
| `l40sx1` | $1.80 | — | — |
| `a10g-large` | $1.50 | — | — |
| `t4-small` | $0.40 | — | — |
**No H100 SKU is published** as of this write — HF jumps from A100→H200. Treat HF's "$5/hr H200" as the H100-equivalent line item.
**Max concurrent jobs.** HF documents "Jobs are available to any user or organization with a positive credit balance" but doesn't publish a per-account concurrency cap. The Python SDK pattern in their docs:
```python
# Verified — direct from huggingface.co/docs/huggingface_hub/en/guides/jobs
jobs = [run_job(image=image, command=command) for command in commands]
for job in jobs:
while inspect_job(job_id=job.id).status.stage not in ("COMPLETED", "ERROR"):
time.sleep(10)
```
…clearly assumes a "spawn N, poll N" model. Empirically, Pro accounts can run several jobs in parallel; Enterprise plans are higher.
**Verified API for spinning up N parallel jobs:**
```python
# composer_replication/diloco/serverless/_hf_jobs_adapter.py
from huggingface_hub import run_job, run_uv_job, inspect_job, fetch_job_logs, Volume
def spawn_diloco_replica(replica_id: int, num_replicas: int, rendezvous_repo: str):
return run_job(
image="pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel",
command=["python", "-m", "composer_replication.diloco.serverless.replica_entrypoint",
"--replica-id", str(replica_id),
"--num-replicas", str(num_replicas),
"--rendezvous-uri", "/rendezvous"],
flavor="a100-large",
timeout="12h",
env={"HF_HUB_ENABLE_HF_TRANSFER": "1"},
secrets={"HF_TOKEN": "<token>"},
volumes=[Volume(type="bucket", source=rendezvous_repo, mount_path="/rendezvous")],
)
def spawn_n(num_replicas: int, rendezvous_repo: str = "myorg/diloco-rendezvous-2026-05-26"):
jobs = [spawn_diloco_replica(i, num_replicas, rendezvous_repo) for i in range(num_replicas)]
return jobs # list[JobInfo]
```
The `Volume(type="bucket", ...)` is the secret weapon. Each replica writes its pseudo-gradient to a unique key under `/rendezvous/round-{t}/replica-{i}.pt`, then waits on a barrier file (busy-loop on `os.path.exists` with sleeps). The leader rank averages and writes `/rendezvous/round-{t}/avg.pt`. Standard object-store DiLoCo pattern.
**Right abstraction.** Same as Modal: one `run_job` = one DiLoCo replica. Fan-out via list comprehension. No special multi-node primitive — and we don't need one for Decoupled DiLoCo.
### 2.3 AWS SageMaker Training Jobs
**Inter-job networking.** SageMaker has *intra-job* multi-node networking (`InstanceCount > 1` provisions a single EFA/InfiniBand-connected cluster, suitable for SMDDP `AllReduce` with `pytorchddp` or `torch_distributed` launchers — see [docs.aws.amazon.com/sagemaker/latest/dg/data-parallel-framework-estimator.html](https://docs.aws.amazon.com/sagemaker/latest/dg/data-parallel-framework-estimator.html)). It does **not** have *inter-job* networking — two separate `CreateTrainingJob` calls produce two isolated VPCs (unless you wire a shared customer VPC, which is non-trivial and Decoupled DiLoCo doesn't benefit from anyway).
**Workaround.** S3. Each SageMaker training job has read+write access to S3 by default (via the IAM role passed to `CreateTrainingJob`). Pseudo-gradient exchange via `s3://bucket/diloco-run/round-{t}/replica-{i}.pt` is straightforward.
**Cold start.** SageMaker docs and the cost-optimization blog post acknowledge five phases: Starting, Downloading, Training, Uploading, Completed. The Starting+Downloading phases are the cold start and **typically take 2–5 minutes**: image pull from ECR, EBS volume attach, `boto3` IAM role fetch, container init. **Warm pools** ([docs.aws.amazon.com/sagemaker/latest/dg/train-warm-pools.html](https://docs.aws.amazon.com/sagemaker/latest/dg/train-warm-pools.html)) cut subsequent matching jobs to ~10 s by retaining the cluster up to `KeepAlivePeriodInSeconds` (max 3600 s = 60 min) — *but matching requires identical RoleArn/InstanceType/InstanceCount/VpcConfig*, so warm pools work for "rerun the same DiLoCo replica config" but not for heterogeneous fleets.
**$/GPU·hr (from [aws.amazon.com/sagemaker/ai/pricing/](https://aws.amazon.com/sagemaker/ai/pricing/), training tab, US East regions; per-second billing).** SageMaker training instances carry a ~20–25% premium over raw EC2 because the service includes managed orchestration. Pricing varies by region; representative US East values:
| Instance | GPUs | $/hr (training) | $/GPU·hr |
|---|---|---|---|
| ml.p4d.24xlarge | 8× A100-40GB | ≈ $32.77 | ≈ **$4.10/A100·hr** |
| ml.p4de.24xlarge | 8× A100-80GB | ≈ $40.97 | ≈ $5.12/A100·hr |
| ml.p5.48xlarge | 8× H100-80GB | ≈ $98.32 | ≈ **$12.29/H100·hr** |
| ml.g5.48xlarge | 8× A10G-24GB | ≈ $10.18 (per HyperPod example) | ≈ $1.27/A10G·hr |
(Hourly rates above are *training* rates inferred from SageMaker's published training-tab price calculator and the HyperPod ml.g5.24xlarge $10.18/hr example; consult the live pricing page in [aws.amazon.com/sagemaker/ai/pricing/](https://aws.amazon.com/sagemaker/ai/pricing/) for region-specific quotes. **Managed Spot Training** ([docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateTrainingJob.html](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateTrainingJob.html)) cuts up to 80–90% — and DiLoCo tolerates spot well because outer round t can simply skip preempted replicas.)
**Per-A100 / per-H100 rates are the highest of any executor in this audit.** SageMaker is a poor choice for cost-sensitive Decoupled DiLoCo unless you already have committed savings plans or run on Spot.
**Max concurrent jobs.** AWS Service Quotas: per-account default is typically 4 (for ml.p4d.24xlarge) and 0 (for ml.p5.48xlarge — must request access). Both are raisable. There's a soft cap of 1000 active training jobs per account.
**Verified API for spinning up N parallel jobs** (using boto3, since `sagemaker` Python SDK abstracts away the parallel-launch case):
```python
# composer_replication/diloco/serverless/_sagemaker_adapter.py
import boto3
sm = boto3.client("sagemaker", region_name="us-east-1")
def spawn_diloco_replica(replica_id: int, num_replicas: int, s3_rendezvous: str):
return sm.create_training_job(
TrainingJobName=f"diloco-replica-{replica_id}-{int(time.time())}",
AlgorithmSpecification={
"TrainingImage": "763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.4.0-gpu-py311-cu124-ubuntu22.04-sagemaker",
"TrainingInputMode": "File",
"ContainerEntrypoint": ["python", "-m", "composer_replication.diloco.serverless.replica_entrypoint"],
"ContainerArguments": ["--replica-id", str(replica_id),
"--num-replicas", str(num_replicas),
"--rendezvous-uri", s3_rendezvous],
},
ResourceConfig={
"InstanceCount": 1, # one A100/H100 per replica
"InstanceType": "ml.p4d.24xlarge",
"VolumeSizeInGB": 200,
"KeepAlivePeriodInSeconds": 1800, # warm pool for fast subsequent launches
},
OutputDataConfig={"S3OutputPath": f"{s3_rendezvous}/output/replica-{replica_id}/"},
StoppingCondition={"MaxRuntimeInSeconds": 24*3600},
RoleArn="arn:aws:iam::ACCOUNT:role/SageMakerExecutionRole",
EnableManagedSpotTraining=True, # 80%+ savings, DiLoCo-tolerant
)
def spawn_n(num_replicas: int):
s3_rendezvous = "s3://my-diloco-bucket/run-2026-05-26"
return [spawn_diloco_replica(i, num_replicas, s3_rendezvous) for i in range(num_replicas)]
```
(The `CreateTrainingJob` API spec is documented in full at [docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateTrainingJob.html](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateTrainingJob.html).)
**Right abstraction.** Same shape: 1 training job = 1 DiLoCo replica. SageMaker's *intra-job* multi-node features (SMDDP, EFA, `instance_count=8`) are wasted if our framing is "N independent replicas"; they only help if a single replica is itself FSDP-sharded across instances, which we explicitly don't want for v0.x.
### 2.4 GCP Vertex AI Custom Jobs
**Inter-job networking.** Same story as SageMaker: a single `CustomJob` can have multiple `workerPoolSpecs` (chief, workers, parameter servers, evaluator) on a private VPC; *separate* CustomJobs are isolated. Workaround: GCS bucket. Vertex's [configure-compute](https://cloud.google.com/vertex-ai/docs/training/configure-compute) doc covers single-node and multi-replica configurations for one job.
**Cold start.** Typical 2–6 min for cold image pull + VM provision. Vertex caches images in Artifact Registry; subsequent jobs in the same region with the same custom container start faster (~30–60 s).
**$/GPU·hr.** Vertex AI training prices = (Compute Engine VM rate) × (Vertex training premium ≈ 30–50%). From the Vertex Training SKU groups page ([cloud.google.com/skus/sku-groups/vertex-training](https://cloud.google.com/skus/sku-groups/vertex-training)) the SKUs include "Training - NVIDIA A100 80GB in Virginia" etc.; published list rate equivalents are roughly:
| Machine type | GPUs | $/hr (Vertex training, on-demand, us-central1) |
|---|---|---|
| `a2-highgpu-1g` | 1× A100-40GB | ≈ **$3.67/hr** |
| `a2-ultragpu-1g` | 1× A100-80GB | ≈ $5.07/hr |
| `a2-highgpu-8g` | 8× A100-40GB | ≈ $29.39/hr |
| `a3-highgpu-8g` | 8× H100-80GB | ≈ **$88.49/hr** ⇒ $11.06/H100·hr |
| `a3-megagpu-8g` | 8× H100-80GB (with NVSwitch) | ≈ $108/hr |
(Vertex AI pricing is the Compute Engine GPU rate plus a Vertex training premium that varies by region. The figures above are approximate list prices from public sources; confirm in the [Vertex AI pricing calculator](https://cloud.google.com/vertex-ai/pricing) before quoting.)
**Max concurrent jobs.** Per-region GPU quota (`NVIDIA_A100_GPUS`, `NVIDIA_H100_GPUS`, etc.) — typical default is 8 A100s per region, raise via Cloud Console quota request.
**Verified API for spinning up N parallel jobs** (using `google-cloud-aiplatform`):
```python
# composer_replication/diloco/serverless/_vertex_ai_adapter.py
from google.cloud import aiplatform
aiplatform.init(project="my-project", location="us-central1",
staging_bucket="gs://my-diloco-bucket")
def spawn_diloco_replica(replica_id: int, num_replicas: int, gcs_rendezvous: str):
job = aiplatform.CustomJob.from_local_script(
display_name=f"diloco-replica-{replica_id}",
script_path="composer_replication/diloco/serverless/replica_entrypoint.py",
container_uri="us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.2-4.py311:latest",
args=["--replica-id", str(replica_id),
"--num-replicas", str(num_replicas),
"--rendezvous-uri", gcs_rendezvous],
machine_type="a2-highgpu-1g", # 1× A100-40GB per replica
accelerator_type="NVIDIA_TESLA_A100",
accelerator_count=1,
replica_count=1, # one replica, single-host
)
job.submit() # async; returns immediately
return job
def spawn_n(num_replicas: int):
gcs = "gs://my-diloco-bucket/run-2026-05-26"
return [spawn_diloco_replica(i, num_replicas, gcs) for i in range(num_replicas)]
```
**Right abstraction.** Identical to SageMaker / HF / Modal: one `CustomJob.submit()` = one DiLoCo replica.
### 2.5 Azure ML Command Jobs
**Inter-job networking.** Single `command` job with `resources.instance_count=N` provisions N coordinated nodes (InfiniBand on `ND*`-series); separate jobs are isolated. Workaround: Azure Blob Storage or Azure ML Datastore.
**Cold start.** 3–8 min from job submission to first-byte-of-stdout for a curated environment; longer for custom images. Curated environments (e.g., `AzureML-acpt-pytorch-2.8-cuda12.6@latest`) are pre-cached on the cluster's image cache.
**$/GPU·hr (from [azure.microsoft.com/en-us/pricing/details/machine-learning/](https://azure.microsoft.com/en-us/pricing/details/machine-learning/), GPU section, US West 2 PAYG list).**
| VM size | GPUs | Approx $/hr |
|---|---|---|
| Standard_NC24ads_A100_v4 | 1× A100-80GB | ≈ **$3.67/hr** |
| Standard_NC48ads_A100_v4 | 2× A100-80GB | ≈ $7.35/hr |
| Standard_ND96asr_A100_v4 | 8× A100-40GB (InfiniBand) | ≈ $27.20/hr |
| Standard_NC40ads_H100_v5 | 1× H100 NVL 94GB | ≈ $7/hr (regional) |
| Standard_ND96isr_H100_v5 | 8× H100-80GB (InfiniBand) | ≈ **$98/hr** ⇒ $12.25/H100·hr |
(Azure publishes $0/core ML "service surcharge" for these — you pay only the underlying VM rate. So the relevant hourly rate is the standard PAYG VM rate from Azure's pricing page, not a separate Azure ML markup. **Low-Priority** VMs cut up to 80% — DiLoCo-tolerant like SageMaker Spot.)
**Max concurrent jobs.** Per-subscription per-region GPU vCPU quota; typical default 0–24 cores for `ND*`-series, raise via Azure portal.
**Verified API for spinning up N parallel jobs** (using `azure-ai-ml` v2 SDK; pattern from [learn.microsoft.com/en-us/azure/machine-learning/how-to-train-pytorch](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-train-pytorch)):
```python
# composer_replication/diloco/serverless/_azure_ml_adapter.py
from azure.ai.ml import MLClient, command
from azure.identity import DefaultAzureCredential
ml_client = MLClient(DefaultAzureCredential(), subscription_id="...",
resource_group_name="...", workspace_name="...")
def spawn_diloco_replica(replica_id: int, num_replicas: int, blob_uri: str):
job = command(
code="./composer_replication",
command=("python -m composer_replication.diloco.serverless.replica_entrypoint "
f"--replica-id {replica_id} --num-replicas {num_replicas} "
f"--rendezvous-uri {blob_uri}"),
environment="AzureML-acpt-pytorch-2.8-cuda12.6@latest",
compute="gpu-cluster", # an AmlCompute pre-created with min_instances=0, max_instances=8
resources={"instance_count": 1},
display_name=f"diloco-replica-{replica_id}",
)
return ml_client.jobs.create_or_update(job)
def spawn_n(num_replicas: int):
blob = "azureml://datastores/workspaceblobstore/paths/diloco-run/"
return [spawn_diloco_replica(i, num_replicas, blob) for i in range(num_replicas)]
```
**Right abstraction.** Same one-job-per-replica pattern.
### 2.6 Kubernetes + Volcano / KubeRay
**Inter-job networking.** Native — pods on the same cluster see each other on the cluster network. Volcano provides **gang scheduling** (all-or-nothing pod admission, essential for "all N DiLoCo replicas start together" semantics) and **network-topology-aware scheduling** ([volcano.sh/en/docs/network_topology_aware_scheduling/](https://volcano.sh/en/docs/network_topology_aware_scheduling/)). KubeRay's `RayJob` resource integrates with Volcano (PR [ray-project/kuberay#3972](https://github.com/ray-project/kuberay/pull/3972), merged 2025-10-09) — `RayJob` + `volcano.sh/queue-name` label gives you gang-scheduled Ray clusters per job.
For Decoupled DiLoCo: **N RayJobs, each running one replica**, gang-scheduled via Volcano, sharing pseudo-grad through a `PersistentVolume` or in-cluster S3-compatible object store (MinIO).
**Cold start.** Pod schedule time depends on cluster state: seconds (pre-pulled image, free GPU node) to minutes (image pull + GPU node autoscale). Predictable on a steady-state cluster.
**$/GPU·hr.** **Whatever the underlying K8s cluster pays.** This is the *cheapest* tier in this audit if the user already runs a GPU K8s cluster (e.g., RunPod K8s, Lambda Cloud, OCI K8s, on-prem). Examples:
- RunPod community cloud K8s: ~$1.19/hr A100-80GB, ~$1.99/hr H100.
- Lambda K8s: ~$1.29/hr A100-40GB, ~$2.49/hr H100-80GB.
- On-prem owned hardware: amortized $0.50–$1.00 per A100/H100 hour.
**Max concurrent jobs.** Cluster capacity. Volcano's queue-based admission control + Kubernetes-native quotas govern this.
**Verified API for spinning up N parallel jobs** (Volcano `Job` + KubeRay pattern from the docs):
```yaml
# k8s manifest, one per DiLoCo replica
apiVersion: batch.volcano.sh/v1alpha1
kind: Job
metadata: {name: diloco-replica-0}
spec:
minAvailable: 1
schedulerName: volcano
queue: diloco-queue
tasks:
- replicas: 1
name: replica
template:
spec:
containers:
- name: trainer
image: myorg/composer-replication:latest
command: ["python", "-m", "composer_replication.diloco.serverless.replica_entrypoint",
"--replica-id", "0", "--num-replicas", "4",
"--rendezvous-uri", "s3://minio.cluster.local/diloco/"]
resources:
limits: {nvidia.com/gpu: 1}
restartPolicy: OnFailure
```
…and the framework's `K8sExecutor` adapter does `kubectl apply -f` (or uses the Python K8s client) for each of N rendered manifests.
**Right abstraction.** Either one `volcano.batch.Job` per replica (simple, no Ray) or one `RayJob` per replica (overkill for DiLoCo, but useful if you want Ray Tune integration). One pod = one DiLoCo replica.
### 2.7 RunPod / Lambda / Vast.ai (honourable mentions)
Not in the original candidate list, but worth one paragraph each because they're the price-leaders for serverless GPUs:
- **RunPod Serverless / Pods.** Cheap on-demand A100/H100 (~$1.19–$2.17/hr A100-80GB; ~$1.99–$4.18/hr H100). REST API `POST /v2/{endpoint}/run` for serverless; SDK `runpod` for pods. No native multi-job network — same S3 workaround. **Strong third adapter candidate** for a cost-optimised deployment.
- **Lambda Cloud (Lambda Labs).** Bare metal hourly rentals, not a true serverless API. Programmatic launch via `lambdalabs` API. Outside the "serverless" framing.
- **Vast.ai.** Bidding-style spot market. API-driven launches. Cheapest per A100·hr in the market, but variable availability.
We do **not** include these as v0 adapters but document them as "next-up after Modal + HF" if the user wants further price compression.
---
## 3. The right abstraction: `composer_replication.diloco.serverless`
### 3.1 The core interface
```python
# composer_replication/diloco/serverless/_protocol.py
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Iterator, Protocol
@dataclass(frozen=True)
class ReplicaSpec:
"""One DiLoCo replica's launch config. Mirrors `make_diloco_outer_loop()`'s
args (see composer_replication/diloco/__init__.py) plus a rendezvous_uri
for the object-store all-reduce backend."""
replica_id: int
num_replicas: int
rendezvous_uri: str # s3://, gs://, az://, hf://, file://
model_id: str # e.g. "Qwen/Qwen2.5-7B"
inner_optimizer: dict[str, Any] # serializable; reconstructed in worker
sync_every: int = 500
outer_lr: float = 0.7
outer_momentum: float = 0.9
outer_rounds: int = 100
extra_env: dict[str, str] | None = None
@dataclass(frozen=True)
class ReplicaHandle:
replica_id: int
backend: str # "modal" | "hfjobs" | "sagemaker" | ...
job_id: str
log_url: str | None = None
@dataclass(frozen=True)
class ReplicaResult:
replica_id: int
status: str # "completed" | "failed" | "preempted"
final_checkpoint_uri: str | None
metrics: dict[str, Any]
class ServerlessExecutor(Protocol):
"""Protocol any serverless backend implements to host Decoupled DiLoCo."""
def launch_replicas(self, specs: list[ReplicaSpec]) -> list[ReplicaHandle]: ...
def poll(self, handles: list[ReplicaHandle]) -> list[ReplicaHandle]: ...
def stream_logs(self, handle: ReplicaHandle) -> Iterator[str]: ...
def cancel(self, handles: list[ReplicaHandle]) -> None: ...
def collect(self, handles: list[ReplicaHandle], *,
timeout: float | None = None) -> list[ReplicaResult]: ...
@property
def backend_name(self) -> str: ...
@property
def supports_inter_replica_network(self) -> bool:
"""True iff backend natively connects replicas (e.g., Modal i6pn).
False = pseudo-grad must use rendezvous_uri object store. Default rendezvous
is *always* object-store regardless; this flag only unlocks an opt-in
same-backend fast path (see ModalExecutor(use_clustered_rdma=True))."""
...
```
Concrete adapters inherit from a small `BaseExecutor(ABC)` for cross-cutting retry/log/timeout, paralleling `composer_replication.trainer.composer_trainer`. `launch_replicas()` is partial-failure tolerant: on partial submit it returns handles for the K successful replicas with the failed one carrying `job_id=""` and a logged warning; the caller is responsible for cleanup via `cancel()`.
### 3.2 The object-store all-reduce (the secret weapon)
The whole point of "decoupled" DiLoCo is that the cross-replica primitive is just object-store I/O. We implement it at the framework layer, *not* at the executor layer, so every adapter gets it for free:
```python
# composer_replication/diloco/serverless/_rendezvous.py
import time, torch, fsspec
class ObjectStoreAllReduce:
"""Drop-in for `torchft.Manager.allreduce` over a shared object store.
Each round t:
(1) replica i writes {uri}/round-{t}/replica-{i}.pt
(2) all replicas barrier on count == num_replicas
(3) rank 0 averages, writes {uri}/round-{t}/avg.pt
(4) others read avg.pt, copy_ into the in-place tensor
(5) rank 0 GCs round-(t-1)
fsspec-backed so one path covers s3://, gs://, az://, hf://, file://.
"""
def __init__(self, replica_id, num_replicas, rendezvous_uri,
fsspec_kwargs=None, poll_s=2.0, timeout_s=600.0):
self.replica_id, self.num_replicas = replica_id, num_replicas
self.uri = rendezvous_uri.rstrip("/")
self.fs, _ = fsspec.url_to_fs(self.uri, **(fsspec_kwargs or {}))
self.poll, self.timeout, self._round = poll_s, timeout_s, 0
def allreduce(self, tensor):
t = self._round
my = f"{self.uri}/round-{t}/replica-{self.replica_id}.pt"
avg = f"{self.uri}/round-{t}/avg.pt"
with self.fs.open(my, "wb") as f:
torch.save(tensor.cpu(), f)
deadline = time.time() + self.timeout
while time.time() < deadline:
existing = [p for p in self.fs.ls(f"{self.uri}/round-{t}/")
if p.endswith(".pt") and "/replica-" in p]
if len(existing) >= self.num_replicas: break
time.sleep(self.poll)
else:
raise TimeoutError(f"barrier timeout at round {t}")
if self.replica_id == 0:
tensors = [torch.load(self.fs.open(f"{self.uri}/round-{t}/replica-{i}.pt", "rb"),
map_location="cpu") for i in range(self.num_replicas)]
torch.save(torch.stack(tensors).mean(dim=0), self.fs.open(avg, "wb"))
deadline = time.time() + self.timeout
while time.time() < deadline:
if self.fs.exists(avg):
tensor.copy_(torch.load(self.fs.open(avg, "rb"), map_location=tensor.device))
break
time.sleep(self.poll)
else:
raise TimeoutError(f"avg.pt timeout at round {t}")
if self.replica_id == 0 and t > 0:
try: self.fs.rm(f"{self.uri}/round-{t-1}/", recursive=True)
except Exception: pass
self._round += 1
return _DummyWork()
def should_commit(self): return True
def start_quorum(self, *_, **__): pass
@property
def current_step(self): return self._round
class _DummyWork:
def wait(self): pass
def get_future(self): pass
```
The `ObjectStoreAllReduce` mocks the torchft `Manager` interface — exactly what `make_diloco_outer_loop` already takes (see `composer_replication/diloco/__init__.py` lines 64–125). **No changes to the existing DiLoCo wrapper needed.**
### 3.3 Replica entrypoint
This is the script every adapter runs in its container:
```python
# composer_replication/diloco/serverless/replica_entrypoint.py
"""Run one Decoupled DiLoCo replica. Designed to be invoked as
python -m composer_replication.diloco.serverless.replica_entrypoint \
--replica-id N --num-replicas K --rendezvous-uri s3://... \
--model-id Qwen/Qwen2.5-7B --sync-every 500 --outer-rounds 100
"""
import argparse, os, torch
from composer_replication.diloco import make_diloco_outer_loop
from composer_replication.diloco.serverless._rendezvous import ObjectStoreAllReduce
def main() -> None:
p = argparse.ArgumentParser()
p.add_argument("--replica-id", type=int, required=True)
p.add_argument("--num-replicas", type=int, required=True)
p.add_argument("--rendezvous-uri", required=True)
p.add_argument("--model-id", required=True)
p.add_argument("--sync-every", type=int, default=500)
p.add_argument("--outer-rounds", type=int, default=100)
p.add_argument("--outer-lr", type=float, default=0.7)
args = p.parse_args()
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(args.model_id, torch_dtype=torch.bfloat16).cuda()
inner_opt = torch.optim.AdamW(model.parameters(), lr=4e-4)
manager = ObjectStoreAllReduce(replica_id=args.replica_id,
num_replicas=args.num_replicas,
rendezvous_uri=args.rendezvous_uri)
outer = make_diloco_outer_loop(
manager=manager, model_fragments=[model], inner_optimizer=inner_opt,
outer_lr=args.outer_lr, outer_momentum=0.9, nesterov=True,
sync_every=args.sync_every,
)
with outer:
for outer_round in range(args.outer_rounds):
for inner_step in range(args.sync_every):
# caller plugs in their data + loss; for v0 we use a sketch.
inner_opt.zero_grad(); ...; inner_opt.step()
# outer-loop sync fires automatically at sync_every step boundary.
# Push final checkpoint to rendezvous_uri/final/replica-N.pt
...
if __name__ == "__main__":
main()
```
### 3.4 Package layout
> **Realised in v0.1 (Wave 17 update):** ADR-005 shipped a flatter
> layout than the proposal below. The actual `composer_replication/diloco/serverless/`
> tree:
>
> ```
> composer_replication/
> └── diloco/
> ├── __init__.py # existing: make_diloco_outer_loop, torchft import
> └── serverless/
> ├── __init__.py # re-exports all public classes (Wave 17a)
> ├── executor.py # ServerlessExecutor Protocol + ReplicaHandle
> │ # + LocalProcessExecutor concrete adapter
> ├── allreduce.py # ObjectStoreAllReduce + MockManager
> ├── modal.py # ModalExecutor (skeleton — see __init__ docstring)
> ├── hf_jobs.py # HFJobsExecutor (skeleton — uses huggingface_hub.run_job)
> ├── replica_entrypoint.py # script each replica runs
> └── tests/ # multi-process file:// rendezvous tests
> ```
>
> No leading underscores, no `_protocol`/`_base`/`_rendezvous` split,
> and Modal/HFJobs are flat modules rather than subpackages. The full
> public re-export surface (verified by
> `tests/test_serverless_local.py::test_public_reexports_include_all_executors`):
>
> ```python
> from composer_replication.diloco.serverless import (
> ServerlessExecutor, # Protocol — implement to add your own backend
> LocalProcessExecutor, # multi-process local replicas (CPU/GPU)
> ModalExecutor, # Modal cloud — skeleton in modal.py
> HFJobsExecutor, # HuggingFace Jobs — skeleton in hf_jobs.py
> ObjectStoreAllReduce, # fsspec-backed allreduce (s3://, gs://, file://, hf://)
> MockManager, # torchft.Manager-shaped duck-type
> ReplicaHandle, # opaque handle returned by launch_replicas
> )
> ```
>
> Wave 16's user-journey reviewer caught that earlier versions of this
> `__init__.py` defined `ModalExecutor` and `HFJobsExecutor` in their
> respective modules but failed to re-export them from the package
> namespace. Wave 17a fixed the re-exports and added a regression
> test. The proposal below predates that fix.
```
composer_replication/
└── diloco/
├── __init__.py # existing: make_diloco_outer_loop, torchft import
└── serverless/
├── __init__.py # re-exports
├── _protocol.py # ServerlessExecutor Protocol, ReplicaSpec, ReplicaHandle, ReplicaResult
├── _base.py # BaseExecutor(ABC) — common retry/log/timeout logic
├── _rendezvous.py # ObjectStoreAllReduce (the cross-cutting allreduce)
├── replica_entrypoint.py # the script every adapter runs in-container
├── modal/
│ ├── __init__.py # ModalExecutor
│ └── adapter.py
├── hfjobs/
│ ├── __init__.py # HFJobsExecutor
│ └── adapter.py
└── runpod/ # optional v0.1+
├── __init__.py
└── adapter.py
```
**v0 ships:** `Modal` + `HFJobs`. Both inherit from `BaseExecutor`, both delegate cross-replica state to `ObjectStoreAllReduce`. Symmetric implementation surface ≈ 250 lines per adapter.
**v0.1+ candidates** (add when needed): SageMaker, Vertex AI, Azure ML, RunPod, K8s/Volcano. The `Protocol` is stable; adding adapters is incremental.
### 3.5 What the user writes
```python
from composer_replication.diloco.serverless import (
ModalExecutor, HFJobsExecutor, ReplicaSpec
)
specs = [
ReplicaSpec(replica_id=i, num_replicas=4,
rendezvous_uri="s3://my-diloco-runs/2026-05-26/",
model_id="Qwen/Qwen2.5-7B",
inner_optimizer={"name": "AdamW", "lr": 4e-4},
sync_every=500, outer_rounds=100)
for i in range(4)
]
# Option A: all four replicas on Modal A100s
executor = ModalExecutor(gpu="A100-40GB", region=None, preemptible=True)
handles = executor.launch_replicas(specs)
results = executor.collect(handles)
# Option B: heterogeneous fleet — 2 on Modal, 2 on HF Jobs
modal_ex = ModalExecutor(gpu="A100-40GB")
hf_ex = HFJobsExecutor(flavor="a100-large")
modal_handles = modal_ex.launch_replicas(specs[:2])
hf_handles = hf_ex.launch_replicas(specs[2:])
# both groups read+write the SAME s3://... rendezvous URI — they DiLoCo together.
results = modal_ex.collect(modal_handles) + hf_ex.collect(hf_handles)
```
The "heterogeneous fleet" pattern is the **point** of Decoupled DiLoCo as articulated in the user brief. Modal + HF together is a meaningful test that tells us both adapters work and the rendezvous protocol is backend-agnostic.
---
## 4. Cross-cutting design decisions
### 4.1 Why object-store rendezvous is the default (even on Modal)
Even though Modal supports `@modal.experimental.clustered` with RDMA, **the framework default is object-store-based pseudo-gradient exchange.** Reasons:
1. **Backend portability.** Same code runs on Modal, HF, SageMaker, Vertex, Azure, K8s. Adding a new backend is implementing 6 methods (`launch_replicas`, `poll`, `stream_logs`, `cancel`, `collect`, `backend_name`) — *zero* changes to the rendezvous layer.
2. **Cost asymmetry.** RDMA-class networking on Modal requires `@clustered(rdma=True)` which gates on 8 GPUs/node and tighter scheduling — *more* expensive than 4 separate `@function` invocations of 1 GPU each.
3. **DiLoCo's communication is ridiculous overkill for RDMA.** 2 GB every 10 minutes = ~3 Mbps average. S3 GET/PUT at 10 MB/s does it in ~3 min — well under the 10 min outer-round budget.
4. **Failure decoupling.** A clustered-RDMA failure aborts the whole job (gang-scheduled). Object-store rendezvous tolerates a missing replica (skip its tensor in the average) — better matches DiLoCo's natural fault tolerance.
The opt-in escape hatch: `ModalExecutor(use_clustered_rdma=True)` dispatches to `@modal.experimental.clustered(rdma=True)` and skips object-store. This is for the user who wants Modal-only, max-throughput, single-region runs. It's *not* the default and *not* what we test against.
### 4.2 Rendezvous URI scheme support
`fsspec` covers all the storage backends we need:
| Scheme | Backend | Used for |
|---|---|---|
| `s3://` | `s3fs` | SageMaker default; cheapest for AWS-centric runs |
| `gs://` | `gcsfs` | Vertex AI default |
| `az://` | `adlfs` | Azure ML default |
| `hf://` | `huggingface_hub.HfFileSystem` | HF Jobs preferred (Volume mount makes it look like local fs already) |
| `file://` | builtin | local single-host tests; CI |
The framework picks the *right* default per-executor (Modal → `s3://`, HF → `hf://`, SageMaker → `s3://`, etc.) but always allows override.
### 4.3 Failure model
**Replica failure mid-round.** The barrier in `ObjectStoreAllReduce` has a configurable timeout (default 600 s). If a replica doesn't write its file by then, rank-0 (the averager) has two options governed by `replica_failure_policy`:
- `"strict"` (default): TimeoutError → all replicas abort. Resume from last committed checkpoint.
- `"skip"`: rank-0 averages over what's there, includes a `--num-survivors=K` annotation in `avg.pt`. Other replicas read this and continue. DiLoCo paper §4.5 reports robustness to occasional missing workers; this matches that.
**Whole-cluster failure.** Outer rounds checkpoint to `{rendezvous_uri}/checkpoint-{t}/`; restart sets `args.restart_from=T` and skips ahead.
### 4.4 What we explicitly do NOT do
- **No cross-job NCCL.** Even on Modal, even with `clustered`, the framework uses object-store rendezvous. (Modal `clustered` is exposed only via the explicit opt-in flag.)
- **No DDP/FSDP across replicas.** Each replica is its own self-contained DDP/FSDP world; replicas talk to each other only via the outer-loop. This is the *core* of DiLoCo.
- **No "control plane" service.** No coordinator process, no scheduler container. The object store *is* the coordinator (writes are the messages, file-existence is the synchronization). This is what makes the system work across heterogeneous executors with no shared infra.
- **No Modal-specific or HF-specific dependencies in `composer_replication.diloco`.** Adapter dependencies (`modal`, `huggingface_hub`) are imported lazily inside the adapter modules, exactly how `torchft` is imported lazily in `composer_replication/diloco/__init__.py` today.
---
## 5. Risks and mitigations
| Risk | Likelihood | Mitigation |
|---|---|---|
| Object-store latency dominates outer-round wallclock for large models | M | For 70B+, add `fsspec` parallel-upload (multipart) + bf16 quantize on-write. Most outer rounds are 7B-scale where 2 GB transfer is well under 1 min. |
| Rank-0 replica crashes mid-average → orphaned barrier | L | Add a `lock-{t}.json` heartbeat with TTL; any non-zero replica that sees a stale lock can take over. v1+. |
| Modal + HF cost arbitrage misleading because preemption rates differ | M | Track preemption-rate per backend, surface in `ReplicaResult.metrics`. User-visible. |
| HF Jobs has no public per-account concurrency cap → may hit a hidden limit at N=8 | L | Add exponential-backoff retry around `run_job`; cap `max_concurrent_launches` configurable per executor. |
| AWS / GCP / Azure premiums make their adapters effectively price-uncompetitive | H (already true) | Be honest in docs (this doc). Recommend Modal + HF for cost-sensitive users; cloud-vendor adapters for users who *must* run there for compliance or credits. |
| Rendezvous bucket becomes a security choke point (model weights exposed) | M | Document that `rendezvous_uri` should be a private bucket with replica-only IAM/principals. Provide `RendezvousAccessPolicy` helper that emits boto3/gcloud/az IAM JSON. |
| Modal `@experimental.clustered` API churn (it's experimental) | M | Default path doesn't depend on `clustered`. Fall-back path uses regular `@function`. Document the opt-in clearly. |
| torchft sign-convention regression | L | Already pinned with the unit test in spike 008 (see `spikes/008-streaming-diloco/tests/test_diloco_smoke.py::test_diloco_pseudogradient_sign_convention`). The serverless layer doesn't touch this — it only swaps in a different `Manager.allreduce` impl. |
---
## 6. Validation plan
Three smoke tests, in order of cost:
1. **Spike 009-A (free, ≤30 min):** `LocalProcessExecutor` + `ObjectStoreAllReduce` with `file://` rendezvous. Two in-process replicas DiLoCo-train a 0.5B model on MNIST-equivalent text data. Asserts the rendezvous protocol works.
2. **Spike 009-B (Modal, ≤$5):** `ModalExecutor` × 2 replicas, A100-40GB each, Qwen2.5-0.5B, 50 inner steps × 2 outer rounds. Asserts the Modal adapter launches, replicas find each other through S3 rendezvous, and pseudo-gradients average correctly. Cost: ~30 min × $2.10 × 2 = $2.10 + setup overhead, comfortable under cap.
3. **Spike 009-C (heterogeneous, ≤$10):** 1 Modal A100 + 1 HF Jobs `a100-large`. Same model, 2 outer rounds. Validates that rendezvous works across backends — the key claim of Decoupled DiLoCo. Cost: ~30 min × ($2.10 + $2.50) = ~$2.30, plus per-job startup.
Each spike has a verdict.md following the conventions from `spikes/008-streaming-diloco/`.
---
## 7. References (primary sources, all cited above)
- **DiLoCo paper:** Douillard et al., "DiLoCo: Distributed Low-Communication Training of Language Models," arXiv:2311.08105 (2023). <https://arxiv.org/abs/2311.08105>
- **Streaming DiLoCo paper:** Liu et al., "Streaming DiLoCo with overlapping communication," 2025. <https://arxiv.org/abs/2501.18512>
- **torchft `local_sgd.DiLoCo`:** <https://github.com/meta-pytorch/torchft/blob/main/torchft/local_sgd.py>
- **Modal multi-node clusters:** <https://modal.com/docs/guide/multi-node-training>
- **Modal cluster networking (i6pn):** <https://modal.com/docs/guide/private-networking>
- **Modal pricing:** <https://modal.com/pricing>
- **Modal GPU options:** <https://modal.com/docs/guide/gpu>
- **HF Jobs overview:** <https://huggingface.co/docs/hub/en/jobs>
- **HF Jobs pricing:** <https://huggingface.co/docs/hub/jobs-pricing>
- **HF Jobs Python API:** <https://huggingface.co/docs/huggingface_hub/en/guides/jobs>
- **HF Jobs reference:** <https://huggingface.co/docs/huggingface_hub/main/en/package_reference/jobs>
- **AWS SageMaker pricing:** <https://aws.amazon.com/sagemaker/ai/pricing/>
- **AWS SageMaker `CreateTrainingJob` API:** <https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateTrainingJob.html>
- **AWS SageMaker SMDDP:** <https://docs.aws.amazon.com/sagemaker/latest/dg/data-parallel-framework-estimator.html>
- **AWS SageMaker warm pools:** <https://docs.aws.amazon.com/sagemaker/latest/dg/train-warm-pools.html>
- **GCP Vertex AI compute config:** <https://cloud.google.com/vertex-ai/docs/training/configure-compute>
- **GCP Vertex AI training SKUs:** <https://cloud.google.com/skus/sku-groups/vertex-training>
- **GCP Vertex AI pricing:** <https://cloud.google.com/vertex-ai/pricing>
- **Azure ML pricing:** <https://azure.microsoft.com/en-us/pricing/details/machine-learning/>
- **Azure ML PyTorch SDK v2 guide:** <https://learn.microsoft.com/en-us/azure/machine-learning/how-to-train-pytorch>
- **Azure NDasrA100_v4 spec:** <https://learn.microsoft.com/en-us/azure/virtual-machines/sizes/gpu-accelerated/ndasra100v4-series>
- **Azure NCads H100 v5 spec:** <https://learn.microsoft.com/en-us/azure/virtual-machines/ncads-h100-v5>
- **Volcano:** <https://volcano.sh/en/docs/unified_scheduling/>
- **Volcano network-topology-aware scheduling:** <https://volcano.sh/en/docs/network_topology_aware_scheduling/>
- **KubeRay + Volcano integration:** <https://docs.ray.io/en/latest/cluster/kubernetes/k8s-ecosystem/volcano.html>
- **KubeRay RayJob+Volcano PR:** <https://github.com/ray-project/kuberay/pull/3972>
Internal references (in this repo):
- `docs/research/MODAL_RECONNAISSANCE.md` — pricing/cold-start audit for Modal smoke runs.
- `docs/research/DILOCO_RECONNAISSANCE.md` — DiLoCo implementation candidates audit.
- `docs/adrs/ADR-001-gpu-venue.md` — local-vs-cloud GPU decision for smoke phase.
- `docs/adrs/ADR-003-diloco-impl.md` — torchft choice + sign convention.
- `composer_replication/diloco/__init__.py` — existing `make_diloco_outer_loop` wrapper this design plugs into without modification.
- `spikes/008-streaming-diloco/` — the existing in-process DiLoCo smoke that the serverless adapter inherits sign-convention test from.