Reinforcement Learning
Transformers
English
post-training
distillation
agentic-coding
composer-2.5
cursor
kimi-k2
grpo
dapo
diloco
openenv
trl
verl
research
methodology
Instructions to use Codeseys/composer-replication-framework with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use Codeseys/composer-replication-framework with Transformers:
# Load model directly from transformers import AutoModel model = AutoModel.from_pretrained("Codeseys/composer-replication-framework", dtype="auto") - Notebooks
- Google Colab
- Kaggle
| # 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. | |