feat(weight-sync): delta weight sync

[YSunLIN]

feat(weight-sync): shard-local delta sync for full-weight rollout updates

Summary

Motivation

In disaggregated RL (dedicated rollout engine, separate/colocate layouts), every training step ends with a weight sync: the trained FSDP model is pushed into the rollout engine so the next rollout samples fresh weights. Today all three full-weight transports (NCCLWeightSync, IPCWeightSync, TensorWeightSync) do this densely — FullWeightSync._iter_full_tensors() redistributes every FSDP DTensor shard to Replicate (a whole-mesh all-gather), buckets the full tensors, and ships every byte of the model on every step. At scale this dominates the sync back-edge: sync time grows linearly with parameter count, and for large MoE models it can cost minutes per step.

The key observation: under typical RL learning rates, the vast majority of BF16 weight bytes are bit-for-bit unchanged step-over-step. Only a small fraction (~1–3%) of parameter elements actually move each step. Shipping the whole model to communicate a 1–3% change is the waste we want to remove.

Positioning

The existing dense path is not just "the simple version" — it is the all-gather-then-stage pattern that scales with model size regardless of how little changed. A delta path instead scales with how much changed:

• Dense (today): all-gather each shard → full tensors on rank 0 → bucket → broadcast. Cost ∝ model size, every step.
• Delta (this work): each rank diffs its own local shard against a pinned snapshot → gathers only changed (index, value) pairs → broadcast the sparse payload. Cost ∝ changed bytes. No full-tensor all-gather, no rank-0 full-model staging.

Delta is a sibling full-weight handler, not a rewrite of the dense ones: it reuses the same Remote BROADCAST dispatch, the same bucketed flush mechanism, and the same name_remap/track_prefix routing. It applies only to full base weights — the LoRA family already ships just the tiny adapter, so delta is out of scope there by construction.

Design

1. Config surface. A single additive delta_sync: bool = false kwarg on FullWeightSync (threaded through all three transports). When false, behavior is byte-identical to today. When true, the handler consults a DeltaWeightEncoder instead of pushing dense tensors.

2. Encoder interface (decision engine).
DeltaWeightEncoder owns the snapshot + diff, per shard:

• seed(name, shard) — record the first-sync dense baseline (no delta emitted; the seeding sync is a full dense pass so both sides share a baseline).
• encode(name, shard) -> SparseDelta — bit-exact integer diff of shard against its pinned snapshot, refresh the snapshot, return the changed positions as SparseDelta(indices: int32 [nnz], values: shard-dtype [nnz]).

The diff is bit-exact (integer comparison of the raw bytes), so the reconstructed weights on the rollout side are identical to a dense push — no thresholding, no numeric drift. Snapshots are pinned CPU tensors of the local shard only, so memory overhead is one extra shard-sized copy per rank, independent of world size.

3. Two-phase protocol.

• Seed (first sync): dense pass through the existing bucketed wire, no position data; establishes the baseline snapshot on both trainer and rollout.
• Steady state: each rank diffs its shard locally, encodes changed positions, and rank 0 assembles the disjoint per-rank pieces, buckets them deterministically, and broadcasts. The rollout side applies deltas in place on live weights — no full-model mirror is ever materialized.

4. Integrity net. Per-flush checksums (reusing
transfer/checksum.py::fingerprint_tensor) verify on the receiver that the applied result matches the trainer's post-step weights, catching any lost or misordered delta without a full read-back.

5. Backend integration. sync() in each transport gains a delta branch:
when self._delta_encoder is set, iterate shard-local deltas instead of _iter_full_tensors(), and drive the matching sparse-apply receiver (update_weights_from_distributed / update_weights_from_ipc gain a sparse payload path on the engine side). The dense branch is unchanged.

Test Plan

• Dense-parity: existing delta_sync=false recipes — logprob + weight fingerprint parity vs the current dense push, must remain byte-identical.
• Delta path (once landed): after N steps, assert (a) the seed sync is dense, (b) steady-state payloads carry only changed positions, (c) the reconstructed rollout weights fingerprint-match a dense push of the same step's weights, (d) end-to-end rollout/loss parity vs the dense baseline.
• SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failure.
• Throughput smoke test comparing sync wall-clock vs the dense baseline across model sizes (recipe/hardware TBD).
• Current commit: all changed files parse (ast.parse); docstring guard passes; no wired runtime path beyond the fail-closed guard on delta_sync=true.

Compatibility / Risk

• Config: additive; delta_sync defaults to false, dense semantics preserved across all three transports.
• Behavior: none at default. Opting into delta_sync fails closed until the encoder lands.
• Numerics: bit-exact diff — reconstructed weights are identical to a dense push; the per-flush checksum gates this.
• Scope: full-weight only; LoRA sync is unaffected by design.
• No checkpoint or data-format changes.

Reviewer Notes

• This is a draft. Please review the interface shape first — the DeltaWeightEncoder.seed/encode contract and the SparseDelta (indices, values) wire shape are what everything downstream depends on.
• Open question: should the snapshot live in the encoder (as scaffolded) or in the backend (which already owns the FSDP shard lifecycle)? Encoder-owned keeps the diff self-contained; backend-owned avoids a second shard copy if the backend already pins one.
• Open question: MoE / fused-expert params (_iter_full_tensors_ep) produce gathered per-expert HF tensors, not raw local shards — delta over those needs a slot table or a dense fallback for the expert block. Out of scope for the first slice; the seed/steady split makes a per-param dense fallback cheap.
• AI-assisted; duplicate-work check: no open PR touches the full-weight sync transports or adds a delta/snapshot encoder.