diff --git a/docs/advance/delta_weight_sync.md b/docs/advance/delta_weight_sync.md new file mode 100644 index 00000000000..cb3c41f37c6 --- /dev/null +++ b/docs/advance/delta_weight_sync.md @@ -0,0 +1,107 @@ +# Delta Weight Sync + +Last updated: 07/10/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 is an explicit **dense** pass — the raw weights stream through the same + bucketed wire with no positions attached (values only), populating the trainer-side snapshot as they + go — so a dummy-initialized rollout gets a correct base without any sparse-encoding overhead. + 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). + +**Supported training engines**: the shard export requires ``Shard(0)`` DTensor parameters, which both +FSDP versions provide: + +- **FSDP2** (``fully_shard``, ``actor.strategy=fsdp2``): native DTensor params; the export never stages + the whole shard on the GPU (``state_dict()`` is reference-only, shards move lazily per parameter). +- **FSDP1** (``actor.strategy=fsdp``, the default): verl configures ``SHARDED_STATE_DICT``, whose export + also emits per-rank ``Shard(0)`` DTensors. FSDP1's state-dict export runs through the unshard + machinery, so the whole-shard GPU staging round trip is kept for it (it is skipped for FSDP2). + Single-GPU FSDP1 uses ``FULL_STATE_DICT`` (plain tensors) and degrades to the replicated/rank-0 path — + still correct, just not shard-parallel. + +Other shard dimensions than ``Shard(0)`` are not supported and raise. + +> **Config note**: the training engine reads the **top-level** ``actor_rollout_ref.actor.strategy``; +> setting only ``actor.fsdp_config.strategy`` does *not* select FSDP2. + +## 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, FSDP1/FSDP2 training engines. + +## 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. diff --git a/docs/advance/one_step_off.md b/docs/advance/one_step_off.md index 99170d75edc..165f7cf785d 100644 --- a/docs/advance/one_step_off.md +++ b/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 | diff --git a/docs/index.rst b/docs/index.rst index d98d874a2cf..095bfdc478f 100644 --- a/docs/index.rst +++ b/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 diff --git a/docs/workers/engine_workers.rst b/docs/workers/engine_workers.rst index 673abd188b3..fa4f0fdbbc7 100644 --- a/docs/workers/engine_workers.rst +++ b/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`` diff --git a/tests/checkpoint_engine/sharded_delta_multigpu_check.py b/tests/checkpoint_engine/sharded_delta_multigpu_check.py new file mode 100644 index 00000000000..4a5caf8d94a --- /dev/null +++ b/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() diff --git a/tests/checkpoint_engine/test_delta_sync.py b/tests/checkpoint_engine/test_delta_sync.py new file mode 100644 index 00000000000..dcdb33a1dac --- /dev/null +++ b/tests/checkpoint_engine/test_delta_sync.py @@ -0,0 +1,247 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Bit-identity round-trip tests for the delta sync sender / decoder. + +These exercise the framework-agnostic core (``DeltaState`` + ``encode_chunk`` ++ ``decode_chunk``) without requiring a delta-capable SGLang build. They +verify that: + +* The first call seeds the snapshot and emits no flushes. +* The diff is dtype-agnostic and lossless: every changed element is sent + with the exact trainer-side bytes, while unchanged elements are not in + the payload at all. +* Both encodings (indices, gap-deltas) reconstruct the changed values to + bit-identity at the receiver. +""" + +from __future__ import annotations + +import pytest +import torch + +from verl.checkpoint_engine.delta_sync import ( + DeltaState, + iter_delta_flushes, +) +from verl.checkpoint_engine.delta_sync.encode import decode_chunk + + +def _weights_gen(named: list[tuple[str, torch.Tensor]]): + for n, t in named: + yield n, t + + +def _round_trip_one_flush(flush, expected_changed: dict[str, torch.Tensor]) -> None: + """Decode ``flush`` and verify changed positions match the trainer's bytes.""" + decoded = decode_chunk( + flush.encoding, + flush.positions_cpu.numpy().tobytes(), + flush.values_gpu, + flush.params, + ) + for name, expected in expected_changed.items(): + assert name in decoded, f"{name} missing from receiver payload" + recv = decoded[name] + mask = ~torch.isnan(recv) + # NaN means "not in payload"; for changed positions we must reconstruct + # the exact trainer-side bytes. + assert torch.equal(recv[mask].view(expected.dtype), expected[mask].view(expected.dtype)) + + +def _make_named(dtype=torch.bfloat16) -> list[tuple[str, torch.Tensor]]: + torch.manual_seed(0) + return [ + ("layer.0.weight", torch.randn(64, 32, dtype=dtype)), + ("layer.1.weight", torch.randn(32, 16, dtype=dtype)), + ] + + +@pytest.mark.parametrize("encoding", ["indices", "deltas"]) +def test_first_call_seeds_no_flushes(encoding: str): + state = DeltaState() + named = _make_named() + flushes = list( + iter_delta_flushes( + _weights_gen(named), + state, + encoding=encoding, + bucket_bytes=1 << 20, + ) + ) + assert flushes == [] + assert state.seeded + + +@pytest.mark.parametrize("encoding", ["indices", "deltas"]) +def test_round_trip_bit_identical(encoding: str): + state = DeltaState() + named = _make_named() + + # Seed. + list( + iter_delta_flushes( + _weights_gen(named), + state, + encoding=encoding, + bucket_bytes=1 << 20, + ) + ) + + # Mutate a sparse set of positions in each tensor. + changed_truth: dict[str, torch.Tensor] = {} + new_named = [] + for name, t in named: + new = t.clone() + flat = new.view(-1) + idx = torch.tensor([1, 17, 200, 511], dtype=torch.int64) % flat.numel() + flat[idx] = flat[idx] + 0.5 + new_named.append((name, new)) + changed_truth[name] = new + + # Real sender path. + flushes = list( + iter_delta_flushes( + _weights_gen(new_named), + state, + encoding=encoding, + bucket_bytes=1 << 20, + ) + ) + assert flushes, "expected at least one flush after a real change" + + # Concatenate decoded views; later flushes overwrite earlier ones for the + # same parameter (matches the receiver's apply-then-overwrite semantics). + for flush in flushes: + _round_trip_one_flush(flush, changed_truth) + + +def test_no_change_emits_no_flush(): + state = DeltaState() + named = _make_named() + list( + iter_delta_flushes( + _weights_gen(named), + state, + encoding="indices", + bucket_bytes=1 << 20, + ) + ) + # Identical second pass. + flushes = list( + iter_delta_flushes( + _weights_gen(named), + state, + encoding="indices", + bucket_bytes=1 << 20, + ) + ) + assert flushes == [] + + +def test_dtype_agnostic_diff_fp32(): + state = DeltaState() + named = _make_named(dtype=torch.float32) + list( + iter_delta_flushes( + _weights_gen(named), + state, + encoding="indices", + bucket_bytes=1 << 20, + ) + ) + new_named = [] + truth: dict[str, torch.Tensor] = {} + for name, t in named: + new = t.clone() + new.view(-1)[3] += 1e-3 + new_named.append((name, new)) + truth[name] = new + flushes = list( + iter_delta_flushes( + _weights_gen(new_named), + state, + encoding="indices", + bucket_bytes=1 << 20, + ) + ) + assert flushes + for flush in flushes: + _round_trip_one_flush(flush, truth) + + +def _apply_flushes_to_mirror(mirror: dict[str, torch.Tensor], flushes) -> None: + """Reproduce SGLang's masked delta apply (``_apply_delta_payload``) in-process. + + Decodes each flush and overwrites only the changed (non-NaN) positions into + the running full-weight mirror -- exactly what SGLang's masked-copy loader + does to its live weights when the engine sends + ``update_weights_from_distributed(load_format="delta")``. No GPU / NCCL / + SGLang needed: the transport only moves bytes, so the lossless guarantee is + fully exercised by decode + combine. + """ + for flush in flushes: + decoded = decode_chunk( + flush.encoding, + flush.positions_cpu.numpy().tobytes(), + flush.values_gpu, + flush.params, + ) + for name, recv in decoded.items(): + mask = ~torch.isnan(recv) + mirror[name][mask] = recv[mask] + + +@pytest.mark.parametrize("encoding", ["indices", "deltas"]) +def test_delta_result_equals_full_sync(encoding: str): + """The weights a rollout ends up with via delta sync must be byte-identical + to what the old full-weight path delivers (i.e. the trainer's current W). + + full path -> rollout receives every tensor verbatim == W_new + delta path -> rollout starts from W_old mirror, applies only the changed + positions == W_old + delta + Assert the two are bit-equal for every parameter, across multiple steps. + """ + state = DeltaState() + named = _make_named() + + # Seed the trainer snapshot from W0 and initialize the rollout mirror to the + # same W0 (both sides loaded the identical init checkpoint). + list(iter_delta_flushes(_weights_gen(named), state, encoding=encoding, bucket_bytes=1 << 20)) + mirror = {n: t.clone() for n, t in named} + + cur_named = named + for step in range(3): + # Trainer takes a step: mutate a sparse, step-dependent set of positions. + new_named = [] + for i, (name, t) in enumerate(cur_named): + new = t.clone() + flat = new.view(-1) + idx = torch.tensor([1 + step, 17, 200 + 5 * i, 511], dtype=torch.int64) % flat.numel() + flat[idx] = flat[idx] + 0.25 * (step + 1) + new_named.append((name, new)) + cur_named = new_named + + # full path reference: the rollout would just receive W_new in full. + full_result = {n: t.clone() for n, t in new_named} + + # delta path: sender diffs vs snapshot, receiver applies onto its mirror. + flushes = list( + iter_delta_flushes(_weights_gen(new_named), state, encoding=encoding, bucket_bytes=1 << 20) + ) + _apply_flushes_to_mirror(mirror, flushes) + + for name, full in full_result.items(): + assert torch.equal(mirror[name].view(torch.uint8), full.view(torch.uint8)), ( + f"delta != full at step {step} for {name}" + ) diff --git a/tests/checkpoint_engine/test_sglang_loader.py b/tests/checkpoint_engine/test_sglang_loader.py new file mode 100644 index 00000000000..a80728f87a1 --- /dev/null +++ b/tests/checkpoint_engine/test_sglang_loader.py @@ -0,0 +1,159 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Bit-identity tests for the SGLang custom-weight-loader delta apply. + +Drives the real sender machinery (``DeltaState`` + ``iter_delta_flushes``) and +feeds each flush through :func:`sglang_loader.apply_delta` against a stand-in +model whose ``load_weights`` mimics SGLang's ``param.copy_(loaded)`` semantics. +Verifies the masked in-place apply: changed positions land bit-exactly, and +positions outside the delta payload are never touched. +""" + +from __future__ import annotations + +import json + +import pytest +import torch + +from verl.checkpoint_engine.delta_sync import DeltaState, iter_delta_flushes +from verl.checkpoint_engine.delta_sync.sglang_loader import apply_delta + + +class _FakeModel: + """Holds live params; load_weights lands on param.copy_(loaded), like SGLang.""" + + def __init__(self, named: list[tuple[str, torch.Tensor]]): + self.params = {n: t.clone() for n, t in named} + + def load_weights(self, chunk): + for name, tensor in chunk: + self.params[name].copy_(tensor) + + +def _make_named(dtype=torch.bfloat16) -> list[tuple[str, torch.Tensor]]: + torch.manual_seed(0) + return [ + ("layer.0.weight", torch.randn(64, 32, dtype=dtype)), + ("layer.1.weight", torch.randn(32, 16, dtype=dtype)), + ] + + +def _flush_to_named_tensors(flush) -> list[tuple[str, torch.Tensor]]: + """Package one DeltaFlush exactly like DeltaCheckpointEngine.update_weights_via_server.""" + spec = { + "encoding": flush.encoding, + "params": [vars(p) for p in flush.params], + "checksum": int(flush.checksum), + } + spec_t = torch.frombuffer(bytearray(json.dumps(spec).encode()), dtype=torch.uint8) + return [ + ("__delta_spec__", spec_t), + ("__positions__", flush.positions_cpu.clone()), + ("__values__", flush.values_gpu.clone()), + ] + + +@pytest.mark.parametrize("encoding", ["indices", "deltas"]) +def test_masked_apply_bit_identical(encoding: str): + state = DeltaState() + named = _make_named() + list(iter_delta_flushes(iter(named), state, encoding=encoding, bucket_bytes=1 << 20)) # seed + + model = _FakeModel(named) + + # Mutate a sparse set of positions in each tensor. + new_named = [] + for name, t in named: + new = t.clone() + flat = new.view(-1) + idx = torch.tensor([1, 17, 200, 511], dtype=torch.int64) % flat.numel() + flat[idx] = flat[idx] + 0.5 + new_named.append((name, new)) + + flushes = list(iter_delta_flushes(iter(new_named), state, encoding=encoding, bucket_bytes=1 << 20)) + assert flushes + for flush in flushes: + apply_delta(model, _flush_to_named_tensors(flush)) + + for name, expected in new_named: + got = model.params[name] + assert torch.equal(got.view(torch.int16), expected.view(torch.int16)), f"{name} not bit-identical" + + +def test_untouched_positions_preserved(): + """Positions absent from the delta must keep the model's LIVE values (not the + trainer snapshot's) -- proves the apply is masked, not a full overwrite.""" + state = DeltaState() + named = _make_named() + list(iter_delta_flushes(iter(named), state, encoding="indices", bucket_bytes=1 << 20)) + + model = _FakeModel(named) + # Poison one untouched position in the live model; a full overwrite would revert it. + sentinel_name = named[0][0] + model.params[sentinel_name].view(-1)[3] = 42.0 + + new_named = [] + for name, t in named: + new = t.clone() + new.view(-1)[7] = new.view(-1)[7] + 1.0 # change only position 7 + new_named.append((name, new)) + + flushes = list(iter_delta_flushes(iter(new_named), state, encoding="indices", bucket_bytes=1 << 20)) + for flush in flushes: + apply_delta(model, _flush_to_named_tensors(flush)) + + live = model.params[sentinel_name].view(-1) + assert live[3].item() == 42.0, "masked apply must not touch positions outside the delta" + assert live[7] == new_named[0][1].view(-1)[7] + + +def test_checksum_mismatch_raises(): + state = DeltaState() + named = _make_named() + list(iter_delta_flushes(iter(named), state, encoding="indices", bucket_bytes=1 << 20)) + new_named = [(n, t + 0.5) for n, t in named] + flushes = list(iter_delta_flushes(iter(new_named), state, encoding="indices", bucket_bytes=1 << 20)) + named_tensors = _flush_to_named_tensors(flushes[0]) + named_tensors[2][1].view(torch.uint8)[0] ^= 0xFF # corrupt one value byte + with pytest.raises(RuntimeError, match="checksum"): + apply_delta(_FakeModel(named), named_tensors) + + +def test_dense_flush_applies_full_tensors(): + """Dense (first-sync) flushes carry values only; the loader must apply them verbatim.""" + named = _make_named() + model = _FakeModel([(n, torch.zeros_like(t)) for n, t in named]) # dummy init + + params, pieces, val_off = [], [], 0 + for name, t in named: + flat = t.contiguous().view(-1) + params.append({ + "name": name, "dtype": str(t.dtype).replace("torch.", ""), "shape": list(t.shape), + "pos_start": 0, "pos_end": 0, "pos_width": 4, + "val_start": val_off, "val_end": val_off + flat.numel(), + }) + pieces.append(flat) + val_off += flat.numel() + values = torch.cat(pieces) + + from verl.checkpoint_engine.delta_sync.encode import checksum + + spec = {"encoding": "dense", "params": params, + "checksum": int(checksum(torch.empty(0, dtype=torch.uint8), values))} + spec_t = torch.frombuffer(bytearray(json.dumps(spec).encode()), dtype=torch.uint8) + apply_delta(model, [("__delta_spec__", spec_t), ("__values__", values)]) + + for name, expected in named: + assert torch.equal(model.params[name].view(torch.int16), expected.view(torch.int16)), name diff --git a/tests/checkpoint_engine/test_sharded_delta.py b/tests/checkpoint_engine/test_sharded_delta.py new file mode 100644 index 00000000000..dde4fa05f7a --- /dev/null +++ b/tests/checkpoint_engine/test_sharded_delta.py @@ -0,0 +1,65 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU unit tests for the sharded-delta primitives. + +The full sharded path (DTensor shards + gather-v across ranks vs the full-gather diff) is +validated bit-identically in a multi-GPU check; see +``tests/checkpoint_engine/sharded_delta_multigpu_check.py`` (run with torchrun). These +tests cover the process-local pieces that CI can run without a process group. +""" + +from __future__ import annotations + +import pytest +import torch + +from verl.checkpoint_engine.delta_sync.sharded import local_shard_view, shard_delta_indices + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_shard_delta_indices_matches_bytewise_diff(dtype): + torch.manual_seed(0) + # A "shard" of some parameter, whose flat start in the full param is `offset`. + shard = torch.randn(1000, dtype=dtype) + new = shard.clone() + changed = torch.tensor([3, 17, 500, 999], dtype=torch.int64) + new[changed] = new[changed] + 0.5 + offset = 4096 # this shard begins at flat position 4096 within the full param + + gidx, gval = shard_delta_indices(new, shard, offset) + + # positions are (offset + local changed index), bytewise-exact values + assert torch.equal(gidx.sort().values, (changed + offset).sort().values) + order = torch.argsort(gidx) + got_pos = (gidx[order] - offset).to(torch.int64) + assert torch.equal(gval[order].view(torch.int16 if dtype == torch.bfloat16 else torch.int32), + new[got_pos].view(torch.int16 if dtype == torch.bfloat16 else torch.int32)) + + +def test_shard_delta_indices_no_change_is_empty(): + shard = torch.randn(256, dtype=torch.bfloat16) + gidx, gval = shard_delta_indices(shard.clone(), shard, offset=0) + assert gidx.numel() == 0 + assert gval.numel() == 0 + + +def test_local_shard_view_plain_tensor(): + # A non-DTensor (replicated / unsharded) param: whole tensor is local, offset 0. + t = torch.randn(64, 8, dtype=torch.bfloat16) + local, offset, contributes = local_shard_view(t) + assert offset == 0 + assert local.shape == (64 * 8,) + assert torch.equal(local, t.reshape(-1)) + # outside a process group, rank 0 is assumed -> contributes + assert contributes is True diff --git a/verl/checkpoint_engine/__init__.py b/verl/checkpoint_engine/__init__.py index e23036b921d..45c4b029356 100644 --- a/verl/checkpoint_engine/__init__.py +++ b/verl/checkpoint_engine/__init__.py @@ -64,3 +64,11 @@ __all__ += ["MooncakeCheckpointEngine"] except ImportError: MooncakeCheckpointEngine = None + +try: + from .delta_checkpoint_engine import DeltaCheckpointEngine, DeltaShardedCheckpointEngine + + __all__ += ["DeltaCheckpointEngine", "DeltaShardedCheckpointEngine"] +except ImportError: + DeltaCheckpointEngine = None + DeltaShardedCheckpointEngine = None diff --git a/verl/checkpoint_engine/base.py b/verl/checkpoint_engine/base.py index e19c9f96684..a2af5051667 100644 --- a/verl/checkpoint_engine/base.py +++ b/verl/checkpoint_engine/base.py @@ -170,6 +170,16 @@ def finalize(self): """ raise NotImplementedError + def pop_sync_metrics(self) -> dict: + """Return and clear metrics describing the last completed weight sync. + + Backends that track per-sync statistics (e.g. the delta engines' changed + ratio and wire payload size) override this. The sender worker returns it + from ``update_weights`` so the trainer can merge the stats into that + step's metrics (and thus into wandb). + """ + return {} + @abstractmethod async def send_weights( self, @@ -321,6 +331,11 @@ def __init__( @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False) async def update_weights(self, global_steps: int = None): + # Engines that push weights straight into the inference server (e.g. delta's + # sparse in-place apply) drive the server themselves instead of yielding tensors. + if hasattr(self.checkpoint_engine, "update_weights_via_server"): + await self.checkpoint_engine.update_weights_via_server(self.server_adapter, global_steps=global_steps) + return weights = self.checkpoint_engine.receive_weights(global_steps=global_steps) await self.server_adapter.update_weights(weights, global_steps=global_steps) @@ -477,7 +492,7 @@ async def update_weights(self, global_steps: int = None): # 0. update weights for sync training with colocated actor and rollout if self.backend == "naive": ray.get(self.actor_wg.update_weights(global_steps=global_steps, mode=self.backend)) - return + return {} # 1. abort and save all unfinished requests for partial rollout await self.abort_replicas() @@ -496,10 +511,16 @@ async def update_weights(self, global_steps: int = None): self.build_process_group(rollout) # 5. update weights of all workers - ray.get( + results = ray.get( actor_wg.update_weights(global_steps=global_steps, mode=self.backend) + rollout.update_weights(global_steps=global_steps) ) + # The sender workers return the engine's per-sync metrics (empty for + # backends that don't track any); merge and hand them to the trainer. + sync_metrics: dict = {} + for result in results[: actor_wg.world_size]: + if isinstance(result, dict): + sync_metrics.update(result) # 6. finalize all workers ray.get( @@ -513,6 +534,8 @@ async def update_weights(self, global_steps: int = None): # 8. resume all unfinished requests for partial rollout await self.resume_generation_replicas() + return sync_metrics + async def split_weight_chunks( weights: Generator[tuple[str, torch.Tensor], None, None], bucket_size: int diff --git a/verl/checkpoint_engine/delta_checkpoint_engine.py b/verl/checkpoint_engine/delta_checkpoint_engine.py new file mode 100644 index 00000000000..41e1e667316 --- /dev/null +++ b/verl/checkpoint_engine/delta_checkpoint_engine.py @@ -0,0 +1,625 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Delta weight-sync checkpoint engine (NCCL transport) for DISAGGREGATED rollout. + +Puts the delta on the trainer->rollout wire: the trainer byte-diffs against a +pinned-CPU snapshot and broadcasts only the changed ``(position, value)`` pairs +over the same ``ray.util.collective`` NCCL group the full-weight +:class:`NCCLCheckpointEngine` uses (actor rank0 -> rollout CheckpointEngineWorkers). +Each rollout worker then hands its local copy of the sparse payload to its +colocated SGLang TP worker via same-GPU ``update_weights_from_tensor`` IPC, where +the verl-shipped :mod:`.delta_sync.sglang_loader` (registered through SGLang's +stock ``--custom-weight-loader`` hook — no SGLang fork or patch needed) decodes +and masked-applies it *in place* onto 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. + +The first sync broadcasts a full delta (every position) so a dummy-initialized +rollout gets a correct base; subsequent syncs are sparse. +""" + +from __future__ import annotations + +import json +import logging +import os +from typing import Generator +from unittest.mock import patch + +import ray.util.collective as collective +import torch +import zmq + +with patch("importlib.metadata.distributions", return_value=[]): + import cupy as cp + +from .delta_sync import DeltaState, iter_delta_flushes +from .delta_sync.encode import DeltaParam, checksum as _checksum +from .delta_sync.sglang_loader import LOADER_FQN +from .delta_sync.sharded import gather_dense_to_rank0, gather_v_to_rank0, shard_delta_indices +from .delta_sync.wrapper import DeltaFlush + +from .base import CheckpointEngineRegistry +from .nccl_checkpoint_engine import MasterMetadata, NCCLCheckpointEngine + +logger = logging.getLogger(__name__) + + +@CheckpointEngineRegistry.register("delta") +class DeltaCheckpointEngine(NCCLCheckpointEngine): + """NCCL delta transport. Reuses NCCLCheckpointEngine's group/zmq machinery; + overrides send/receive to move only changed positions+values.""" + + # Cap on changed elements per DeltaParam entry. The receiver-side decode + # densifies per entry with an int64 index transient (8 B/element), so an + # uncapped entry (e.g. a 7B model's whole embedding on the full seed, ~545M + # elements) would spike several GiB at once. Oversized per-param deltas are + # sliced into multiple entries (the masked apply is sequential, so splitting + # is transparent); 64M elements bounds the transient to ~512 MiB. + MAX_ENTRY_ELEMS = 64 << 20 + + def __init__(self, *args, encoding: str = "indices", **kwargs) -> None: + super().__init__(*args, **kwargs) + self.encoding = encoding + self._state = DeltaState() # trainer-side snapshot for diffing + self._sync_metrics: dict = {} # rank0 sender stats for the last sync (see pop_sync_metrics) + + def prepare(self) -> MasterMetadata | None: + # Delta broadcasts small per-flush buffers directly, so skip the parent's + # 2 * bucket_size fixed buffers. Still hand back the master zmq endpoint + # that build_topology() distributes to the rollout workers. + return MasterMetadata(zmq_ip=self.ip, zmq_port=self.listen_port) if self.is_master else None + + # ---- trainer side ---- + # ---- shared STREAMING wire ---- + # Broadcast each flush the moment it is produced and free it, instead of materializing every + # flush up front. Peak device memory stays ~2 buckets (like NCCLCheckpointEngine's send/recv + # buffers) rather than the whole model -- required for large models where the first (full-seed) + # sync would otherwise hold the entire delta on rank 0. Wire: one zmq manifest + NCCL broadcast + # per flush, with an ``is_last`` flag so the receiver loops until the stream ends. Each + # CheckpointEngineWorker then hands its local copy of the sparse payload to its colocated + # SGLang TP worker (same-GPU IPC), where the verl-shipped custom weight loader applies it + # in place -- no full-model staging anywhere on the rollout side. + def _publish_flush(self, flush: DeltaFlush, first: bool, is_last: bool) -> None: + meta = { + "is_full": first, + "encoding": self.encoding, + "is_last": is_last, + "terminal_empty": False, + "pos_numel": int(flush.positions_cpu.numel()), + "val_numel": int(flush.values_gpu.numel()), + "val_dtype": str(flush.values_gpu.dtype).replace("torch.", ""), + "spec": { + "encoding": self.encoding, + "params": [vars(p) for p in flush.params], + "checksum": int(flush.checksum), + }, + } + self.socket.send_string(self.topic, flags=zmq.SNDMORE) + self.socket.send_pyobj(meta) + pos_u8 = flush.positions_cpu.to("cuda", non_blocking=True).contiguous().view(torch.uint8) + val_u8 = flush.values_gpu.contiguous().view(torch.uint8) + pos_cp = cp.empty(pos_u8.numel(), dtype=cp.uint8) + val_cp = cp.empty(val_u8.numel(), dtype=cp.uint8) + pos_cp[:] = cp.asarray(pos_u8) + val_cp[:] = cp.asarray(val_u8) + collective.broadcast(pos_cp, src_rank=0, group_name=self.group_name) + collective.broadcast(val_cp, src_rank=0, group_name=self.group_name) + + def _publish_dense_flush(self, params: list[DeltaParam], values: torch.Tensor, is_last: bool) -> None: + """Publish a dense (full-coverage, positions-free) flush -- used by the first sync.""" + values = values.contiguous() + empty_pos = torch.empty(0, dtype=torch.uint8, device=values.device) + meta = { + "is_full": True, + "encoding": "dense", + "is_last": is_last, + "terminal_empty": False, + "pos_numel": 0, + "val_numel": int(values.numel()), + "val_dtype": str(values.dtype).replace("torch.", ""), + "spec": { + "encoding": "dense", + "params": [vars(p) for p in params], + "checksum": int(_checksum(empty_pos, values)), + }, + } + self.socket.send_string(self.topic, flags=zmq.SNDMORE) + self.socket.send_pyobj(meta) + val_u8 = values.view(torch.uint8) + val_cp = cp.empty(val_u8.numel(), dtype=cp.uint8) + val_cp[:] = cp.asarray(val_u8) + collective.broadcast(val_cp, src_rank=0, group_name=self.group_name) + + def _publish_terminal(self, first: bool) -> None: + """End-of-stream marker when zero flushes were produced (no broadcast, just a signal).""" + meta = {"is_full": first, "encoding": self.encoding, "is_last": True, "terminal_empty": True} + self.socket.send_string(self.topic, flags=zmq.SNDMORE) + self.socket.send_pyobj(meta) + + def pop_sync_metrics(self) -> dict: + metrics, self._sync_metrics = self._sync_metrics, {} + return metrics + + def _stream_flushes(self, flush_iter, first: bool, global_steps, tag: str) -> tuple[int, int, int]: + """Stream flushes with a 1-flush lookahead so the final flush carries ``is_last``; each flush + is freed right after it is broadcast, bounding peak memory to ~2 flushes. + + Returns ``(n_flushes, changed_elems, wire_bytes)`` for sync metrics.""" + pending = None + n = 0 + total = 0 + wire_bytes = 0 + for f in flush_iter: + if pending is not None: + self._publish_flush(pending, first, is_last=False) + n += 1 + total += int(pending.values_gpu.numel()) + wire_bytes += int(pending.positions_cpu.numel()) + int(pending.values_gpu.nbytes) + pending = f + if pending is not None: + self._publish_flush(pending, first, is_last=True) + n += 1 + total += int(pending.values_gpu.numel()) + wire_bytes += int(pending.positions_cpu.numel()) + int(pending.values_gpu.nbytes) + else: + self._publish_terminal(first) + logger.info( + "delta-nccl send v=%s %s flushes=%d nnz=%d (streamed)", + global_steps, tag, n, total, + ) + return n, total, wire_bytes + + async def send_weights( + self, + weights: Generator[tuple[str, torch.Tensor], None, None], + global_steps: int | None = None, + ): + assert self.rank <= 0, "Trainer workers other than rank 0 should not send weights." + if self.rank < 0: + # Non-src actor ranks must still iterate to drive the FSDP all-gather. + for _name, _w in weights: + pass + return + + if not self._state.seeded: + # First sync: dense, single streaming pass -- each full tensor is + # snapshotted and sent raw (values only, no positions), so nothing + # materializes the whole model and the init semantics are explicit. + bucket: list = [] + bucket_bytes = 0 + pending = None + n_flushes = 0 + total_elems = 0 + wire_bytes = 0 + + def _emit(is_last): + nonlocal pending, n_flushes, wire_bytes + if pending is not None: + self._publish_dense_flush(pending[0], pending[1], is_last=is_last) + n_flushes += 1 + wire_bytes += int(pending[1].nbytes) + pending = None + + def _seal(): + nonlocal bucket, bucket_bytes, pending + if not bucket: + return + _emit(is_last=False) + params = [] + val_off = 0 + for name, dtype_str, shape, flat in bucket: + n = int(flat.numel()) + params.append( + DeltaParam(name=name, dtype=dtype_str, shape=list(shape), + pos_start=0, pos_end=0, pos_width=4, + val_start=val_off, val_end=val_off + n) + ) + val_off += n + values = torch.cat([flat for *_, flat in bucket]) + pending = (params, values) + bucket = [] + bucket_bytes = 0 + + for name, tensor in weights: + tensor = tensor.detach() + self._state.seed_param(name, tensor) + flat = tensor.contiguous().view(-1) + total_elems += int(flat.numel()) + bucket.append((name, str(tensor.dtype).replace("torch.", ""), list(tensor.shape), flat)) + bucket_bytes += int(flat.numel()) * flat.element_size() + if bucket_bytes >= self.bucket_size: + _seal() + _seal() + if pending is not None: + _emit(is_last=True) + else: + self._publish_terminal(True) + if total_elems: + self._sync_metrics = { + "checkpoint_engine/changed_ratio": 1.0, + "checkpoint_engine/changed_elems": float(total_elems), + "checkpoint_engine/payload_mbytes": wire_bytes / (1 << 20), + "checkpoint_engine/flushes": float(n_flushes), + } + logger.info("delta-nccl send v=%s DENSE-SEED flushes=%d elems=%d", global_steps, n_flushes, total_elems) + return + + first = False + weights_iter = weights + + total_elems = 0 + + def _counted(it): + nonlocal total_elems + for name, tensor in it: + total_elems += tensor.numel() + yield name, tensor + + flush_gen = iter_delta_flushes( + _counted(weights_iter), self._state, encoding=self.encoding, bucket_bytes=self.bucket_size + ) + n_flushes, changed, wire_bytes = self._stream_flushes( + flush_gen, first, global_steps, "FULL" if first else "delta" + ) + if total_elems: + self._sync_metrics = { + "checkpoint_engine/changed_ratio": changed / total_elems, + "checkpoint_engine/changed_elems": float(changed), + "checkpoint_engine/payload_mbytes": wire_bytes / (1 << 20), + "checkpoint_engine/flushes": float(n_flushes), + } + + # ---- rollout worker side ---- + def receive_weights(self, global_steps: int | None = None): + raise RuntimeError( + "delta engine applies weights inside SGLang via update_weights_via_server; " + "it does not yield tensors to the server adapter" + ) + + async def update_weights_via_server(self, server_adapter, global_steps: int | None = None) -> None: + """Rollout-side apply loop: hand each sparse flush to the colocated SGLang worker. + + Every CheckpointEngineWorker receives the broadcast payload into its own GPU buffer + (one bucket, freed right after dispatch) and forwards it — same-GPU CUDA IPC via + ``update_weights_from_tensor`` — to its SGLang TP worker, where the verl-shipped + ``sglang_loader.apply_delta`` (registered through ``--custom-weight-loader``) + decodes and masked-applies it in place. No full-model mirror is staged anywhere: + SGLang's own live weights are the base. + """ + assert self.rank > 0, "Rank 0 should not receive weights." + await server_adapter._init_server_adapter() + engine = getattr(server_adapter, "_engine", None) + assert getattr(server_adapter, "_pd_role", None) is None, ( + "delta checkpoint engine does not support PD disaggregation" + ) + applied = 0 + while True: + self.socket.recv_string() + meta = self.socket.recv_pyobj() + if meta.get("terminal_empty"): + break + + dense = meta.get("encoding") == "dense" + val_dtype = getattr(torch, meta["val_dtype"]) + elem = torch.empty(0, dtype=val_dtype).element_size() + val_u8 = torch.empty(meta["val_numel"] * elem, dtype=torch.uint8, device="cuda") + if dense: + pos = None + collective.broadcast(val_u8, src_rank=0, group_name=self.group_name) + else: + pos = torch.empty(meta["pos_numel"], dtype=torch.uint8, device="cuda") + collective.broadcast(pos, src_rank=0, group_name=self.group_name) + collective.broadcast(val_u8, src_rank=0, group_name=self.group_name) + val = val_u8.view(val_dtype) + + spec_bytes = json.dumps(meta["spec"]).encode() + spec_t = torch.frombuffer(bytearray(spec_bytes), dtype=torch.uint8).to("cuda") + named = [("__delta_spec__", spec_t), ("__values__", val)] + if pos is not None: + named.insert(1, ("__positions__", pos)) + await self._dispatch_flush_to_sglang( + server_adapter, + engine, + named, + flush_cache=bool(meta["is_last"]), + ) + applied += 1 + del pos, val_u8, val, spec_t + if meta["is_last"]: + break + + if engine is not None and server_adapter._is_server_tp_leader() and global_steps is not None: + await server_adapter.server_actor.set_global_steps.remote(global_steps) + logger.info("delta apply v=%s flushes=%d (in-place via sglang loader)", global_steps, applied) + + @staticmethod + async def _dispatch_flush_to_sglang(server_adapter, engine, params_batch, flush_cache: bool) -> None: + """Hand one sparse flush to the colocated SGLang server via same-GPU CUDA IPC. + + SPMD across the replica's CheckpointEngineWorkers: every rank serializes its local + GPU copy, TP0 gathers the IPC handles and posts one ``update_weights_from_tensor`` + with ``load_format`` pointing at the verl loader. Same flow as + ``sglang.srt.weight_sync.utils.update_weights``, inlined so the radix cache is + flushed only on the stream's last flush instead of on every bucket (the request's + ``flush_cache`` defaults to True). Checksum is re-verified inside each SGLang + worker (fail loud). + """ + import torch.distributed as dist + from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput + from sglang.srt.model_executor.model_runner import LocalSerializedTensor + from sglang.srt.utils import MultiprocessingSerializer + from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions + + monkey_patch_torch_reductions() + mesh = server_adapter.device_mesh["infer_tp"] + tp_size = mesh.mesh.size()[0] + serialized = [(name, MultiprocessingSerializer.serialize(t.detach())) for name, t in params_batch] + + gathered = [None for _ in range(tp_size)] if mesh.get_local_rank() == 0 else None + dist.gather_object( + obj=serialized, + object_gather_list=gathered, + dst=mesh.mesh.tolist()[0], + group=mesh.get_group(), + ) + if mesh.get_local_rank() != 0: + return + + named_tensors = [ + (group[0][0], LocalSerializedTensor(values=[part[1] for part in group])) + for group in zip(*gathered, strict=True) + ] + req = UpdateWeightsFromTensorReqInput( + serialized_named_tensors=[ + MultiprocessingSerializer.serialize(named_tensors) for _ in range(tp_size) + ], + load_format=LOADER_FQN, + flush_cache=flush_cache, + ) + await engine.update_weights_from_tensor(req) + + +@CheckpointEngineRegistry.register("delta_sharded") +class DeltaShardedCheckpointEngine(DeltaCheckpointEngine): + """Delta over NCCL, but the diff is computed on each rank's local FSDP shard. + + Instead of ``full_tensor()``-ing every parameter and diffing on rank 0 (parent), each + actor rank keeps a pinned-CPU snapshot of only *its* shard, byte-diffs the shard, and + only the changed ``(within-param position, value)`` pairs are gathered to rank 0 -- so + the all-gather volume drops to the sparsity ratio and rank 0 no longer holds a + full-model snapshot. The assembled result is bit-identical to the parent's diff, so the + wire and the receiver (:meth:`DeltaCheckpointEngine.update_weights_via_server`) are + reused unchanged. + + ``send_weights`` here expects the SHARDED generator ``get_per_tensor_param_shard()`` + (``(name, local_flat_shard, within_param_offset, full_numel, full_shape, contributes)``) + rather than the full-tensor generator. + """ + + def __init__(self, *args, encoding: str = "indices", **kwargs) -> None: + super().__init__(*args, encoding=encoding, **kwargs) + self._shard_snap: dict[str, torch.Tensor] = {} # name -> pinned-CPU shard snapshot + self._shard_seeded = False + + def _assemble_flush(self, per_param: list) -> DeltaFlush: + """Build one DeltaFlush (indices encoding) from rank 0's gathered per-param deltas. + + ``per_param``: list of ``(name, dtype_str, shape, global_idx, values)`` where + ``global_idx`` are within-parameter flat positions (== what the receiver decodes). + + Positions stay on the GPU end to end (int32 pieces -> one cat -> uint8 view); + the wire broadcasts from the GPU anyway, and a host round-trip here + (``.cpu().numpy().tobytes()`` + join) dominated the whole send at scale + (~2.4s/sync at 7B steady state, ~83s on the full seed). + """ + idx_pieces: list[torch.Tensor] = [] + val_pieces: list[torch.Tensor] = [] + params: list[DeltaParam] = [] + pos_off = val_off = 0 + for name, dtype_str, shape, idx, val in per_param: + nnz = int(idx.numel()) + idx_pieces.append(idx.to(torch.int32)) + val_pieces.append(val) + params.append( + DeltaParam(name=name, dtype=dtype_str, shape=list(shape), + pos_start=pos_off, pos_end=pos_off + nnz * 4, pos_width=4, + val_start=val_off, val_end=val_off + nnz) + ) + pos_off += nnz * 4 + val_off += nnz + + values_gpu = torch.cat(val_pieces) if val_pieces else torch.empty(0, dtype=self.rollout_dtype, device="cuda") + positions_u8 = ( + torch.cat(idx_pieces).contiguous().view(torch.uint8) + if idx_pieces + else torch.empty(0, dtype=torch.uint8, device=values_gpu.device) + ) + cks = _checksum(positions_u8, values_gpu) + return DeltaFlush(encoding=self.encoding, params=params, + positions_cpu=positions_u8, values_gpu=values_gpu, checksum=cks) + + def _send_dense_seed(self, weights, global_steps=None): + """First sync: assemble and broadcast the raw weights, bucketed, positions-free. + + Explicit dense semantics instead of the previous bitflip-forced full diff: + no per-element indices on the wire (values only), no whole-parameter + (idx, val) spike on rank 0 -- its peak is one assembled parameter plus a + bucket -- and the snapshot is populated as the stream goes. + """ + is_r0 = self.is_master + bucket: list = [] # (name, dtype_str, full_shape, flat_full_tensor) + bucket_bytes = 0 + pending = None # (params, values) awaiting emission (1-flush lookahead for is_last) + n_flushes = 0 + total_elems = 0 + wire_bytes = 0 + + def _emit(is_last): + nonlocal pending, n_flushes, wire_bytes + if pending is not None: + self._publish_dense_flush(pending[0], pending[1], is_last=is_last) + n_flushes += 1 + wire_bytes += int(pending[1].nbytes) + pending = None + + def _seal(): + nonlocal bucket, bucket_bytes, pending + if not bucket: + return + _emit(is_last=False) + params = [] + val_off = 0 + for name, dtype_str, full_shape, flat in bucket: + n = int(flat.numel()) + params.append( + DeltaParam(name=name, dtype=dtype_str, shape=list(full_shape), + pos_start=0, pos_end=0, pos_width=4, + val_start=val_off, val_end=val_off + n) + ) + val_off += n + values = torch.cat([flat for *_, flat in bucket]) + pending = (params, values) + bucket = [] + bucket_bytes = 0 + + for name, local, offset, full_numel, full_shape, contributes in weights: + local = local.detach().contiguous().view(-1) + snap = self._shard_snap.get(name) + if snap is None or snap.numel() != local.numel(): + snap = torch.empty_like(local, device="cpu", pin_memory=True) + snap.copy_(local, non_blocking=True) + self._shard_snap[name] = snap + + shard = local if contributes else torch.empty(0, dtype=local.dtype, device=local.device) + full = gather_dense_to_rank0(shard, offset if contributes else 0, full_numel) + if is_r0: + total_elems += int(full_numel) + bucket.append((name, str(local.dtype).replace("torch.", ""), list(full_shape), full)) + bucket_bytes += int(full.nbytes) + if bucket_bytes >= self.bucket_size: + _seal() + + self._shard_seeded = True + if not is_r0: + return + _seal() + if pending is not None: + _emit(is_last=True) + else: + self._publish_terminal(True) + if total_elems: + self._sync_metrics = { + "checkpoint_engine/changed_ratio": 1.0, + "checkpoint_engine/changed_elems": float(total_elems), + "checkpoint_engine/payload_mbytes": wire_bytes / (1 << 20), + "checkpoint_engine/flushes": float(n_flushes), + } + logger.info( + "delta-sharded send v=%s DENSE-SEED flushes=%d elems=%d", + global_steps, n_flushes, total_elems, + ) + + async def send_weights(self, weights, global_steps=None): + # All actor ranks participate (gather-v is collective); only torch rank 0 broadcasts. + # rank 0 accumulates the gathered per-param deltas into bucket_size-sized flushes and streams + # each one as soon as it fills (then frees it), so peak memory is ~2 buckets rather than the + # whole model. + assert self.rank <= 0, "Trainer workers other than rank 0 should not send weights." + if not self._shard_seeded: + return self._send_dense_seed(weights, global_steps) + is_r0 = self.is_master + first = False # the dense first sync is handled by _send_dense_seed + bucket: list = [] + bucket_bytes = 0 + pending = None # a DeltaFlush awaiting emission (1-flush lookahead so the last flush is is_last) + n_flushes = 0 + changed_elems = 0 + total_elems = 0 + wire_bytes = 0 + + def _emit(is_last): + nonlocal pending, n_flushes + if pending is not None: + self._publish_flush(pending, first, is_last=is_last) + n_flushes += 1 + pending = None + + def _seal(): + nonlocal bucket, bucket_bytes, pending + if not bucket: + return + _emit(is_last=False) # another bucket follows, so the prior one is not the last + pending = self._assemble_flush(bucket) + bucket = [] + bucket_bytes = 0 + + for name, local, offset, _full_numel, full_shape, contributes in weights: + local = local.detach().contiguous().view(-1) + snap = self._shard_snap.get(name) + if snap is None or snap.numel() != local.numel(): + snap = torch.empty_like(local, device="cpu", pin_memory=True) + if contributes: + base = snap.to(local.device, non_blocking=True) + gidx, gval = shard_delta_indices(local, base, offset) + else: + # replicated param owned by another rank; contribute nothing but keep lockstep. + gidx = torch.empty(0, dtype=torch.int64, device=local.device) + gval = torch.empty(0, dtype=local.dtype, device=local.device) + snap.copy_(local, non_blocking=True) # update snapshot to current shard + self._shard_snap[name] = snap + + aidx, aval = gather_v_to_rank0(gidx, gval) + if is_r0: + total_elems += int(_full_numel) + if is_r0 and aidx is not None and aidx.numel() > 0: + changed_elems += int(aidx.numel()) + wire_bytes += int(aidx.numel()) * (4 + aval.element_size()) + # Slice oversized per-param deltas so one entry never exceeds + # MAX_ENTRY_ELEMS (bounds the receiver-side decode transient). + for s in range(0, aidx.numel(), self.MAX_ENTRY_ELEMS): + e = min(s + self.MAX_ENTRY_ELEMS, aidx.numel()) + bucket.append( + (name, str(local.dtype).replace("torch.", ""), list(full_shape), aidx[s:e], aval[s:e]) + ) + bucket_bytes += (e - s) * 8 + (e - s) * aval.element_size() + if bucket_bytes >= self.bucket_size: + _seal() + + self._shard_seeded = True + if os.environ.get("VERL_DELTA_DEBUG"): + pinned = sum(t.numel() * t.element_size() for t in self._shard_snap.values()) + print( + f"[delta-debug] engine rank={self.rank} SNAPSHOT: {len(self._shard_snap)} pinned-CPU shard snaps, " + f"{pinned / (1 << 20):.1f} MiB total (expect ~model_bytes/fsdp_size per rank when truly sharded)", + flush=True, + ) + if not is_r0: + return + _seal() # seal the final partial bucket into `pending` + if pending is not None: + _emit(is_last=True) + else: + self._publish_terminal(first) + if total_elems: + self._sync_metrics = { + "checkpoint_engine/changed_ratio": changed_elems / total_elems, + "checkpoint_engine/changed_elems": float(changed_elems), + "checkpoint_engine/payload_mbytes": wire_bytes / (1 << 20), + "checkpoint_engine/flushes": float(n_flushes), + } + logger.info( + "delta-sharded send v=%s %s flushes=%d (streamed)", + global_steps, "FULL" if first else "delta", n_flushes, + ) diff --git a/verl/checkpoint_engine/delta_sync/__init__.py b/verl/checkpoint_engine/delta_sync/__init__.py new file mode 100644 index 00000000000..5004a3aac09 --- /dev/null +++ b/verl/checkpoint_engine/delta_sync/__init__.py @@ -0,0 +1,48 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Delta weight sync primitives for the checkpoint engine. + +This package contains the framework-agnostic core of "send only the elements +that changed" weight sync: the pinned-CPU snapshot, bytewise diff, encode, and +bucket logic. :class:`~verl.checkpoint_engine.delta_checkpoint_engine.DeltaCheckpointEngine` +consumes these flushes and broadcasts them over NCCL; the rollout worker decodes +them back into full tensors locally. + +Design follows THUDM/slime's delta-sync implementation +(``slime/backends/megatron_utils/update_weight/update_weight_from_distributed_delta.py``). +""" + +from .delta_state import DeltaState, ParamDiff +from .encode import ( + DeltaBucket, + DeltaEncodingName, + EncodedChunk, + bytewise_diff_mask, + checksum, + encode_chunk, +) +from .wrapper import DeltaFlush, iter_delta_flushes + +__all__ = [ + "DeltaBucket", + "DeltaEncodingName", + "DeltaFlush", + "DeltaState", + "EncodedChunk", + "ParamDiff", + "bytewise_diff_mask", + "checksum", + "encode_chunk", + "iter_delta_flushes", +] diff --git a/verl/checkpoint_engine/delta_sync/delta_state.py b/verl/checkpoint_engine/delta_sync/delta_state.py new file mode 100644 index 00000000000..eebd8632d0a --- /dev/null +++ b/verl/checkpoint_engine/delta_sync/delta_state.py @@ -0,0 +1,188 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Pinned-CPU snapshot of broadcast weights for delta sync. + +``DeltaState`` holds one pinned-host copy per HF tensor name that has been +broadcast. Each sync prefetches the relevant snapshot tiles to the GPU on a +side stream (overlapped with the previous chunk's compute), the diff/encode +runs on the default stream, and the updated values are copied back to host +on another side stream. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass +class ParamDiff: + """One per-parameter compute output. + + ``values`` is a view of the current GPU tensor (no copy); ``mask`` is a + same-shape bool tensor marking the positions whose bytes differ from the + snapshot. + """ + + name: str + values: torch.Tensor + mask: torch.Tensor + + +class DeltaState: + """Pinned-CPU snapshot store plus side streams for D2H / H2D pipelining. + + Lifecycle: + + * The first sync seeds the snapshot from current model weights and skips + the engine RPC -- the receiver is assumed to have loaded the same HF + checkpoint at init. + * Each subsequent sync: ``prefetch_snapshot`` for chunk N+1 is issued on + the H2D stream before chunk N's compute; ``compute_diffs`` waits on the + prefetch event; ``update_snapshot_async`` writes the just-sent values + back to the pinned host buffer on the D2H stream. + * ``flush_snapshot`` blocks until all enqueued D2H copies have landed. + """ + + def __init__(self) -> None: + self.snapshot: dict[str, torch.Tensor] = {} + self.d2h_stream: torch.cuda.Stream | None = None + self.h2d_stream: torch.cuda.Stream | None = None + self._snapshot_dirty = False + # Pipelining via side streams and pinned host memory is only worthwhile + # (and only available) when CUDA is. The CPU-only path keeps the same + # data flow but degrades to plain blocking copies, which keeps unit + # tests runnable on CPU CI. + self._cuda = torch.cuda.is_available() + + @property + def seeded(self) -> bool: + return len(self.snapshot) > 0 + + def seed(self, named_tensors: list[tuple[str, torch.Tensor]]) -> None: + """Populate the snapshot from the current model state. + + Costs roughly one full D2H of every parameter; expected to be called + exactly once, at the first ``update_weights`` invocation. + """ + for name, tensor in named_tensors: + self._allocate(name, tensor) + for name, tensor in named_tensors: + self.snapshot[name].copy_(tensor.detach(), non_blocking=False) + self._snapshot_dirty = False + + def seed_param(self, name: str, tensor: torch.Tensor) -> None: + """Streaming variant of :meth:`seed`: snapshot a single parameter.""" + self._allocate(name, tensor) + self.snapshot[name].copy_(tensor.detach(), non_blocking=False) + self._snapshot_dirty = False + + def _allocate(self, name: str, tensor: torch.Tensor) -> None: + if name in self.snapshot: + return + # pin_memory requires CUDA; fall back to a regular host buffer in CPU + # CI / smoke tests. + self.snapshot[name] = torch.empty_like( + tensor, device=torch.device("cpu"), pin_memory=self._cuda + ) + + def prefetch_snapshot( + self, named_tensors: list[tuple[str, torch.Tensor]] + ) -> tuple[list[torch.Tensor], "torch.cuda.Event | None"]: + """Start an async H2D copy of snapshot tiles matching ``named_tensors``. + + Returns the GPU-resident previous-step tensors and a CUDA event that + ``compute_diffs`` waits on before reading them. On CPU (no CUDA) the + copy is blocking and the event is ``None``. + """ + if self._cuda: + if self.h2d_stream is None: + self.h2d_stream = torch.cuda.Stream() + prev_gpu: list[torch.Tensor] = [] + with torch.cuda.stream(self.h2d_stream): + for name, tensor in named_tensors: + if name not in self.snapshot: + raise KeyError( + f"missing snapshot for {name!r}; the first update_weights call " + "must seed the snapshot before any chunk is prefetched" + ) + prev_gpu.append( + self.snapshot[name].to( + device=tensor.device, non_blocking=True + ) + ) + event = self.h2d_stream.record_event() + return prev_gpu, event + # CPU fallback: snapshot already lives on host, no async transfer needed. + prev_host: list[torch.Tensor] = [] + for name, tensor in named_tensors: + if name not in self.snapshot: + raise KeyError( + f"missing snapshot for {name!r}; the first update_weights call " + "must seed the snapshot before any chunk is prefetched" + ) + prev_host.append(self.snapshot[name].to(device=tensor.device)) + return prev_host, None + + def compute_diffs( + self, + named_tensors: list[tuple[str, torch.Tensor]], + prefetched: tuple[list[torch.Tensor], "torch.cuda.Event | None"], + ) -> list[ParamDiff]: + """Wait for the prefetched H2D copy, then per-param bytewise diff.""" + from .encode import bytewise_diff_mask # local import avoids cycle + + prev_gpu, event = prefetched + if event is not None: + event.wait() + return [ + ParamDiff(name=name, values=current, mask=bytewise_diff_mask(current, prev)) + for (name, current), prev in zip(named_tensors, prev_gpu, strict=True) + ] + + def update_snapshot_async( + self, named_tensors: list[tuple[str, torch.Tensor]] + ) -> None: + """Enqueue a D2H copy of ``named_tensors`` into the pinned snapshot. + + Non-blocking on CUDA; the caller must invoke :meth:`flush_snapshot` + before the next sync. On CPU the copy is blocking. + """ + if not self._cuda: + for name, tensor in named_tensors: + self._allocate(name, tensor) + self.snapshot[name].copy_(tensor.detach()) + self._snapshot_dirty = False + return + if self.d2h_stream is None: + self.d2h_stream = torch.cuda.Stream() + event = torch.cuda.current_stream().record_event() + with torch.cuda.stream(self.d2h_stream): + self.d2h_stream.wait_event(event) + for name, tensor in named_tensors: + self._allocate(name, tensor) + self.snapshot[name].copy_(tensor.detach(), non_blocking=True) + self._snapshot_dirty = True + + def flush_snapshot(self) -> None: + """Block until every enqueued D2H snapshot copy has landed.""" + if not self._snapshot_dirty: + return + if self._cuda: + if self.d2h_stream is not None: + self.d2h_stream.synchronize() + else: + torch.cuda.synchronize() + self._snapshot_dirty = False diff --git a/verl/checkpoint_engine/delta_sync/encode.py b/verl/checkpoint_engine/delta_sync/encode.py new file mode 100644 index 00000000000..cccd610b001 --- /dev/null +++ b/verl/checkpoint_engine/delta_sync/encode.py @@ -0,0 +1,382 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Position+value encoding and bucketing for delta sync. + +Both encodings share one on-wire layout (a uint8 positions blob plus a +parameter-dtype values tensor with a per-parameter manifest); decoders +dispatch on metadata. + +* ``indices`` -- int32 absolute positions (4 bytes / nnz) +* ``deltas`` -- uint16 gap-deltas with uint32 per-parameter fallback + +Values are sent verbatim in the parameter's dtype regardless of encoding. +""" + +from __future__ import annotations + +import itertools +from dataclasses import dataclass, field, replace +from typing import Literal + +import numpy as np +import torch + +from .delta_state import ParamDiff + +DeltaEncodingName = Literal["indices", "deltas"] + + +# ---------- diff ---------------------------------------------------------- + + +def bytewise_diff_mask(current: torch.Tensor, snapshot: torch.Tensor) -> torch.Tensor: + """Per-element bool mask: True where ``current`` and ``snapshot`` differ. + + Dtype-agnostic via view-as-integer; no arithmetic. + """ + es = current.element_size() + int_dtype = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64}.get(es) + if int_dtype is None: + raise ValueError(f"unsupported element size {es}") + return current.view(int_dtype) != snapshot.view(int_dtype) + + +# ---------- wire format --------------------------------------------------- + + +@dataclass +class DeltaParam: + """Per-parameter manifest entry for a single chunk / bucket. + + Offsets are byte offsets into the surrounding ``__positions__`` blob and + element offsets into the surrounding ``__values__`` tensor. + """ + + name: str + dtype: str + shape: list[int] + pos_start: int + pos_end: int + pos_width: int # 2 or 4 + val_start: int + val_end: int + + +@dataclass +class EncodedChunk: + """One HF chunk after position+value encoding, before bucket merging. + + ``pos_u8`` stays on the values' device (a host round trip here dominated + the whole send at scale); the wire broadcasts from the GPU anyway. + """ + + pos_u8: torch.Tensor + val_tensor: torch.Tensor + params: list[DeltaParam] + nnz: int + + @classmethod + def empty(cls) -> "EncodedChunk": + return cls( + pos_u8=torch.empty(0, dtype=torch.uint8), + val_tensor=torch.empty(0, dtype=torch.bfloat16), + params=[], + nnz=0, + ) + + +# ---------- checksum ------------------------------------------------------ + + +def checksum(positions: torch.Tensor, values: torch.Tensor) -> int: + """Wire-corruption check; sender computes pre-flush, receiver post-recv. + + Uses ``torch.hash_tensor`` (XOR-reduce over uint64 bitcast); one reduction + plus one ``.item()`` sync per argument. + """ + p = int(torch.hash_tensor(positions).item()) if positions.numel() else 0 + v = int(torch.hash_tensor(values).item()) if values.numel() else 0 + return p ^ (v << 1) + + +# ---------- encode -------------------------------------------------------- + + +def _sparse_boundaries( + diffs: list[ParamDiff], +) -> tuple[torch.Tensor, list[int], torch.Tensor, list[int]]: + """One concat -> one nonzero -> one searchsorted -> one tolist(). + + Collapses per-parameter host syncs to a single one per chunk. + """ + device = diffs[0].values.device + sizes = [d.values.numel() for d in diffs] + cum = list(itertools.accumulate(sizes)) + cum_t = torch.tensor(cum, dtype=torch.int64, device=device) + + big_values = torch.cat([d.values.contiguous().view(-1) for d in diffs], dim=0) + big_mask = torch.cat([d.mask.contiguous().view(-1) for d in diffs], dim=0) + big_idx = big_mask.nonzero(as_tuple=False).view(-1) + big_val = big_values[big_idx] + bounds = torch.searchsorted(big_idx, cum_t).tolist() + return big_val, bounds, big_idx, cum + + +def _encode_indices(diffs: list[ParamDiff]) -> EncodedChunk: + if not diffs: + return EncodedChunk.empty() + big_val, bounds, big_idx, cum = _sparse_boundaries(diffs) + pos_pieces: list[torch.Tensor] = [] + val_pieces: list[torch.Tensor] = [] + params: list[DeltaParam] = [] + pos_byte_off = val_off = 0 + prev_b = 0 + prev_param_start = 0 + for i, d in enumerate(diffs): + b = bounds[i] + nnz = b - prev_b + if nnz > 0: + local_idx = (big_idx[prev_b:b] - prev_param_start).to(torch.int32) + pos_pieces.append(local_idx) + val_pieces.append(big_val[prev_b:b]) + params.append( + DeltaParam( + name=d.name, + dtype=str(d.values.dtype).replace("torch.", ""), + shape=list(d.values.shape), + pos_start=pos_byte_off, + pos_end=pos_byte_off + nnz * 4, + pos_width=4, + val_start=val_off, + val_end=val_off + nnz, + ) + ) + pos_byte_off += nnz * 4 + val_off += nnz + prev_b = b + prev_param_start = cum[i] + if not params: + return EncodedChunk.empty() + positions = torch.cat(pos_pieces, dim=0) + values = torch.cat(val_pieces, dim=0) + return EncodedChunk( + pos_u8=positions.contiguous().view(torch.uint8), + val_tensor=values, + params=params, + nnz=val_off, + ) + + +def _encode_deltas(diffs: list[ParamDiff]) -> EncodedChunk: + """Gap-encode sorted positions. + + Store ``idx[k] - idx[k-1] - 1`` with ``idx[-1] := -1`` so the first delta + equals the first index. Each parameter downcasts to uint16 if the max gap + fits, else uint32. At ~2% density on bf16 weights the typical max gap is + ~300, so uint16 normally suffices; the uint32 fallback covers pathological + inputs without correctness risk. The receiver inverts via + ``idx = cumsum(delta + 1) - 1``. + """ + if not diffs: + return EncodedChunk.empty() + big_val, bounds, big_idx, cum = _sparse_boundaries(diffs) + + kept: list[tuple[ParamDiff, int]] = [] + per_param_deltas: list[torch.Tensor] = [] + val_pieces: list[torch.Tensor] = [] + prev_b = 0 + prev_param_start = 0 + for i, d in enumerate(diffs): + b = bounds[i] + nnz = b - prev_b + if nnz > 0: + local_idx = big_idx[prev_b:b] - prev_param_start # int64, sorted + prev = torch.cat( + [ + torch.tensor( + [-1], dtype=local_idx.dtype, device=local_idx.device + ), + local_idx[:-1], + ] + ) + per_param_deltas.append(local_idx - prev - 1) + val_pieces.append(big_val[prev_b:b]) + kept.append((d, nnz)) + prev_b = b + prev_param_start = cum[i] + + if not kept: + return EncodedChunk.empty() + + max_per_param = ( + torch.stack([d.max() for d in per_param_deltas]).cpu().tolist() + ) + pos_u8_pieces: list[torch.Tensor] = [] + pos_byte_off = val_off = 0 + params: list[DeltaParam] = [] + for (d, nnz), deltas, max_d in zip( + kept, per_param_deltas, max_per_param, strict=True + ): + width = 2 if int(max_d) <= 65535 else 4 + int_dtype = torch.int16 if width == 2 else torch.int32 + piece = deltas.to(int_dtype).contiguous().view(torch.uint8) + pos_u8_pieces.append(piece) + params.append( + DeltaParam( + name=d.name, + dtype=str(d.values.dtype).replace("torch.", ""), + shape=list(d.values.shape), + pos_start=pos_byte_off, + pos_end=pos_byte_off + piece.numel(), + pos_width=width, + val_start=val_off, + val_end=val_off + nnz, + ) + ) + pos_byte_off += piece.numel() + val_off += nnz + + values = torch.cat(val_pieces, dim=0) + return EncodedChunk( + pos_u8=torch.cat(pos_u8_pieces, dim=0), + val_tensor=values, + params=params, + nnz=val_off, + ) + + +def encode_chunk(diffs: list[ParamDiff], encoding: DeltaEncodingName) -> EncodedChunk: + """Encode one chunk of per-parameter diffs into an :class:`EncodedChunk`.""" + if encoding == "indices": + return _encode_indices(diffs) + if encoding == "deltas": + return _encode_deltas(diffs) + raise ValueError(f"unsupported delta encoding: {encoding!r}") + + +# ---------- decoder (for tests / receivers that live in Python) ---------- + + +def decode_chunk( + encoding: DeltaEncodingName, + pos_bytes: bytes, + val_tensor: torch.Tensor, + params: list[DeltaParam], +) -> dict[str, torch.Tensor]: + """Reverse of :func:`encode_chunk`. + + Returns full-shape parameter tensors with NaN at unchanged positions. The + rollout worker calls this on each received flush and overwrites only the + non-NaN positions into its full-weight mirror before pushing to the engine. + """ + out: dict[str, torch.Tensor] = {} + pos_np = np.frombuffer(pos_bytes, dtype=np.uint8) + for p in params: + param_dtype = getattr(torch, p.dtype) + numel = 1 + for s in p.shape: + numel *= s + flat = torch.full( + (numel,), float("nan"), dtype=param_dtype, device=val_tensor.device + ) + vals = val_tensor[p.val_start : p.val_end] + chunk_bytes = pos_np[p.pos_start : p.pos_end] + n_elems = len(chunk_bytes) // p.pos_width + if n_elems == 0: + out[p.name] = flat.view(*p.shape) + continue + if p.pos_width == 4: + idx = torch.from_numpy( + chunk_bytes.view(np.int32).copy() + ).to(dtype=torch.int64, device=val_tensor.device) + if encoding != "indices": + idx = (idx + 1).cumsum(dim=0) - 1 + elif p.pos_width == 2: + idx = torch.from_numpy( + chunk_bytes.view(np.uint16).copy() + ).to(dtype=torch.int64, device=val_tensor.device) + idx = (idx + 1).cumsum(dim=0) - 1 + else: + raise ValueError(f"unsupported pos_width {p.pos_width}") + flat.index_copy_(0, idx, vals.to(param_dtype)) + out[p.name] = flat.view(*p.shape) + return out + + +# ---------- bucket -------------------------------------------------------- + + +@dataclass +class DeltaBucket: + """Accumulates encoded chunks for one flush. + + Per-parameter offsets in incoming chunks are rebased into the bucket's + growing position blob and value tensor on ``add``. + """ + + pos_pieces: list[torch.Tensor] = field(default_factory=list) + val_pieces: list[torch.Tensor] = field(default_factory=list) + params: list[DeltaParam] = field(default_factory=list) + pos_total: int = 0 + val_total: int = 0 + byte_size: int = 0 + + @property + def has_updates(self) -> bool: + return bool(self.pos_pieces) + + def should_flush_before_add(self, chunk: EncodedChunk, byte_limit: int) -> bool: + chunk_bytes = ( + chunk.pos_u8.numel() + + chunk.val_tensor.numel() * chunk.val_tensor.element_size() + ) + return self.has_updates and self.byte_size + chunk_bytes > byte_limit + + def add(self, chunk: EncodedChunk) -> None: + for p in chunk.params: + self.params.append( + replace( + p, + pos_start=p.pos_start + self.pos_total, + pos_end=p.pos_end + self.pos_total, + val_start=p.val_start + self.val_total, + val_end=p.val_end + self.val_total, + ) + ) + self.pos_pieces.append(chunk.pos_u8) + self.val_pieces.append(chunk.val_tensor) + self.pos_total += chunk.pos_u8.numel() + self.val_total += chunk.val_tensor.numel() + self.byte_size += ( + chunk.pos_u8.numel() + + chunk.val_tensor.numel() * chunk.val_tensor.element_size() + ) + + def merged_positions(self) -> torch.Tensor: + if not self.pos_pieces: + return torch.empty(0, dtype=torch.uint8) + return torch.cat(self.pos_pieces, dim=0) + + def merged_values(self) -> torch.Tensor: + if not self.val_pieces: + return torch.empty(0, dtype=torch.bfloat16) + return torch.cat(self.val_pieces, dim=0) + + def clear(self) -> None: + self.pos_pieces.clear() + self.val_pieces.clear() + self.params.clear() + self.pos_total = 0 + self.val_total = 0 + self.byte_size = 0 diff --git a/verl/checkpoint_engine/delta_sync/sglang_loader.py b/verl/checkpoint_engine/delta_sync/sglang_loader.py new file mode 100644 index 00000000000..42d5c2fe247 --- /dev/null +++ b/verl/checkpoint_engine/delta_sync/sglang_loader.py @@ -0,0 +1,168 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""In-process sparse delta apply for SGLang, loaded via ``--custom-weight-loader``. + +SGLang's ``update_weights_from_tensor`` supports pluggable loaders: when the +request's ``load_format`` names an import path registered in +``--custom-weight-loader``, SGLang ``dynamic_import``s it **inside every TP +worker process** and calls ``loader(model, named_tensors)``. That is exactly +the hook a sparse delta needs: the delta payload is decoded and applied *in +place* onto SGLang's live weights (masked overwrite of only the changed +positions), so the receiver never stages a full-model mirror anywhere — +peak memory is one decode chunk, independent of model size — and no SGLang +fork or patch is required. + +Wire contract (what the delta checkpoint engine sends per flush): + +* ``__delta_spec__`` — uint8 tensor holding a JSON manifest + ``{"encoding", "params": [DeltaParam-dict...], "checksum"}``. +* ``__positions__`` — uint8 blob of packed positions (per-param slices are + byte offsets ``pos_start:pos_end``; ``indices`` packs little-endian int32 + absolute positions, ``deltas`` packs uint16/uint32 gaps per ``pos_width``). +* ``__values__`` — the changed values in the flush's (uniform) dtype; + per-param slices are element offsets ``val_start:val_end``. + +Register at server launch (verl config):: + + +actor_rollout_ref.rollout.engine_kwargs.sglang.custom_weight_loader='["verl.checkpoint_engine.delta_sync.sglang_loader.apply_delta"]' +""" + +from __future__ import annotations + +import json +import math +from contextlib import contextmanager + +import torch + +# Cap on the densified tensors handed to one model.load_weights call, matching +# SGLang's own delta-apply chunking default. +CHUNK_BYTES = 512 << 20 + +# Import path callers pass as both --custom-weight-loader and load_format. +LOADER_FQN = "verl.checkpoint_engine.delta_sync.sglang_loader.apply_delta" + + +def apply_delta(model, named_tensors) -> None: + """Decode one sparse delta flush and masked-apply it onto ``model`` in place.""" + from .encode import checksum as _checksum + + tensors = dict(named_tensors) + spec = json.loads(bytes(tensors["__delta_spec__"].cpu().numpy().tobytes()).decode()) + values = tensors["__values__"] + positions = tensors.get("__positions__") + if positions is None: # dense flush (first sync) carries values only + positions = torch.empty(0, dtype=torch.uint8, device=values.device) + + got = _checksum(positions, values) + if got != int(spec["checksum"]): + raise RuntimeError( + f"delta checksum mismatch in sglang loader: got {got}, expected {spec['checksum']}; " + "indicates corruption between sender encode and receiver apply" + ) + + if spec["encoding"] == "dense": + _apply_dense(model, spec["params"], values) + return + + encoding = spec["encoding"] + with _masked_copy(): + chunk: list[tuple[str, torch.Tensor]] = [] + chunk_bytes = 0 + for p in spec["params"]: + t = _decode_one(encoding, positions, values, p) + nbytes = t.numel() * t.element_size() + if chunk and chunk_bytes + nbytes > CHUNK_BYTES: + model.load_weights(chunk) + chunk, chunk_bytes = [], 0 + chunk.append((p["name"], t)) + chunk_bytes += nbytes + if chunk: + model.load_weights(chunk) + + +def _apply_dense(model, params: list[dict], values: torch.Tensor) -> None: + """Apply a dense (full-coverage) flush: plain chunked load, no masking needed.""" + chunk: list[tuple[str, torch.Tensor]] = [] + chunk_bytes = 0 + for p in params: + dtype = getattr(torch, p["dtype"]) + t = values[p["val_start"] : p["val_end"]].to(dtype).view(p["shape"]) + nbytes = t.numel() * t.element_size() + if chunk and chunk_bytes + nbytes > CHUNK_BYTES: + model.load_weights(chunk) + chunk, chunk_bytes = [], 0 + chunk.append((p["name"], t)) + chunk_bytes += nbytes + if chunk: + model.load_weights(chunk) + + +def _decode_one(encoding: str, positions: torch.Tensor, values: torch.Tensor, p: dict) -> torch.Tensor: + """Densify one param's sparse delta into a full-shape NaN-masked tensor. + + ``indices`` positions are reinterpreted via an int32 view (8 B/element + transient) rather than a per-byte int64 unpack (32 B/element), so even a + full-seed flush of a large embedding stays within a few GiB. + """ + numel = math.prod(p["shape"]) + dtype = getattr(torch, p["dtype"]) + flat = torch.full((numel,), float("nan"), dtype=dtype, device=values.device) + vals = values[p["val_start"] : p["val_end"]] + if vals.numel() == 0: + return flat.view(p["shape"]) + + pos_b = positions[p["pos_start"] : p["pos_end"]] + if encoding == "indices": + # pos_start is always int32-aligned (each entry packs nnz * 4 bytes). + idx = pos_b.view(torch.int32).to(torch.int64) + elif encoding == "deltas": + w = int(p["pos_width"]) + bv = pos_b.view(-1, w).to(torch.int64) + if w == 2: + gaps = bv[:, 0] | (bv[:, 1] << 8) + else: + gaps = bv[:, 0] | (bv[:, 1] << 8) | (bv[:, 2] << 16) | (bv[:, 3] << 24) + idx = (gaps + 1).cumsum(dim=0) - 1 + else: + raise ValueError(f"unsupported delta encoding: {encoding!r}") + + flat.index_copy_(0, idx, vals.to(dtype)) + return flat.view(p["shape"]) + + +@contextmanager +def _masked_copy(): + """Temporarily make ``Tensor.copy_`` skip NaN positions in the source. + + SGLang's per-model ``load_weights`` ultimately lands on ``param.copy_(loaded)`` + (possibly on a narrowed TP slice). Under this context a NaN-masked source + overwrites only the changed positions; fully dense sources (e.g. the first + full-seed flush) take the original fast path untouched. + """ + orig_copy = torch.Tensor.copy_ + + def masked_copy_(self, src, *args, **kwargs): + if isinstance(src, torch.Tensor) and src.is_floating_point() and self.shape == src.shape: + mask = ~torch.isnan(src) + if not bool(mask.all()): + self[mask] = src[mask].to(self.dtype) + return self + return orig_copy(self, src, *args, **kwargs) + + torch.Tensor.copy_ = masked_copy_ + try: + yield + finally: + torch.Tensor.copy_ = orig_copy diff --git a/verl/checkpoint_engine/delta_sync/sharded.py b/verl/checkpoint_engine/delta_sync/sharded.py new file mode 100644 index 00000000000..8f0894ef0d0 --- /dev/null +++ b/verl/checkpoint_engine/delta_sync/sharded.py @@ -0,0 +1,229 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Sharded delta: diff each rank's local FSDP shard, gather only the changes to rank 0. + +The default delta path all-gathers the full parameter (``DTensor.full_tensor()``) and +byte-diffs it against a full-model pinned-CPU snapshot on rank 0. This module instead lets +each rank keep a pinned snapshot of only *its* shard, byte-diff the shard locally, and +gather just the changed ``(within-parameter position, value)`` pairs to rank 0 -- so the +all-gather volume drops to the sparsity ratio (~1-3%) and rank 0 no longer needs a +full-model snapshot. The gathered result is bit-identical to the full-tensor diff, so the +downstream encode + broadcast and the receiver are unchanged. + +Scope: FSDP2 ``Shard(0)`` DTensors (the common case) + replicated / non-DTensor params. +Other shard dims are strided in the flattened layout and raise NotImplementedError. +""" + +from __future__ import annotations + +import torch +import torch.distributed as dist +from torch.distributed.tensor import DTensor +from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + +_DTYPE_INT = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64} + + +def _prod(xs) -> int: + n = 1 + for x in xs: + n *= int(x) + return n + + +def local_shard_view(param: torch.Tensor, mesh_rank0_only: bool = True): + """Return (local_flat_shard, within_param_flat_offset, contributes). + + * For a ``Shard(0)`` DTensor: the rank's local rows, flattened, and their flat offset + into the full (flattened) parameter, computed purely locally from the DTensor spec + (``compute_local_shape_and_global_offset`` does no collective). + * For a replicated / plain tensor: the whole tensor is present on every rank, so only + one rank should contribute it -- ``contributes`` is False on the others to avoid + double-counting (gated on the mesh coordinate, else global rank 0). + """ + if not isinstance(param, DTensor): + return param.reshape(-1), 0, (dist.get_rank() == 0 if dist.is_initialized() else True) + + placements = param.placements + for p in placements: + if p.is_shard() and p.dim != 0: + raise NotImplementedError( + f"sharded delta only supports Shard(0) (FSDP2 default); got placements={placements}" + ) + + # A parameter is replicated along any Replicate mesh dim (e.g. the ulysses/SP dim of a + # 2D FSDP mesh). Every rank on that dim holds the *same* shard, so only the coord-0 rank + # should contribute -- otherwise the gather double-counts it. + coord = param.device_mesh.get_coordinate() + contributes = True + if coord is not None: + for mesh_dim, p in enumerate(placements): + if p.is_replicate() and coord[mesh_dim] != 0: + contributes = False + break + + if all(p.is_replicate() for p in placements): + return param.to_local().reshape(-1), 0, contributes + + _, global_offset = compute_local_shape_and_global_offset( + param.shape, param.device_mesh, param.placements + ) + inner = _prod(param.shape[1:]) + offset = int(global_offset[0]) * inner + return param.to_local().reshape(-1), offset, contributes + + +def shard_delta_indices( + local_new: torch.Tensor, + local_snap: torch.Tensor, + offset: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Byte-diff a local shard against its snapshot; return (global_positions, values). + + Positions are int64 indices into the *full flattened parameter* (offset + local index). + Dtype-agnostic, bytewise (view-as-int), no arithmetic -- matches ``bytewise_diff_mask``. + """ + es = local_new.element_size() + int_dtype = _DTYPE_INT.get(es) + if int_dtype is None: + raise ValueError(f"unsupported element size {es}") + mask = local_new.view(int_dtype) != local_snap.view(int_dtype) + local_idx = mask.nonzero(as_tuple=False).view(-1) + values = local_new[local_idx] + global_idx = local_idx.to(torch.int64) + offset + return global_idx, values + + +def gather_v_to_rank0( + local_idx: torch.Tensor, + local_val: torch.Tensor, + group=None, +) -> tuple[torch.Tensor, torch.Tensor] | tuple[None, None]: + """Gather variable-length (idx, val) from every rank to rank 0 (pad-to-max ``dist.gather``). + + Returns (idx, val) on rank 0 (concatenated, padding stripped) and (None, None) elsewhere. + The counts are exchanged with a fixed-size all_gather; the payloads then ride a single + padded gather to rank 0 -- so nothing all-gathers the full tensor. + """ + rank = dist.get_rank(group) + world = dist.get_world_size(group) + dev = local_idx.device + + n = int(local_idx.numel()) + cnt = torch.tensor([n], dtype=torch.long, device=dev) + counts = [torch.zeros(1, dtype=torch.long, device=dev) for _ in range(world)] + dist.all_gather(counts, cnt, group=group) + counts = [int(c.item()) for c in counts] + max_n = max(counts) if counts else 0 + if max_n == 0: + return (torch.empty(0, dtype=torch.int64, device=dev), + torch.empty(0, dtype=local_val.dtype, device=dev)) if rank == 0 else (None, None) + + idx_pad = torch.zeros(max_n, dtype=torch.int64, device=dev) + val_pad = torch.zeros(max_n, dtype=local_val.dtype, device=dev) + idx_pad[:n] = local_idx + val_pad[:n] = local_val + + idx_list = [torch.zeros(max_n, dtype=torch.int64, device=dev) for _ in range(world)] if rank == 0 else None + val_list = [torch.zeros(max_n, dtype=local_val.dtype, device=dev) for _ in range(world)] if rank == 0 else None + dist.gather(idx_pad, idx_list, dst=0, group=group) + dist.gather(val_pad, val_list, dst=0, group=group) + + if rank != 0: + return None, None + idx = torch.cat([idx_list[r][: counts[r]] for r in range(world)]) + val = torch.cat([val_list[r][: counts[r]] for r in range(world)]) + return idx, val + + +def gather_dense_to_rank0( + local_val: torch.Tensor, + offset: int, + full_numel: int, + group=None, +) -> torch.Tensor | None: + """Assemble a full flat parameter on rank 0 from each rank's contiguous shard. + + Each rank contributes ``(offset, values)`` (empty on non-contributing ranks); + rank 0 places every shard at its flat offset. Only the values ride the wire -- + no per-element indices -- so the dense first sync carries none of the sparse + encoding overhead and rank 0 peaks at one full parameter. + """ + rank = dist.get_rank(group) + world = dist.get_world_size(group) + dev = local_val.device + + n = int(local_val.numel()) + meta = torch.tensor([n, offset], dtype=torch.long, device=dev) + metas = [torch.zeros(2, dtype=torch.long, device=dev) for _ in range(world)] + dist.all_gather(metas, meta, group=group) + counts = [int(m[0].item()) for m in metas] + offsets = [int(m[1].item()) for m in metas] + max_n = max(counts) if counts else 0 + if max_n == 0: + return torch.empty(0, dtype=local_val.dtype, device=dev) if rank == 0 else None + + val_pad = torch.zeros(max_n, dtype=local_val.dtype, device=dev) + val_pad[:n] = local_val + val_list = [torch.zeros(max_n, dtype=local_val.dtype, device=dev) for _ in range(world)] if rank == 0 else None + dist.gather(val_pad, val_list, dst=0, group=group) + if rank != 0: + return None + full = torch.empty(full_numel, dtype=local_val.dtype, device=dev) + for r in range(world): + if counts[r]: + full[offsets[r] : offsets[r] + counts[r]] = val_list[r][: counts[r]] + return full + + +def gather_v_grouped_to_rank0( + local_idx: torch.Tensor, + local_val: torch.Tensor, + group=None, +) -> list[tuple[torch.Tensor, torch.Tensor]] | None: + """Like :func:`gather_v_to_rank0`, but returns the payloads *per rank* instead of concatenated. + + Returns, on rank 0, a list ``[(idx_r, val_r) for r in range(world)]`` of each rank's sparse + ``(local-shard-position, value)`` pairs (padding stripped); ``None`` on the other ranks. Keeping + the per-rank boundary lets rank 0 rebuild each rank's shard buffer and feed the *native* + ``_weight_merge_across_tp`` -- so no per-parameter global-position math is needed on our side. + """ + rank = dist.get_rank(group) + world = dist.get_world_size(group) + dev = local_idx.device + + n = int(local_idx.numel()) + cnt = torch.tensor([n], dtype=torch.long, device=dev) + counts = [torch.zeros(1, dtype=torch.long, device=dev) for _ in range(world)] + dist.all_gather(counts, cnt, group=group) + counts = [int(c.item()) for c in counts] + max_n = max(counts) if counts else 0 + + if max_n == 0: + return [(torch.empty(0, dtype=torch.int64, device=dev), + torch.empty(0, dtype=local_val.dtype, device=dev)) for _ in range(world)] if rank == 0 else None + + idx_pad = torch.zeros(max_n, dtype=torch.int64, device=dev) + val_pad = torch.zeros(max_n, dtype=local_val.dtype, device=dev) + idx_pad[:n] = local_idx + val_pad[:n] = local_val + + idx_list = [torch.zeros(max_n, dtype=torch.int64, device=dev) for _ in range(world)] if rank == 0 else None + val_list = [torch.zeros(max_n, dtype=local_val.dtype, device=dev) for _ in range(world)] if rank == 0 else None + dist.gather(idx_pad, idx_list, dst=0, group=group) + dist.gather(val_pad, val_list, dst=0, group=group) + + if rank != 0: + return None + return [(idx_list[r][: counts[r]], val_list[r][: counts[r]]) for r in range(world)] diff --git a/verl/checkpoint_engine/delta_sync/wrapper.py b/verl/checkpoint_engine/delta_sync/wrapper.py new file mode 100644 index 00000000000..7ca97b8ba34 --- /dev/null +++ b/verl/checkpoint_engine/delta_sync/wrapper.py @@ -0,0 +1,167 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Adapter from verl's ``(name, tensor)`` generator API to delta flushes. + +The verl rollout API speaks "generator of (HF name, tensor)". This module +wraps that generator with the delta machinery and yields one +:class:`DeltaFlush` per bucket boundary -- the checkpoint engine broadcasts +those flushes over NCCL directly. + +The first call has to seed the snapshot (no flushes emitted, no engine RPCs); +subsequent calls produce one or more sparse flushes carrying only the +positions and values that actually changed. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Generator, Iterable + +import torch + +from .delta_state import DeltaState +from .encode import ( + DeltaBucket, + DeltaEncodingName, + DeltaParam, + checksum as _checksum, + encode_chunk, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class DeltaFlush: + """One ready-to-dispatch flush. + + * ``positions_cpu`` is a uint8 positions blob. The base encode path + produces it on the host; the sharded engines keep it on the GPU (the wire + broadcasts from the GPU, so a host round-trip would be pure overhead). + * ``values_gpu`` stays on the GPU until the checkpoint engine broadcasts it + over NCCL. + * ``params`` carries the per-parameter manifest the receiver needs to + decode the blob (sent alongside the data over the zmq side-channel). + """ + + encoding: DeltaEncodingName + params: list[DeltaParam] + positions_cpu: torch.Tensor + values_gpu: torch.Tensor + checksum: int + + @property + def nnz(self) -> int: + return self.values_gpu.numel() + + @property + def wire_bytes(self) -> int: + return ( + self.positions_cpu.numel() + + self.values_gpu.numel() * self.values_gpu.element_size() + ) + + +def _chunked( + iterable: Iterable[tuple[str, torch.Tensor]], chunk_params: int +) -> Iterable[list[tuple[str, torch.Tensor]]]: + chunk: list[tuple[str, torch.Tensor]] = [] + for item in iterable: + chunk.append(item) + if len(chunk) >= chunk_params: + yield chunk + chunk = [] + if chunk: + yield chunk + + +def _materialize_flush( + bucket: DeltaBucket, encoding: DeltaEncodingName +) -> DeltaFlush: + positions_u8 = bucket.merged_positions() + values_gpu = bucket.merged_values() + params = list(bucket.params) + bucket.clear() + # Positions already live on the values' device (encode keeps them there); + # checksum and the NCCL broadcast both consume them in place. + cks = _checksum(positions_u8, values_gpu) + return DeltaFlush( + encoding=encoding, + params=params, + positions_cpu=positions_u8, + values_gpu=values_gpu, + checksum=cks, + ) + + +def iter_delta_flushes( + weights: Generator[tuple[str, torch.Tensor], None, None], + state: DeltaState, + *, + encoding: DeltaEncodingName, + bucket_bytes: int, + chunk_params: int = 16, +) -> Iterable[DeltaFlush]: + """Wrap ``weights`` into a stream of delta flushes. + + On the very first invocation the snapshot is seeded from ``weights`` and + no flush is emitted -- callers must catch the empty iterator and skip + engine dispatch (the receiver is assumed to share the init checkpoint). + + Args: + weights: ``(name, tensor)`` generator from + ``actor.engine.get_per_tensor_param()``. + state: persistent snapshot/streams owned by the worker. + encoding: positions encoding. + bucket_bytes: flush threshold; matches + ``checkpoint_engine.update_weights_bucket_megabytes`` semantics. + chunk_params: how many parameters per encode chunk. Tuning knob for + the 1-step H2D prefetch lookahead. + """ + if not state.seeded: + seed = list(weights) + state.seed(seed) + logger.info("DeltaState seeded with %d HF tensors", len(seed)) + return + + bucket = DeltaBucket() + pending_chunk: list[tuple[str, torch.Tensor]] | None = None + pending_prefetch: tuple[list[torch.Tensor], torch.cuda.Event] | None = None + + for hf_chunk in _chunked(weights, chunk_params): + next_prefetch = state.prefetch_snapshot(hf_chunk) + if pending_chunk is not None and pending_prefetch is not None: + diffs = state.compute_diffs(pending_chunk, prefetched=pending_prefetch) + state.update_snapshot_async(pending_chunk) + chunk = encode_chunk(diffs, encoding) + if chunk.params: + if bucket.should_flush_before_add(chunk, bucket_bytes): + yield _materialize_flush(bucket, encoding) + bucket.add(chunk) + pending_chunk, pending_prefetch = hf_chunk, next_prefetch + + if pending_chunk is not None and pending_prefetch is not None: + diffs = state.compute_diffs(pending_chunk, prefetched=pending_prefetch) + state.update_snapshot_async(pending_chunk) + chunk = encode_chunk(diffs, encoding) + if chunk.params: + if bucket.should_flush_before_add(chunk, bucket_bytes): + yield _materialize_flush(bucket, encoding) + bucket.add(chunk) + + if bucket.has_updates: + yield _materialize_flush(bucket, encoding) + + state.flush_snapshot() diff --git a/verl/experimental/one_step_off_policy/shell/grpo_0.6b_gsm8k_fsdp2_sglang_delta_2_6.sh b/verl/experimental/one_step_off_policy/shell/grpo_0.6b_gsm8k_fsdp2_sglang_delta_2_6.sh new file mode 100644 index 00000000000..6a0ca11c893 --- /dev/null +++ b/verl/experimental/one_step_off_policy/shell/grpo_0.6b_gsm8k_fsdp2_sglang_delta_2_6.sh @@ -0,0 +1,82 @@ +set -x + +# One-step-off (disaggregated) GRPO with DELTA weight sync. +# +# Same as grpo_0.6b_gsm8k_fsdp2_sglang_2_6.sh, but the trainer -> rollout weight +# sync only broadcasts the parameters that changed since the previous sync, +# instead of the full model. In RL post-training >99% of BF16 weight bytes are +# unchanged step-over-step, so this cuts the disaggregated weight-sync traffic to +# the sparsity ratio while staying bit-exact (per-flush checksum verified). +# +# Enable it with two flags (see the last two lines of the python command): +# actor_rollout_ref.rollout.checkpoint_engine.backend=delta +# +actor_rollout_ref.rollout.checkpoint_engine.engine_kwargs.delta.encoding=indices +# +# Requirements / scope: disaggregated (hybrid_engine=False) + SGLang rollout in +# BF16. Encoding is "indices" (int32 absolute positions) or "deltas" (uint16 gap). + +project_name='GRPO' +exp_name='GRPO-Qwen3-0.6b-gsm8k-fsdp2-sglang-one-step-off-delta-2-6' + +# Paths +RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"} +MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen3-0.6B"} +CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"} +TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/gsm8k/train.parquet"} +TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/gsm8k/test.parquet"} + +NNODES=${NNODES:-1} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +n_gpus_rollout=2 +n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout)) + + +python3 -m verl.experimental.one_step_off_policy.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files="${TRAIN_FILE}" \ + data.val_files="${TEST_FILE}" \ + data.train_batch_size=1152 \ + data.max_prompt_length=512 \ + data.max_response_length=1024 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \ + critic.strategy=fsdp2 \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.hybrid_engine=False \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=192 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.name=sglang \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \ + actor_rollout_ref.rollout.n=5 \ + actor_rollout_ref.rollout.load_format=safetensors \ + actor_rollout_ref.rollout.layered_summon=True \ + actor_rollout_ref.rollout.checkpoint_engine.backend=delta \ + +actor_rollout_ref.rollout.checkpoint_engine.engine_kwargs.delta.encoding=indices \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + algorithm.use_kl_in_reward=False \ + trainer.critic_warmup=0 \ + trainer.val_before_train=True \ + trainer.logger=['console','tensorboard'] \ + trainer.project_name="${project_name}" \ + trainer.experiment_name="${exp_name}" \ + trainer.save_freq=-1 \ + trainer.test_freq=5 \ + trainer.total_epochs=2 \ + trainer.nnodes="${NNODES}" \ + trainer.n_gpus_per_node="${n_gpus_training}" \ + rollout.nnodes="${NNODES}" \ + rollout.n_gpus_per_node="${n_gpus_rollout}" $@ diff --git a/verl/experimental/separation/ray_trainer.py b/verl/experimental/separation/ray_trainer.py index 88e5e887606..de9f14e14ec 100644 --- a/verl/experimental/separation/ray_trainer.py +++ b/verl/experimental/separation/ray_trainer.py @@ -642,7 +642,11 @@ def _fit_update_weights(self): if self.config.trainer.critic_warmup <= self.global_steps: # update weights from trainer to rollout with marked_timer("update_weights", timing_raw, color="red"): - self.checkpoint_manager.update_weights(self.global_steps) + sync_metrics = self.checkpoint_manager.update_weights(self.global_steps) + # Engines may report per-sync stats (e.g. the delta backends' changed + # ratio / wire payload); surface them in this step's metrics. + if sync_metrics: + self.metrics.update(sync_metrics) def _fit_dump_data(self, batch: DataProto): timing_raw = self.timing_raw diff --git a/verl/workers/engine/base.py b/verl/workers/engine/base.py index 3003eab7cf5..6d89e3bdeea 100644 --- a/verl/workers/engine/base.py +++ b/verl/workers/engine/base.py @@ -158,6 +158,25 @@ def get_per_tensor_param(self) -> tuple[Generator[tuple[str, torch.Tensor], None """ raise NotImplementedError + def get_per_tensor_param_shard(self, **kwargs) -> tuple[Generator, Optional[dict]]: + """ + Like :meth:`get_per_tensor_param`, but yields each rank's *local* parameter shard + instead of all-gathering full tensors. + + Used by checkpoint engines that operate on shards (e.g. the ``delta_sharded`` + backends, which byte-diff each rank's shard locally and only gather the changed + elements), so the export never materializes full tensors on any rank. + + The per-item schema is engine-specific: implementations yield the local shard + plus enough placement metadata for the consumer to map shard-local positions + back to the full parameter (see the FSDP and Megatron engine implementations). + + Returns: + Generator: A generator that yields per-parameter local shards with placement metadata. + Optional[dict]: Optional peft config. + """ + raise NotImplementedError + def get_data_parallel_size(self): raise NotImplementedError diff --git a/verl/workers/engine/fsdp/transformer_impl.py b/verl/workers/engine/fsdp/transformer_impl.py index b68944a66fa..1fe197e8af2 100644 --- a/verl/workers/engine/fsdp/transformer_impl.py +++ b/verl/workers/engine/fsdp/transformer_impl.py @@ -814,6 +814,86 @@ def load_checkpoint( if self._is_offload_optimizer: offload_fsdp_optimizer(self.optimizer) + def get_per_tensor_param_shard(self, **kwargs): + """Like :meth:`get_per_tensor_param`, but yields each rank's *local* FSDP shard + instead of all-gathering the full tensor -- used by the ``delta_sharded`` + checkpoint engine, which diffs on shards and gathers only the changes. + + Yields ``(name, local_flat_shard_bf16, within_param_flat_offset, full_numel, + full_shape, contributes)``. Non-LoRA base path only. + """ + import os as _os + + from verl.checkpoint_engine.delta_sync.sharded import local_shard_view + + _debug = bool(_os.environ.get("VERL_DELTA_DEBUG")) + # FSDP1's (SHARDED_)STATE_DICT export runs through the unshard machinery and + # asserts flat params are GPU-resident; FSDP2 state_dict() only collects + # DTensor refs and the generator below stages each shard lazily. + _needs_staging = fsdp_version(self.module) == 1 + _rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + if _debug: + from torch.distributed.fsdp import FSDPModule as _FSDP2Module + from torch.distributed.fsdp import FullyShardedDataParallel as _FSDP1 + + n1 = sum(isinstance(m, _FSDP1) for m in self.module.modules()) + n2 = sum(isinstance(m, _FSDP2Module) for m in self.module.modules()) + print( + f"[delta-debug r{_rank}] MODULE TREE: fsdp1_wrapped={n1} fsdp2_wrapped={n2} " + f"root={type(self.module).__name__} staging={_needs_staging}", + flush=True, + ) + + if _needs_staging and not self._uses_fsdp2_cpu_offload_policy: + load_fsdp_model_to_gpu(self.module) + params = self.module.state_dict() + params = convert_weight_keys(params, getattr(self.module, "_fsdp_wrapped_module", self.module)) + if _needs_staging and self._is_offload_param: + offload_fsdp_model_to_cpu(self.module) + + device = get_device_id() + + def _gen(): + from torch.distributed.tensor import DTensor as _DT + + n_dt = n_plain = n_contrib = 0 + local_elems = full_elems = 0 + logged = 0 + for name, param in params.items(): + full_shape = tuple(param.shape) + full_numel = int(param.numel()) + p = param.to(device, non_blocking=True) + if p.is_floating_point(): + p = p.to(torch.bfloat16, non_blocking=True) + local, offset, contributes = local_shard_view(p) + if _debug: + is_dt = isinstance(param, _DT) + n_dt += is_dt + n_plain += not is_dt + n_contrib += bool(contributes) + local_elems += int(local.numel()) + full_elems += full_numel + if logged < 5: + logged += 1 + placements = getattr(param, "placements", None) + print( + f"[delta-debug r{_rank}] param={name} dtensor={is_dt} placements={placements} " + f"full={list(full_shape)}({full_numel}) local_shard={list(local.shape)}({local.numel()}) " + f"offset={offset} contributes={contributes} src_dev={param.device}", + flush=True, + ) + yield name, local, offset, full_numel, full_shape, contributes + if _debug: + ratio = local_elems / max(full_elems, 1) + print( + f"[delta-debug r{_rank}] EXPORT SUMMARY: params={n_dt + n_plain} dtensor={n_dt} plain={n_plain} " + f"contributing={n_contrib} local/full elems={local_elems}/{full_elems} (ratio={ratio:.4f} " + f"-- ~1/fsdp_size means truly sharded, ~1.0 means degenerate full tensors)", + flush=True, + ) + + return _gen(), None + def get_per_tensor_param(self, layered_summon=False, base_sync_done=False, **kwargs): log_gpu_memory_usage("Before load_fsdp_model_to_gpu", logger=logger) diff --git a/verl/workers/engine_workers.py b/verl/workers/engine_workers.py index c65b99e44d6..82b75bb8650 100644 --- a/verl/workers/engine_workers.py +++ b/verl/workers/engine_workers.py @@ -698,9 +698,15 @@ async def update_weights(self, global_steps: int = None, mode: str = "auto"): # 0. send_weights only for async training with disaggregated trainer and rollout if effective_mode != "naive": - per_tensor_param, _ = self.actor.engine.get_per_tensor_param() + # The sharded delta engine diffs each rank's local FSDP shard (no all-gather), + # so it consumes the sharded param generator instead of the full-tensor one. + if effective_mode == "delta_sharded": + per_tensor_param, _ = self.actor.engine.get_per_tensor_param_shard() + else: + per_tensor_param, _ = self.actor.engine.get_per_tensor_param() await self.checkpoint_engine.send_weights(per_tensor_param, global_steps=global_steps) - return + # Engines are registry-pluggable and not required to subclass CheckpointEngine. + return getattr(self.checkpoint_engine, "pop_sync_metrics", dict)() set_expandable_segments(False) log_gpu_memory_usage("Before resume weights", logger=logger) diff --git a/verl/workers/rollout/sglang_rollout/async_sglang_server.py b/verl/workers/rollout/sglang_rollout/async_sglang_server.py index 171a9b3f2e7..8e0b148357c 100644 --- a/verl/workers/rollout/sglang_rollout/async_sglang_server.py +++ b/verl/workers/rollout/sglang_rollout/async_sglang_server.py @@ -256,6 +256,16 @@ async def launch_server(self, master_address: str = None, master_port: int = Non engine_kwargs = self.config.get("engine_kwargs", {}).get("sglang", {}) or {} attention_backend = engine_kwargs.pop("attention_backend", None) mm_attention_backend = engine_kwargs.pop("mm_attention_backend", None) + # Delta checkpoint engines apply sparse weight updates in place through SGLang's + # custom-weight-loader hook; register the verl loader so the update requests' + # load_format resolves inside the TP workers. + custom_weight_loader = list(engine_kwargs.pop("custom_weight_loader", None) or []) + ce_backend = str((self.config.get("checkpoint_engine", None) or {}).get("backend", "")) + if ce_backend.startswith("delta"): + from verl.checkpoint_engine.delta_sync.sglang_loader import LOADER_FQN + + if LOADER_FQN not in custom_weight_loader: + custom_weight_loader.append(LOADER_FQN) if attention_backend is None: if torch.version.hip is not None: attention_backend = "aiter" @@ -310,6 +320,7 @@ async def launch_server(self, master_address: str = None, master_port: int = Non "json_model_override_args": json.dumps({"quantization_config": fp8_block_quant_kwargs}) if quantization == "fp8" else json.dumps({}), + "custom_weight_loader": custom_weight_loader or None, **engine_kwargs, }