feat(weights): S3-compatible object store for the sparse weight broadcast by faresobeid · Pull Request #4103 · PrimeIntellect-ai/prime-rl

faresobeid · 2026-10-15T13:53:53Z

Stacked on #4085.

Adds [weight_broadcast.delta] path for the sparse filesystem broadcast. It is unset by default (the broadcast dir, as today), or another dir, or s3://bucket/prefix, so that trainer and engines can run in different regions.

utils/weight_store.py: S3Path implements the few Path methods the transport uses (/, name, parent, exists, glob, read_text). read_range / remove work on both Path and S3Path. sparse_weights.write_file gains an S3 branch: there is no rename on S3, so a file is one PUT, or a multipart upload in 64 MiB parts for anchors. Endpoint and credentials come from the standard AWS environment variables. boto3 is imported only when S3 is used.
sparse_weights.py: the manifest, delta part files, rank_{r}.parts markers, anchors and engine layouts are all written with write_file, so a part or marker is visible only once complete. The file format is unchanged. Engines poll the store for part files as the trainer uploads them (on S3 too), and read only the frames they need with ranged GETs (pread on a filesystem). Anchors are read from the safetensors headers plus ranged reads, up to 32 tensors ahead.
vLLM workers read S3 deltas in a plain child process (python -m prime_rl.utils.sparse_weights; vLLM workers are daemonic, so no multiprocessing), and the decoded frames come back through a local temporary file. Inside the worker, boto3 releases and retakes the GIL many times per request, and that crawls while the engine serves: 230-340 s per 30B delta, vs a few seconds from the child process.
The handshake markers stay in the trainer's broadcast dir. The orchestrator waits for the manifest in the store and passes s3://.../step_N to the engines, so inference needs no config. With a separate delta.path (a bucket, the bucket root s3://bucket, or another dir), the trainer deletes old steps there and clears leftover step_* entries at startup, as it does for the broadcast dir. The master logs the delta size from the bytes each rank reports in its parts marker.

Qwen3-30B-A3B bf16, 1 node (trainer 6 GPUs, engine TP2), 20 steps, interval = 10, same commit for both columns. MinIO on node-local disk vs the shared filesystem. Medians over v2-v20, measured on an earlier head:

	filesystem	MinIO
delta	340 MiB	345 MiB
trainer encode (parts written while encoding)	0.73 s	0.69 s
anchor write, 57 GiB (background)	24-29 s	129-131 s
engine read + decompress per worker, from the manifest (includes waiting for parts)	1.94 s	4.30 s
prefetch seen by the orchestrator	2.07 s	4.71 s (8-12 s on v10-v15, next to the anchor upload; 3.3-4.1 s on v16-v20)
apply / verify	0.19 / 0.11 s	0.20 / 0.10 s
v1 (first delta) encode / prefetch	16.1 / 17.9 s	38.3 / 47.2 s

All 20 deltas were checksum-verified on both engine workers in both runs.

Offline, on the same node, an anchor load takes 16.5 s from the filesystem and 63 s from MinIO. Fresh engine (measured before part streaming): replay from the v10 anchor to v15, then steps v16-v20: filesystem 41.7 s / 1.4-1.5 s per step; MinIO 102 s / 3.1-3.3 s per step. A delta applied twice fails the checksum check.

#4085 takes the deltas from the optimizer step. With a separate delta.path, the trainer writes the delta parts straight into the storage step dir as the step produces them (there is no rename on S3). Anchors are written as local files next to the broadcast dir (their slices arrive in any order) and uploaded once complete, then the anchor_rank{r}.done marker. Full checkpoints (a resumed run's first broadcast, SFT online evals) go through the same writer. At startup the trainer also clears old engine layouts under a reused prefix: workers of an earlier run would otherwise leave a stale fp8 layout for the trainer to read.

Latest check (Qwen3-30B-A3B AdamW offload = "optimizer", 10 steps, anchor every 5 steps, MinIO on node-local disk): 20/20 deltas verified on both engine workers, both anchors complete in the bucket; fwd/bwd 9.4 s on normal steps (7.0 s with the broadcast dir) and 12.6 / 18.8 s on anchor steps; orchestrator update 11.3 s. CPU round trips against MinIO pass: delta, subprocess read, an anchor written locally, uploaded and read back by flat range across rank slices, bucket-root paths, engine layout cleanup.

All engines read the same bucket. A cache per region would be a replicated copy of the prefix (for example with bucket replication) plus a storage URL per engine; this PR does not implement it.

Validation

Merged the updated feat(weights): sparse XOR deltas for the filesystem weight broadcast #4085 (now up to date with main), no conflicts. CPU on the merged head: pytest tests/unit/orchestrator tests/unit/test_configs.py tests/unit/utils tests/unit/transports tests/unit/inference -m "not gpu" 502 passed; ruff check, ruff format --check, uv lock --check clean.

Before merge

Merge feat(weights): sparse XOR deltas for the filesystem weight broadcast #4085 first, then retarget this PR to main.

Note

High Risk
Changes the critical trainer→engine weight sync path with new S3 I/O, subprocess delta reading, and split handshake vs data locations; misconfiguration or partial uploads could break updates despite checksum checks.

Overview
Adds [weight_broadcast.delta] path so sparse deltas and anchors can live on a separate directory or s3://bucket/prefix, while the broadcast handshake (finished markers, orchestration) stays on the trainer’s local broadcast dir. Engines receive s3://.../step_N from the orchestrator without extra inference config.

Introduces utils/weight_store.py (S3Path, ranged read_range, multipart upload, remove) and threads StoragePath through sparse_weights, the sparse trainer sender, and the vLLM filesystem worker. Delta/anchor I/O uses ranged reads instead of local pread where needed; write_file uploads atomically to S3. On S3, the trainer streams delta parts straight into the remote step dir (no rename/move from hidden staging dirs); anchors are built locally next to the broadcast dir, then uploaded.

vLLM workers call read_delta_in_subprocess for S3 so boto3’s GIL churn does not block serving during large delta downloads. Startup clears stale step_* under a reused prefix; retention also deletes old steps on the separate storage path.

docs/scaling.md documents cross-region setup (AWS env vars, boto3, orchestrator colocated with trainer).

^{Reviewed by Cursor Bugbot for commit ce19c66. Bugbot is set up for automated code reviews on this repo. Configure here.}

cursor