[checkpoint_engine][rollout] Delta weight sync over NCCL for disaggregated rollout (+ sharded-snapshot variant)

docs/advance/delta_weight_sync.md

@@ -0,0 +1,91 @@
+# Delta Weight Sync
+
+Last updated: 07/09/2026.
+
+## Motivation
+
+In a disaggregated setup (``hybrid_engine=False``) the trainer must broadcast its updated weights to the
+rollout engine after every step. By default this is a full-weight broadcast whose cost grows with model
+size. Because RL updates are highly sparse — under typical learning rates over 99% of BF16 weight bytes
+are unchanged step-over-step — you can instead broadcast only the parameters that changed (a *delta*),
+cutting the weight-sync traffic to the sparsity ratio while staying lossless (bit-exact; a per-flush
+checksum is verified on the receiver).
+
+When to use: disaggregated training where the trainer↔rollout link is the bottleneck (commodity network /
+cross-node / large models). On a fast intra-node NVLink link with a small model the full broadcast is
+already cheap, so delta mainly pays off as the model size and the network distance grow.
+
+## Design
+
+The delta backends plug into the standard checkpoint-engine flow (``CheckpointEngineManager`` →
+``CheckpointEngineWorker``), so they work with any trainer that drives weight sync through the
+checkpoint engine (including the V1 ``separate_async`` trainer).
+
+- **Diff**: the trainer byte-diffs each parameter against a pinned-CPU snapshot of the previous sync.
+  The comparison is bit-exact (integer view inequality), so the reconstruction is lossless by
+  construction — no thresholds, no drift.
+- **Encoding**: changed elements are shipped as a shared ``(positions, values)`` payload plus a
+  per-parameter manifest. ``encoding`` selects the position encoding: ``indices`` (int32 absolute
+  positions, lowest compute) or ``deltas`` (uint16 gap deltas, smaller on the wire).
+- **Transport**: the sparse payload is broadcast over the existing NCCL collective group in
+  bucket-sized flushes (streamed: each flush is sent and freed as it is produced, so sender peak
+  memory stays ~2 buckets regardless of model size).
+- **Apply**: each rollout worker hands its local copy of the sparse payload to its colocated SGLang
+  TP worker via same-GPU ``update_weights_from_tensor`` IPC, where a verl-shipped loader —
+  registered automatically through SGLang's stock ``--custom-weight-loader`` hook, so **no SGLang
+  fork or patch is needed** — verifies the flush checksum (fail loud), densifies each parameter's
+  delta into a NaN-masked tensor, and overwrites only the changed positions *in place* on the live
+  weights. No full-model mirror is staged anywhere on the rollout side: receiver peak memory is one
+  bucket plus one decode chunk, independent of model size.
+- **Seeding**: the first sync broadcasts a full delta (every position), so a dummy-initialized rollout
+  gets a correct base; subsequent syncs are sparse.
+
+## Backends
+
+### ``delta``
+
+Diffs on rank 0 after the ordinary full-tensor gather; rank 0 holds a full-model pinned-CPU snapshot.
+
+```shell
+    actor_rollout_ref.rollout.checkpoint_engine.backend=delta \
+    +actor_rollout_ref.rollout.checkpoint_engine.engine_kwargs.delta.encoding=indices
+```
+
+### ``delta_sharded`` (sharded snapshot)
+
+The ``delta`` backend above still ``full_tensor()``-gathers every parameter to rank 0 before diffing.
+The ``delta_sharded`` backend pushes the diff *below* the all-gather: each actor rank pins a snapshot of
+only **its** FSDP shard, byte-diffs the shard locally, and gathers just the changed ``(position, value)``
+pairs to rank 0 (via the engine's ``get_per_tensor_param_shard()`` export). So the gather volume drops
+from the full parameter to the sparsity ratio (~1–3%), and no rank needs a full-model snapshot — the
+memory and the gather traffic both shard with the world size.
+
+```shell
+    actor_rollout_ref.rollout.checkpoint_engine.backend=delta_sharded \
+    +actor_rollout_ref.rollout.checkpoint_engine.engine_kwargs.delta_sharded.encoding=indices
+```
+
+The assembled delta is **bit-identical** to what ``delta`` produces, so the wire format, the per-flush
+checksum, and the rollout-side receiver are all unchanged. Each rank computes its shard's absolute
+position in the full flattened parameter purely locally (from the DTensor spec, no extra collective).
+Scope: FSDP2 ``Shard(0)`` parameters (the common case) plus replicated / non-DTensor params; other shard
+dimensions are not supported and raise.
+
+## Usage
+
+A runnable example is ``verl/experimental/one_step_off_policy/shell/grpo_0.6b_gsm8k_fsdp2_sglang_delta_2_6.sh`` —
+the SGLang 2+6 disaggregated GRPO recipe with ``backend=delta``.
+
+Current scope: disaggregated (``hybrid_engine=False``) + SGLang rollout in BF16, FSDP2 training engine.
+
+## Roadmap
+
+Planned extensions, in design order:
+
+- **Megatron sharded delta**: diff each rank's mcore shard locally and reuse the native
+  megatron→HF bridge on rank 0, extending ``delta_sharded`` beyond FSDP.
+- **Quantized rollout (fp8 etc.)**: diff the quantized bytes (quantize-then-diff) so a low-precision
+  rollout engine can consume deltas without a bf16 intermediate.
+- **Upstream SGLang delta receiver**: once SGLang ships a native distributed delta receiver, the
+  broadcast can land directly on the TP workers (dropping the same-GPU IPC forward); the current
+  custom-weight-loader apply is already in place and bit-compatible with that wire format.

docs/advance/one_step_off.md

@@ -309,6 +309,13 @@ python3 -m verl.experimental.one_step_off_policy.async_main_ppo \
        > - When `trainer.n_gpus_per_node + rollout.n_gpus_per_node > physical_gpus_per_node`,
        >   the required node count is `trainer.nnodes + rollout.nnodes`
 
+### Delta Weight Sync
+
+For disaggregated runs the trainer→rollout weight broadcast can ship only the changed parameters
+(a *delta*) instead of the full weights. See the dedicated design doc:
+[Delta Weight Sync](delta_weight_sync.md), covering the ``delta`` and ``delta_sharded``
+checkpoint-engine backends, their configuration, and the roadmap (Megatron, fp8).
+
 ## Functional Support
 
 | Category           | Support Situation                                                                                               |

docs/index.rst

@@ -122,6 +122,7 @@ verl is fast with:
    :caption: Async Training
 
    advance/one_step_off
+   advance/delta_weight_sync
    advance/fully_async
    advance/async-on-policy-distill
 

docs/workers/engine_workers.rst

@@ -114,7 +114,10 @@ Key RPCs
      directly. LoRA adapters are merged into base weights up-front when
      ``model.lora.merge=True``.
    - For **disaggregated async training**: send the weights through
-     ``self.checkpoint_engine.send_weights`` instead.
+     ``self.checkpoint_engine.send_weights`` instead. The transport is chosen by
+     ``checkpoint_engine.backend`` (e.g. ``nccl`` for a full-weight broadcast, or
+     ``delta`` to broadcast only the parameters that changed since the previous
+     sync — see :doc:`../advance/one_step_off`).
 
 5. ``save_checkpoint`` / ``load_checkpoint``
 

tests/checkpoint_engine/sharded_delta_multigpu_check.py

@@ -0,0 +1,90 @@
+"""Validate verl.checkpoint_engine.delta_sync.sharded against the full-gather diff.
+
+torchrun --nproc_per_node=4 tests/checkpoint_engine/sharded_delta_multigpu_check.py
+"""
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.tensor import Replicate, Shard, distribute_tensor
+
+from verl.checkpoint_engine.delta_sync.sharded import (
+    gather_v_to_rank0,
+    local_shard_view,
+    shard_delta_indices,
+)
+
+
+def _run_case(shape, placements, mesh, dev, rank, si):
+    """Diff a distributed param via the real sharded module vs a full-gather baseline."""
+    torch.manual_seed(si)
+    full_old = torch.randn(*shape, dtype=torch.bfloat16, device=dev)
+    full_new = full_old.clone()
+    g = torch.Generator(device=dev).manual_seed(100 + si)
+    numel = full_old.numel()
+    k = max(1, numel // 100)
+    pert = torch.randint(0, numel, (k,), device=dev, generator=g)
+    full_new.view(-1)[pert] += torch.randn(k, dtype=torch.bfloat16, device=dev, generator=g) * 0.1
+
+    dt_old = distribute_tensor(full_old, mesh, placements)
+    dt_new = distribute_tensor(full_new, mesh, placements)
+
+    # --- sharded path (real module) ---
+    loc_new, off, contributes = local_shard_view(dt_new)
+    loc_old, off2, _ = local_shard_view(dt_old)
+    assert off == off2
+    if contributes:
+        gidx, gval = shard_delta_indices(loc_new, loc_old, off)
+    else:
+        gidx = torch.empty(0, dtype=torch.int64, device=dev)
+        gval = torch.empty(0, dtype=loc_new.dtype, device=dev)
+    sh_idx, sh_val = gather_v_to_rank0(gidx, gval)
+
+    # --- baseline: full gather + diff ---
+    fo = dt_old.full_tensor().reshape(-1)
+    fn = dt_new.full_tensor().reshape(-1)
+    bmask = fn.view(torch.int16) != fo.view(torch.int16)
+    b_idx = bmask.nonzero(as_tuple=False).view(-1).to(torch.int64)
+    b_val = fn[b_idx]
+
+    if rank != 0:
+        return True
+    so = torch.argsort(sh_idx)
+    bo = torch.argsort(b_idx)
+    idx_ok = torch.equal(sh_idx[so], b_idx[bo])
+    val_ok = torch.equal(sh_val[so].view(torch.int16), b_val[bo].view(torch.int16))
+    ok = idx_ok and val_ok and (sh_idx.numel() == b_idx.numel())
+    tag = "x".join(p.__class__.__name__[0] for p in placements)
+    print(f"[shape={shape} mesh={tuple(mesh.shape)} {tag}] nnz sharded={sh_idx.numel()} "
+          f"full={b_idx.numel()} idx={idx_ok} val={val_ok} -> {'PASS' if ok else 'FAIL'}")
+    return ok
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank, world = dist.get_rank(), dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = torch.device("cuda", rank)
+    all_ok = True
+
+    # 1D FSDP mesh, Shard(0) -- uneven shapes stress the offset math.
+    mesh1d = init_device_mesh("cuda", (world,))
+    for si, shape in enumerate([(4096, 1024), (7, 3), (30001, 8), (128,)]):
+        all_ok = _run_case(shape, [Shard(0)], mesh1d, dev, rank, si) and all_ok
+
+    # 2D FSDP x SP(replicate) mesh -- the ulysses case: weights are Shard(0) on the FSDP dim
+    # and Replicate on the SP dim, so only SP-coord-0 ranks may contribute (no double-count).
+    if world % 2 == 0:
+        mesh2d = init_device_mesh("cuda", (world // 2, 2), mesh_dim_names=("fsdp", "sp"))
+        for si, shape in enumerate([(4096, 1024), (7, 3), (128,)]):
+            all_ok = _run_case(shape, [Shard(0), Replicate()], mesh2d, dev, rank, 50 + si) and all_ok
+
+    if rank == 0:
+        print("=" * 50)
+        print(f"OVERALL: {'ALL PASS ✅' if all_ok else 'FAIL ❌'}")
+        print("=" * 50)
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()