[Sync][E] delta weight sync (slime #1806/#1946/#1991)

.github/workflows/pr-test.yml

@@ -60,7 +60,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        info: [{"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_ppo_critic_only_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}]
+        info: [{"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_ppo_critic_only_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}, {"num_gpus": 4, "test_file": "test_delta_weight_update.py"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}
@@ -482,7 +482,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        info: [{"num_gpus": 0, "test_file": "test_megatron_argument_validation.py"}, {"num_gpus": 0, "test_file": "test_rollout_validation.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_rollout_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_runtime_hook_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_path_loading_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_generate_contracts.py"}, {"num_gpus": 0, "test_file": "test_rm_deepscaler.py"}, {"num_gpus": 0, "test_file": "test_rm_f1.py"}, {"num_gpus": 0, "test_file": "test_rm_gpqa.py"}, {"num_gpus": 0, "test_file": "test_rm_math.py"}, {"num_gpus": 0, "test_file": "test_rm_math_dapo.py"}]
+        info: [{"num_gpus": 0, "test_file": "test_megatron_argument_validation.py"}, {"num_gpus": 0, "test_file": "test_rollout_validation.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_rollout_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_runtime_hook_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_path_loading_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_generate_contracts.py"}, {"num_gpus": 0, "test_file": "test_rm_deepscaler.py"}, {"num_gpus": 0, "test_file": "test_rm_f1.py"}, {"num_gpus": 0, "test_file": "test_rm_gpqa.py"}, {"num_gpus": 0, "test_file": "test_rm_math.py"}, {"num_gpus": 0, "test_file": "test_rm_math_dapo.py"}, {"num_gpus": 0, "test_file": "test_dp_schedule.py"}, {"num_gpus": 0, "test_file": "test_cp_utils.py"}, {"num_gpus": 0, "test_file": "test_metric_report.py"}, {"num_gpus": 0, "test_file": "test_metric_report_dist.py"}, {"num_gpus": 0, "test_file": "test_loss_cp_invariance.py"}, {"num_gpus": 0, "test_file": "test_sample.py"}, {"num_gpus": 0, "test_file": "utils/test_hf_checkpoint_saver.py"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}

.github/workflows/pr-test.yml.j2

@@ -6,6 +6,7 @@
         {'test_file': 'test_qwen3.5_0.8B_gsm8k_short.py', 'num_gpus': 4},
         {'test_file': 'test_qwen2.5_0.5B_ppo_critic_only_short.py', 'num_gpus': 4},
         {'test_file': 'test_qwen2.5_0.5B_fully_async_short.py', 'num_gpus': 4},
+        {'test_file': 'test_delta_weight_update.py', 'num_gpus': 4},
       ],
     },
     'e2e-test-vllm-config': {
@@ -65,6 +66,13 @@
         {'test_file': 'test_rm_gpqa.py', 'num_gpus': 0},
         {'test_file': 'test_rm_math.py', 'num_gpus': 0},
         {'test_file': 'test_rm_math_dapo.py', 'num_gpus': 0},
+        {'test_file': 'test_dp_schedule.py', 'num_gpus': 0},
+        {'test_file': 'test_cp_utils.py', 'num_gpus': 0},
+        {'test_file': 'test_metric_report.py', 'num_gpus': 0},
+        {'test_file': 'test_metric_report_dist.py', 'num_gpus': 0},
+        {'test_file': 'test_loss_cp_invariance.py', 'num_gpus': 0},
+        {'test_file': 'test_sample.py', 'num_gpus': 0},
+        {'test_file': 'utils/test_hf_checkpoint_saver.py', 'num_gpus': 0},
       ],
     },
 

docs/zh/advanced/delta-weight-sync.md

@@ -0,0 +1,51 @@
+# Delta 权重同步(Delta Weight Sync)
+
+> 移植自 THUDM/slime #1806/#1946/#1991,翻译到 vime 的 vLLM rollout。
+
+## 它解决什么
+
+RL 训练每步更新 policy 权重后,需要把新权重同步到 rollout 推理引擎(vLLM)。默认的**全量同步**每次把整套参数广播过去 —— 大模型(尤其大 MoE)或跨数据中心场景下,每步全量传输的带宽开销很大。
+
+**Delta 权重同步只传"自上次同步以来变化的字节"**:训练侧对当前权重和上次广播的 pinned-CPU 快照做**逐字节 diff**,只打包变化的位置 + 新值发出去;接收侧把这些字节**原样覆盖**进 live 模型(无损、无算术、不累积漂移)。
+
+> 注意:这是**互补的带宽优化,不是必需**。**colocate(训练/推理同机)用 CUDA-IPC 直传更快,无需 delta**。delta 面向 **non-colocate**(尤其大模型 / 跨 DC / 带宽受限)。
+
+## 开启
+
+```bash
+--update-weight-mode delta            # 默认 full
+--update-weight-transport nccl        # 或 disk
+--update-weight-encoding deltas_zstd  # indices | deltas | deltas_zstd
+```
+
+### 参数
+
+| 参数 | 说明 |
+|---|---|
+| `--update-weight-mode {full,delta}` | `full`=每次广播整参数(默认);`delta`=只发逐字节变化 |
+| `--update-weight-transport {nccl,disk}` | delta 的每桶载体。`nccl`=NCCL 广播(低延迟、同 DC);`disk`=写 safetensors 落共享盘 + 每轮一次 HTTP 唤醒引擎读取(跨 DC、带宽受限) |
+| `--update-weight-encoding {indices,deltas,deltas_zstd}` | 位置编码。`indices`=int32 绝对下标(最大、计算最省);`deltas`=uint16 间隔差(更小);`deltas_zstd`=`deltas` + zstd 压缩(最小) |
+| `--update-weight-delta-dir <path>` | `disk` 传输时,每轮 delta safetensors 的目录(训练侧与引擎侧共享文件系统) |
+| `--update-weight-delta-keep-files` | 保留各轮 delta 文件(默认清理) |
+| `--update-weight-delta-chunk-bytes <N>` | 接收侧每次 `load_weights` 的字节预算(解码后分块 apply) |
+
+## 工作原理(vime 实现)
+
+1. **训练侧**(`UpdateWeightFromDistributedDelta`,继承全量的 `UpdateWeightFromDistributed`):留一份上次广播权重的 pinned-CPU 快照;每次同步做逐字节 diff → `(positions, values)` + 每参数解码清单 `DeltaSpec`。第一次同步只 seed 快照、不联系引擎(引擎 init 时已加载同一 HF checkpoint)。
+2. **wire 格式**:`__positions__`(uint8 字节偏移 blob)+ `__values__`(参数 dtype 的值)+ JSON `DeltaSpec`(encoding + 每参数切片 + checksum)。`nccl`/`disk` 两种传输共享。
+3. **接收侧**(vime 的 `vLLMColocateWorkerExtension` 上经 `/collective_rpc` 调用的方法 + `delta_receiver.py`):
+   - `nccl`:按训练侧广播顺序在 `model_update_group` 上收 `(positions, values)`;
+   - `disk`:读 + 解压 safetensors;
+   - 解码:positions 解包成下标(`indices` 绝对 / `deltas` 间隔差 cumsum 还原),`index_copy_` 到一个填 NaN 的全形 tensor(NaN=未变);
+   - apply:`with delta_apply_context(model): model.load_weights(chunk)` —— vLLM 的 `load_weights` 照常分片,期间 `torch.Tensor.copy_/fill_` 被临时拦截,**只把非-NaN 位写进模型参数存储**;`post_load_weights`(fp8 scale 等派生量)用原始 copy_ 照常重算。
+4. **无需重建镜像**:接收侧是 vime 运行时注入的 worker-extension 方法 / monkey-patch(经容器内 `pip install -e .` 加载),不改 vLLM 引擎源、不改 Dockerfile。
+
+## 已知限制(重要)
+
+- **投机解码(EAGLE/MTP)的 draft 模型不会被同步刷新**。vLLM 在任何权重同步(全量 / delta)下都**只更新主模型**(`gpu_model_runner.reload_weights` 只 load 主模型的 `named_parameters`;drafter 仅在 init 时加载一次)。因此 RL + spec-decode 下,draft 专属层(EAGLE transformer / MTP head)在同步后会 stale —— 这是 **vLLM 层面的全局限制,与 delta 无关**,全量同步也一样。slime 曾用 #1993 在 sglang 侧给 EAGLE draft-worker 转发 delta 文件来修,但 vLLM 没有等价机制;**此 PR 不实现 draft 覆盖**,待上游(vLLM)补 drafter-on-update 后再跟进。
+- `disk` 传输要求训练侧与引擎侧**共享文件系统**;`nccl` 传输用于同机/同 DC。
+- delta diff/编码有 CPU 开销;收益在"变化字节量 ≪ 整模型"时(大模型 / 跨 DC)才显著。
+
+## 验证
+
+`tests/test_delta_weight_update.py`(e2e):delta 同步后,rollout/train logprob 与全量同步一致;zero-delta(无变化)轮次的 weight_version 仍正确递增(不触发 actor 的 version 一致性断言)。

docs/zh/index.rst

@@ -44,6 +44,7 @@ vime 是 GLM-4.7、GLM-4.6、GLM-4.5 背后的 RL 训练框架。除此之外，
    advanced/vllm-config.md
    advanced/megatron-config.md
    advanced/arch-support-beyond-megatron.md
+   advanced/delta-weight-sync.md
 
 .. toctree::
    :maxdepth: 1

requirements.txt

@@ -13,6 +13,7 @@ pyyaml
 qwen_vl_utils # for VLM
 ray[default]
 ring_flash_attn
+safetensors
 tensorboard
 transformers
 vllm-router>=0.1.14

tests/_cp_dist_helpers.py

@@ -0,0 +1,167 @@
+"""Shared infrastructure for the CP-related multi-process CPU tests.
+
+Why this module exists
+----------------------
+The CP / metric-report / backward-grad-norm tests all want to:
+
+1. Stub ``megatron.core.mpu`` *before* importing
+   ``vime.backends.megatron_utils.cp_utils`` (the CPU CI image has no real
+   megatron).
+2. Spawn ``dp_size * cp_size`` workers with real ``torch.distributed`` and
+   exercise the actual production helpers (``get_sum_of_sample_mean``,
+   ``reduce_train_step_metrics``, ``gather_and_reduce_log_dict``,
+   ``rollout_log_metric_contribution``).
+3. Chunk each sample's response tensor across CP ranks the same way the
+   real forward pass does — using
+   ``get_logits_and_tokens_offset_with_cp`` so the slicing stays in lock-
+   step with the production reducer.
+
+Putting that here keeps the per-feature test files focused on the
+behaviour they check (numerics / report formulas / backward) rather than
+on plumbing.
+
+Mapping to Megatron
+-------------------
+- ``mp.spawn(...)`` + gloo backend mirrors the per-rank entry-point that
+  ``torch.distributed.run`` would create for a real launch.
+- ``dp_cp_group = new_group(range(world_size))`` matches
+  ``parallel_state.get_data_parallel_group(with_context_parallel=True)``
+  (Megatron-LM ``finalize_model_grads.py:437``).  In the no-TP / no-PP
+  CPU test setup the whole world *is* that group.
+- The per-rank CP chunking mirrors what the attention layer feeds into
+  the loss in Megatron: each CP rank only sees its 2-chunk slice of the
+  response tokens (cf. ``cp_utils.get_logits_and_tokens_offset_with_cp``,
+  the same helper used by the real forward pass).
+"""
+
+from __future__ import annotations
+
+import os
+import socket
+import sys
+import types
+
+
+# --- Stub ``megatron.core.mpu`` (must run before cp_utils is imported) ---
+#
+# Both this module and any test file that imports it should *import this
+# helper first*. Doing so installs the stub at import time so that the
+# subsequent ``from vime.backends.megatron_utils.cp_utils import ...`` in
+# the test file binds ``cp_utils.mpu`` to this stub.
+#
+# In spawned workers, ``mp.spawn`` re-imports the test module fresh, which
+# re-runs this stub installation; then the worker mutates the stub's
+# ``get_context_parallel_*`` attributes via ``_stub_megatron_in_worker``
+# below to pin (cp_size, cp_rank) for that worker.
+_fake_mpu = types.ModuleType("megatron.core.mpu")
+_fake_mpu.get_context_parallel_world_size = lambda: 1
+_fake_mpu.get_context_parallel_rank = lambda: 0
+_fake_core = types.ModuleType("megatron.core")
+_fake_core.mpu = _fake_mpu
+_fake_megatron = types.ModuleType("megatron")
+_fake_megatron.core = _fake_core
+sys.modules.setdefault("megatron", _fake_megatron)
+sys.modules.setdefault("megatron.core", _fake_core)
+sys.modules.setdefault("megatron.core.mpu", _fake_mpu)
+
+
+def stub_megatron_in_worker(cp_size: int, cp_rank: int) -> None:
+    """Override ``mpu.get_context_parallel_*`` inside an ``mp.spawn`` worker.
+
+    ``mp.spawn`` pickles the worker function by name and re-imports the
+    test module in the child — that re-runs the top-of-file stub install
+    with ``cp_size=1``. By the time the worker runs, ``cp_utils`` has
+    already bound its module-level ``mpu`` reference to the stub.
+
+    So we must MUTATE the stub module's attributes in place rather than
+    replace ``sys.modules['megatron.core.mpu']`` — replacing the module
+    would leave ``cp_utils.mpu`` pointing at the now-shadowed stub.
+    """
+    from megatron.core import mpu  # the stub installed at import time
+
+    mpu.get_context_parallel_world_size = lambda: cp_size
+    mpu.get_context_parallel_rank = lambda: cp_rank
+
+
+def free_port() -> int:
+    """Pick an unused TCP port for ``init_process_group``'s rendezvous.
+
+    Equivalent to what ``torchrun`` does when ``--master-port`` is not
+    set; we just need a port nothing else is bound to so multiple
+    parametrized test cases can spawn without colliding.
+    """
+    s = socket.socket()
+    s.bind(("", 0))
+    port = s.getsockname()[1]
+    s.close()
+    return port
+
+
+def init_worker_process_group(rank: int, world_size: int, master_port: int):
+    """Stand up gloo ``torch.distributed`` and return the DP*CP group.
+
+    The CPU CI image ships gloo but not NCCL; in the no-TP / no-PP setup
+    the DP-with-CP group is the whole world, mirroring
+    ``parallel_state.get_data_parallel_group(with_context_parallel=True)``
+    in Megatron-LM ``finalize_model_grads.py:437``.
+    """
+    import torch.distributed as _dist
+
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(master_port)
+    _dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)
+    return _dist.new_group(ranks=list(range(world_size)))
+
+
+def cp_chunk_response_tensor(x, total_length: int, response_length: int):
+    """Slice a sample's response tensor to what the current CP rank sees.
+
+    Mirrors the real forward pass: at CP > 1 each rank's attention only
+    consumes the two response-token chunks selected by
+    ``get_logits_and_tokens_offset_with_cp`` (the same helper used by the
+    production reducer in ``cp_utils.get_sum_of_sample_mean``). So the
+    "x" we feed into the reducer on a CP rank must be sliced the same
+    way to keep the numbers honest.
+
+    Importing locally so callers don't pay the import cost before
+    ``stub_megatron_in_worker`` has had a chance to pin (cp_size, cp_rank).
+    """
+    import torch
+
+    from vime.backends.megatron_utils.cp_utils import get_logits_and_tokens_offset_with_cp
+
+    prompt_length = total_length - response_length
+    _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(total_length, response_length)
+    c0 = x[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length]
+    c1 = x[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length]
+    return torch.cat([c0, c1])
+
+
+# ---------------------------------------------------------------------------
+# Shared four-rollout fixture, used by both the metric-report distributed
+# tests and the backward-grad-norm test. Keeping the data in one place so
+# the "train report matches rollout report matches grad-norm baseline"
+# contract is anchored on the same numbers everywhere.
+#
+# Four samples (1 rollout each), total_length=12 (4 prompt + 8 response),
+# loss_mask=all-ones. x values differ by orders of magnitude so any cross-
+# rank summation bug shows up as a visibly wrong number.
+#
+# Per-sample token-mean: 4.5 / 45 / 450 / 4500.
+# Per-rollout-mean report (sum / num_rollouts):
+#       (4.5 + 45 + 450 + 4500) / 4 = 1249.875
+# Per-token-loss report (sum_x / total_tokens):
+#       (36 + 360 + 3600 + 36000) / 32 = 1249.875
+# (the two paths agree by construction so the test expectations stay
+# simple — the *report formulas* are still distinct as exercised inside
+# ``reduce_train_step_metrics``.)
+# ---------------------------------------------------------------------------
+FOUR_ROLLOUT_TOTAL_LENGTHS = [12, 12, 12, 12]
+FOUR_ROLLOUT_RESPONSE_LENGTHS = [8, 8, 8, 8]
+FOUR_ROLLOUT_X_VALUES = [
+    [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
+    [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0],
+    [100.0, 200.0, 300.0, 400.0, 500.0, 600.0, 700.0, 800.0],
+    [1000.0, 2000.0, 3000.0, 4000.0, 5000.0, 6000.0, 7000.0, 8000.0],
+]
+FOUR_ROLLOUT_EXPECTED_REPORT = 1249.875