feat(weight-sync): add engine-group wave scheduling

.github/workflows/pr-test.yml

@@ -617,6 +617,18 @@ jobs:
               "straw": true,
               "test_file": "test_megatron_argument_validation.py"
             },
+            {
+              "num_gpus": 0,
+              "test_file": "test_engine_group_wave.py"
+            },
+            {
+              "num_gpus": 0,
+              "test_file": "test_engine_group_wave_distributed.py"
+            },
+            {
+              "num_gpus": 0,
+              "test_file": "test_weight_sync_callsite_benchmark.py"
+            },
             {
               "num_gpus": 0,
               "test_file": "utils/test_megatron_role_config.py"

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

@@ -73,6 +73,9 @@
       'extra_pip_deps': 'transformers wandb',
       'tests': [
         {'test_file': 'test_megatron_argument_validation.py', 'num_gpus': 0, 'straw': True},
+        {'test_file': 'test_engine_group_wave.py', 'num_gpus': 0},
+        {'test_file': 'test_engine_group_wave_distributed.py', 'num_gpus': 0},
+        {'test_file': 'test_weight_sync_callsite_benchmark.py', 'num_gpus': 0},
         {'test_file': 'utils/test_megatron_role_config.py', 'num_gpus': 0},
         {'test_file': 'test_deep_ep_tms_patch.py', 'num_gpus': 0},
         {'test_file': 'test_stateless_adam.py', 'num_gpus': 0},

docs/en/advanced/weight-sync-wave-scheduling.md

@@ -0,0 +1,104 @@
+# Weight-Sync Engine Waves
+
+Large rollout deployments can make every SGLang engine receive or load the same
+weight bucket at once. That maximizes fan-out, but can also create a short burst
+of GPU, PCIe, NVLink, network, or storage traffic. slime can bound this fan-out
+without adding a hardware-specific time delay:
+
+```bash
+--update-weight-max-inflight-engine-groups 2
+```
+
+An engine group here means one logical SGLang engine. Its tensor-, pipeline-, or
+multi-node workers remain one indivisible group. The scheduler walks the resolved
+engine list in stable order and admits at most the configured number at once.
+For four engines and a limit of three, the waves are therefore `(0, 1, 2)` and
+`(3)`; no particular world size or number of engines is assumed.
+
+The default is `0`, which keeps the existing all-at-once behavior. A value at
+least as large as the engine count is equivalent to the default. Negative values
+are rejected.
+
+## Supported transports
+
+- Non-colocated NCCL weight sync creates one process group per logical engine
+  only when a real bound is requested. Broadcasts for the engines in one wave
+  are launched asynchronously, and the next wave is not admitted until both the
+  NCCL work and engine-side load requests complete. The default continues to use
+  the original aggregate process group.
+- Colocated tensor/IPC sync coordinates the same deterministic waves across all
+  trainer ranks. A trainer control-group barrier closes each wave before source
+  buffers can be reused.
+- Full and delta disk sync apply the bound to checkpoint pulls and engine reload
+  requests.
+- Quantized post-load processing uses the same engine-group bound.
+
+The existing pause, flush, and weight-version boundaries are unchanged: an engine
+is not resumed early merely because its own wave completed. Disk-backed checkpoint
+pulls may run before the pause, but the serving weights are not changed by that
+prefetch. Consequently, wave admission does not introduce a new mixed-version
+serving interval.
+
+## Choosing a value
+
+Start with `0` as the throughput baseline, then compare small limits such as `1`,
+`2`, and `4` on the deployment topology that will run the job. Select the largest
+value that avoids the observed traffic or memory burst without increasing total
+weight-sync time or tail latency. This option limits concurrent engine groups;
+`--update-weight-buffer-size` independently controls the size of each weight
+bucket.
+
+## Production-callsite confirmation
+
+`tools/benchmark_weight_sync_callsite.py` invokes the production
+`update_weights_in_engine_group_waves` function with real process groups and
+synthetic payloads. It is intentionally a callsite/module probe: it does not
+replace the scheduler, change defaults, or claim Ray/SGLang load performance.
+
+Use a four-process launch to retain three real two-rank
+`[trainer, engine]` groups. A/B are engine groups 0 and 1; the third group is
+the competing operation that distinguishes the all-at-once and window-2
+policies. Only use device counts permitted by the test environment.
+
+Each command below creates exactly one independent process-run artifact:
+
+```bash
+torchrun --standalone --nproc-per-node=4 \
+  tools/benchmark_weight_sync_callsite.py \
+  --backend gloo --policy candidate_windowed \
+  --evidence-role selection --run-id selection-windowed-00 \
+  --order ab --output-json /tmp/selection-windowed-00.json
+```
+
+In a minimal PyTorch container without slime's Ray/Megatron control-plane
+dependencies, set `SLIME_CALLSITE_SOURCE_LOAD=1`. That opt-in mode loads the
+exact production source file while stubbing only unused actor/conversion
+imports; the artifact records `source_with_control_plane_stubs` so it cannot be
+silently mixed with an installed-runtime campaign.
+When the checkout's `.git` metadata is not mounted into the container, pass the
+tested revision as `SLIME_BENCHMARK_SOURCE_COMMIT`.
+
+Run every policy in both `selection` and `confirmation` roles in separate
+process launches (five launches per policy/role by default), alternating
+`--order ab` and `--order ba`. Then validate and summarize the immutable raw
+artifacts:
+
+```bash
+python tools/benchmark_weight_sync_callsite.py \
+  --summarize /tmp/weight-sync-callsite/*.json \
+  --min-runs-per-role 5 \
+  --summary-json /tmp/weight-sync-callsite-summary.json
+```
+
+The summary fails closed on duplicate process/run identity, incomplete rank
+coverage, payload mismatch, mixed runtime/message/topology cells, or too few
+independent runs. It reports communication A/B, rank-local pair makespan,
+receiver consumer waits, callsite return, device readiness, and whole-stage
+sync readiness separately. NCCL intervals are labeled `event_bracket`, never
+`kernel_observed`. The tool records PyTorch/CUDA/NCCL versions, launch-order
+configuration, dtype, message geometry, graph state, hostname, device identity,
+and exact process-group membership.
+
+This evidence does not establish end-to-end training throughput, SGLang load
+latency, multi-node behavior, or a production policy winner. The summarizer
+does not automatically select or apply a policy.

docs/en/index.rst

@@ -86,6 +86,7 @@ Start by Use Case
    advanced/pd-disaggregation.md
    advanced/external-rollout-engines.md
    advanced/delta-weight-sync.md
+   advanced/weight-sync-wave-scheduling.md
    advanced/sglang-config.md
    advanced/megatron-config.md
    advanced/arch-support-beyond-megatron.md

docs/zh/advanced/weight-sync-wave-scheduling.md

@@ -0,0 +1,91 @@
+# 权重同步 Engine Wave 调度
+
+大规模 rollout 部署可能让所有 SGLang engine 同时接收或加载同一个权重
+bucket。这样可以最大化 fan-out，但也可能瞬间占满 GPU、PCIe、NVLink、网络
+或存储带宽。slime 可以在不引入硬件相关固定延迟的情况下限制这种 fan-out：
+
+```bash
+--update-weight-max-inflight-engine-groups 2
+```
+
+这里的 engine group 指一个逻辑 SGLang engine。它内部的 tensor parallel、
+pipeline parallel 或多节点 worker 始终作为不可拆分的整体。调度器按解析后的
+engine 列表稳定分波，每波最多放行配置数量。例如四个 engine、上限为三时，
+wave 依次为 `(0, 1, 2)` 和 `(3)`；实现不假设特定 world size 或 engine
+数量。
+
+默认值为 `0`，保持原有 all-at-once 行为。配置值大于等于 engine 数量时也
+等价于默认行为；负数会在启动参数校验阶段被拒绝。
+
+## 支持的传输路径
+
+- 非 colocate NCCL 权重同步只在确实要求并发上限时，为每个逻辑 engine
+  创建独立 process group。同一 wave 内的 broadcast 异步发起，NCCL work
+  与 engine 侧加载请求全部完成后才放行下一波；默认行为仍使用原来的聚合
+  process group。
+- Colocate tensor/IPC 同步在所有 trainer rank 间使用相同的确定性 wave；
+  每波结束用 trainer 控制组 barrier 封口，然后才允许复用源 buffer。
+- Full 和 delta disk 同步会对 checkpoint pull 与 engine reload 请求应用
+  同一并发上限。
+- 量化模型的 post-load 处理也遵守同一 engine-group 上限。
+
+现有的 pause、flush 与权重版本边界保持不变：某个 engine 不会仅仅因为自己的
+wave 先完成就提前恢复 generation。disk checkpoint pull 可以在 pause 前预取，
+但该过程不会修改 serving 权重。因此 wave 放行本身不会引入新的新旧权重混用
+窗口。
+
+## 如何选择上限
+
+先用 `0` 建立吞吐基线，再在实际部署拓扑上比较 `1`、`2`、`4` 等较小上限。
+应选择既能消除观测到的流量或显存峰值、又不会增加权重同步总时间或尾延迟的
+最大值。该参数只限制并发 engine group 数量；每个权重 bucket 的大小仍由
+`--update-weight-buffer-size` 独立控制。
+
+## 生产调用点确认
+
+`tools/benchmark_weight_sync_callsite.py` 直接调用生产函数
+`update_weights_in_engine_group_waves`，使用真实 process group 和合成 payload。
+它是调用点/模块级探针：不会重写调度器、改变默认值，也不声称测到了
+Ray/SGLang 加载性能。
+
+四进程启动会保留三个真实的二 rank `[trainer, engine]` 通信组。A/B 是
+engine group 0 和 1；第三组是区分 all-at-once 与 window-2 的竞争操作。
+只使用当前测试环境允许的设备数量。
+
+下列每条命令只生成一个独立进程运行制品：
+
+```bash
+torchrun --standalone --nproc-per-node=4 \
+  tools/benchmark_weight_sync_callsite.py \
+  --backend gloo --policy candidate_windowed \
+  --evidence-role selection --run-id selection-windowed-00 \
+  --order ab --output-json /tmp/selection-windowed-00.json
+```
+
+若最小 PyTorch 容器缺少 slime 的 Ray/Megatron 控制面依赖，可显式设置
+`SLIME_CALLSITE_SOURCE_LOAD=1`。该模式加载完全相同的生产源文件，只替换本
+探针不会调用的 actor/conversion import；制品会记录
+`source_with_control_plane_stubs`，不能静默混入完整运行时 campaign。
+若容器内没有挂载 checkout 的 `.git` 元数据，请用
+`SLIME_BENCHMARK_SOURCE_COMMIT` 显式传入被测 revision。
+
+对每个 policy 分别以 `selection`、`confirmation` 角色启动独立进程（默认
+每个 policy/role 五次），并交替使用 `--order ab`、`--order ba`。随后校验、
+汇总不可变原始制品：
+
+```bash
+python tools/benchmark_weight_sync_callsite.py \
+  --summarize /tmp/weight-sync-callsite/*.json \
+  --min-runs-per-role 5 \
+  --summary-json /tmp/weight-sync-callsite-summary.json
+```
+
+遇到重复进程/运行身份、rank 覆盖不完整、payload 不一致、混用运行时/
+消息/拓扑 cell，或独立运行不足时，汇总器会 fail closed。它分别报告通信
+A/B、rank-local pair makespan、接收端 consumer wait、调用点返回、设备
+ready 和整个同步阶段 ready。NCCL 区间只标记为 `event_bracket`，绝不冒充
+`kernel_observed`。工具记录 PyTorch/CUDA/NCCL 版本、launch-order 配置、
+dtype、消息几何、graph 状态、主机/设备身份与精确 PG membership。
+
+这些证据不能证明端到端训练吞吐、SGLang 加载延迟、多机行为或某个生产
+策略必胜；汇总器不会自动选择或应用策略。

docs/zh/index.rst

@@ -86,6 +86,7 @@ slime 的设计目标，是让这两大能力彼此强化，同时避免把系
    advanced/pd-disaggregation.md
    advanced/external-rollout-engines.md
    advanced/delta-weight-sync.md
+   advanced/weight-sync-wave-scheduling.md
    advanced/sglang-config.md
    advanced/megatron-config.md
    advanced/arch-support-beyond-megatron.md

slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py

@@ -22,6 +22,7 @@
 from slime.utils import accelerator
 from slime.utils.disk_delta import NUM_WORKERS, checksum, make_tensor_reader, overwrite_encode
 from slime.utils.distributed_utils import get_gloo_group
+from slime.utils.engine_group_wave import run_engine_group_waves
 
 from .update_weight_from_distributed import UpdateWeightFromDistributed
 
@@ -107,7 +108,16 @@ def _capture_baseline(self) -> None:
             os.makedirs(self.delta_dir, exist_ok=True)
             if self._post_write_hook is not None:
                 self._post_write_hook(self.args, self.delta_dir, list(self.rollout_engines))
-            pulls = [engine.pull_weights.remote(target_version=0) for engine in self.rollout_engines]
+            max_inflight_engine_groups = self.args.update_weight_max_inflight_engine_groups
+            if 0 < max_inflight_engine_groups < len(self.rollout_engines):
+                run_engine_group_waves(
+                    self.rollout_engines,
+                    max_inflight_engine_groups,
+                    lambda _index, engine: engine.pull_weights.remote(target_version=0),
+                    ray.get,
+                )
+            else:
+                pulls = [engine.pull_weights.remote(target_version=0) for engine in self.rollout_engines]
         dist.barrier(group=get_gloo_group())
 
         read_hf = make_tensor_reader(self.args.hf_checkpoint)  # index the HF headers once
@@ -118,7 +128,8 @@ def _capture_baseline(self) -> None:
                 self._snapshot[name] = tensor.detach().cpu().contiguous().view(torch.uint8).numpy().reshape(-1)
                 logger.warning("seed: %s absent from hf_checkpoint; seeding from current weights", name)
         if dist.get_rank() == 0:
-            ray.get(pulls)
+            if pulls:
+                ray.get(pulls)
             logger.info(
                 "[disk delta] captured baseline snapshot of %d tensors from %s",
                 len(self._snapshot),
@@ -175,17 +186,23 @@ def _reload_engines(self) -> None:
             self._post_write_hook(self.args, self._version_dir, list(self.rollout_engines))
         dist.barrier(group=get_gloo_group())
         if dist.get_rank() == 0:
-            ray.get([engine.pull_weights.remote(self.weight_version) for engine in self.rollout_engines])
+            max_inflight_engine_groups = self.args.update_weight_max_inflight_engine_groups
+            run_engine_group_waves(
+                self.rollout_engines,
+                max_inflight_engine_groups,
+                lambda _index, engine: engine.pull_weights.remote(self.weight_version),
+                ray.get,
+            )
             ray.get([engine.pause_generation.remote() for engine in self.rollout_engines])
             ray.get([engine.flush_cache.remote() for engine in self.rollout_engines])
-            ray.get(
-                [
-                    engine.update_weights_from_disk.remote(
-                        model_path=self.args.update_weight_local_checkpoint_dir,
-                        weight_version=str(self.weight_version),
-                    )
-                    for engine in self.rollout_engines
-                ]
+            run_engine_group_waves(
+                self.rollout_engines,
+                max_inflight_engine_groups,
+                lambda _index, engine: engine.update_weights_from_disk.remote(
+                    model_path=self.args.update_weight_local_checkpoint_dir,
+                    weight_version=str(self.weight_version),
+                ),
+                ray.get,
             )
             ray.get([engine.continue_generation.remote() for engine in self.rollout_engines])
         dist.barrier(group=get_gloo_group())