[ci] Add e2e test for delta weight update

.github/workflows/pr-test.yml

@@ -39,7 +39,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_fanout_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_fanout_short.py"}, {"num_gpus": 4, "test_file": "test_delta_weight_update.py"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}
@@ -205,7 +205,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        info: [{"enable_eval": "0", "num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_streaming_partial_rollout.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}]
+        info: [{"enable_eval": "0", "num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_streaming_partial_rollout.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.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_fanout_short.py', 'num_gpus': 4},
+        {'test_file': 'test_delta_weight_update.py', 'num_gpus': 4},
       ],
     },
     'e2e-test-sglang-config': {
@@ -25,7 +26,6 @@
         {'test_file': 'test_qwen3_30B_A3B.py', 'num_gpus': 8, 'use_deepep': '1', 'use_fp8_rollout': '1', 'enable_eval': '0'},
         {'test_file': 'test_qwen3.6_35B_A3B_pd_mooncake.py', 'num_gpus': 8, 'use_deepep': '1'},
         {'test_file': 'test_qwen3_30B_A3B_r3.py', 'num_gpus': 8, 'use_deepep': '1', 'use_fp8_rollout': '1', 'enable_eval': '0'},
-        {'test_file': 'test_qwen3_30B_A3B_r3.py', 'num_gpus': 8, 'enable_eval': '0'},
         {'test_file': 'test_qwen3_4B_ppo.py', 'num_gpus': 8, 'enable_eval': '0'},
         {'test_file': 'test_qwen3_4B_ppo_disaggregate.py', 'num_gpus': 8, 'enable_eval': '0'},
         {'test_file': 'test_qwen3_4B_ppo_train_critic_only.py', 'num_gpus': 8, 'enable_eval': '0'},

slime/backends/megatron_utils/update_weight/update_weight_from_distributed_delta.py

@@ -519,6 +519,7 @@ def __init__(
         # on rank 0; ``_finalize_sync`` awaits them at end of sync.
         self._pending_files: list[str] = []
         self._pending_publishes: list = []
+        self._published_any: bool = False
         self._rpc_executor: ThreadPoolExecutor | None = None
         if self.transport == "disk":
             self.delta_dir = args.update_weight_delta_dir
@@ -576,6 +577,11 @@ def update_weights(self) -> None:
         if not self._snapshot_seeded:
             self._seed_snapshot()
             self._snapshot_seeded = True
+            # Pin the engine's recorded version to ours (0) on the seed call so the
+            # CI version-equality check holds before any real sync has happened.
+            if dist.get_rank() == 0 and self.transport == "disk" and self.rollout_engines:
+                weight_version = str(self.weight_version)
+                ray.get([engine.set_weight_version.remote(weight_version) for engine in self.rollout_engines])
             return
 
         self.weight_version += 1
@@ -592,6 +598,7 @@ def update_weights(self) -> None:
         self.density_nnz = self.density_numel = self.wire_bytes = self._flush_idx = 0
         self._pending_files.clear()
         self._pending_publishes.clear()
+        self._published_any = False
         if self.writer is not None:
             self.writer.reset_counters()
         pbar = tqdm(desc=f"[{self._group_name}] Update weights", total=0) if self._is_pp_src_rank else None
@@ -771,6 +778,7 @@ def _publish_batch(self) -> None:
             version_dir = self._version_dir
             engines = list(self.rollout_engines)
             weight_version = str(self.weight_version)
+            self._published_any = True
 
             def _fire_when_committed() -> list:
                 if commit_future is not None:
@@ -808,6 +816,12 @@ def _finalize_sync(self) -> None:
             object_refs = [ref for fut in self._pending_publishes for ref in fut.result()]
             ray.get(object_refs)
             self._pending_publishes.clear()
+            if not self._published_any:
+                # No delta files needed publishing this sync (e.g. all-zero diff).
+                # Engines never saw the new version via update_weights_from_disk, so
+                # bump it explicitly to keep their recorded version in sync with ours.
+                weight_version = str(self.weight_version)
+                ray.get([engine.set_weight_version.remote(weight_version) for engine in self.rollout_engines])
             if not self.args.update_weight_delta_keep_files:
                 shutil.rmtree(self._version_dir, ignore_errors=True)
             ray.get([engine.continue_generation.remote() for engine in self.rollout_engines])

slime/backends/sglang_utils/sglang_engine.py

@@ -353,6 +353,15 @@ def get_weight_version(self):
         response.raise_for_status()
         return response.json()["weight_version"]
 
+    def set_weight_version(self, new_version: str):
+        """Bump the engine's recorded weight version without changing weights.
+
+        Used by the delta-update path when a sync produced no bytes (e.g. an
+        all-zero diff): we still need the engine's version to track the
+        updater's, otherwise the CI version-equality check will trip.
+        """
+        return self._make_request("update_weight_version", {"new_version": str(new_version)})
+
     def release_memory_occupation(self):
         self.flush_cache()
         return self._make_request("release_memory_occupation")

tests/test_delta_weight_update.py

@@ -0,0 +1,143 @@
+"""E2E smoke test for disk-backed delta weight updates.
+
+Runs a tiny Qwen3.5-0.8B job so the first weight update seeds the delta
+snapshot and the post-train update publishes sparse delta files through
+``update_weights_from_disk(load_format="delta", files=...)``.
+"""
+
+import os
+import tempfile
+from pathlib import Path
+
+import slime.utils.external_utils.command_utils as U
+
+
+MODEL_NAME = "Qwen3.5-0.8B"
+MODEL_TYPE = "qwen3.5-0.8B"
+NUM_GPUS = 4
+TORCH_DIST_CKPT = f"/dev/shm/{MODEL_NAME}_torch_dist"
+
+
+def prepare():
+    U.exec_command("mkdir -p /root/models /root/datasets")
+    U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
+    U.hf_download_dataset("zhuzilin/gsm8k")
+    U.convert_checkpoint(
+        model_name=MODEL_NAME,
+        megatron_model_type=MODEL_TYPE,
+        num_gpus_per_node=NUM_GPUS,
+        dir_dst="/dev/shm",
+    )
+
+
+def execute():
+    with tempfile.TemporaryDirectory(prefix="slime_delta_weight_update_") as delta_dir:
+        ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load {TORCH_DIST_CKPT} "
+
+        rollout_args = (
+            "--prompt-data /root/datasets/gsm8k/train.parquet "
+            "--input-key messages "
+            "--label-key label "
+            "--apply-chat-template "
+            "--rollout-shuffle "
+            "--rm-type math "
+            "--num-rollout 1 "
+            "--rollout-batch-size 4 "
+            "--n-samples-per-prompt 4 "
+            "--rollout-max-response-len 1024 "
+            "--rollout-temperature 0.8 "
+            "--over-sampling-batch-size 8 "
+            "--dynamic-sampling-filter-path slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std "
+            "--global-batch-size 16 "
+        )
+
+        perf_args = (
+            "--tensor-model-parallel-size 1 "
+            "--sequence-parallel "
+            "--pipeline-model-parallel-size 1 "
+            "--context-parallel-size 1 "
+            "--expert-model-parallel-size 1 "
+            "--expert-tensor-parallel-size 1 "
+            "--use-dynamic-batch-size "
+            "--max-tokens-per-gpu 9216 "
+        )
+
+        grpo_args = (
+            "--advantage-estimator grpo "
+            "--use-kl-loss "
+            "--kl-loss-coef 0.00 "
+            "--kl-loss-type low_var_kl "
+            # Nonzero entropy coef guarantees a nonzero gradient even when all
+            # rewards in a group tie (advantages=0), so the delta sync writes
+            # real sparse files instead of an empty no-op.
+            "--entropy-coef 0.01 "
+            "--eps-clip 0.2 "
+            "--eps-clip-high 0.28 "
+        )
+
+        optimizer_args = (
+            "--optimizer adam "
+            "--lr 1e-6 "
+            "--lr-decay-style constant "
+            "--weight-decay 0.1 "
+            "--adam-beta1 0.9 "
+            "--adam-beta2 0.98 "
+        )
+
+        sglang_args = (
+            "--rollout-num-gpus-per-engine 1 "
+            "--rollout-num-gpus 3 "
+            "--sglang-mem-fraction-static 0.7 "
+            "--sglang-cuda-graph-max-bs 32 "
+            "--sglang-enable-metrics "
+        )
+
+        delta_args = (
+            "--update-weight-mode delta "
+            "--update-weight-transport disk "
+            "--update-weight-encoding deltas "
+            f"--update-weight-delta-dir {delta_dir} "
+            "--update-weight-delta-keep-files "
+        )
+
+        ci_args = "--ci-test "
+
+        misc_args = (
+            "--attention-dropout 0.0 "
+            "--hidden-dropout 0.0 "
+            "--accumulate-allreduce-grads-in-fp32 "
+            "--attention-softmax-in-fp32 "
+            "--attention-backend flash "
+            "--loss-mask-type qwen3_5 "
+            "--actor-num-nodes 1 "
+            "--actor-num-gpus-per-node 1 "
+        )
+
+        train_args = (
+            f"{ckpt_args} "
+            f"{rollout_args} "
+            f"{optimizer_args} "
+            f"{grpo_args} "
+            f"{U.get_default_wandb_args(__file__)} "
+            f"{perf_args} "
+            f"{sglang_args} "
+            f"{delta_args} "
+            f"{ci_args} "
+            f"{misc_args} "
+        )
+
+        U.execute_train(
+            train_args=train_args,
+            num_gpus_per_node=NUM_GPUS,
+            megatron_model_type=MODEL_TYPE,
+        )
+
+        delta_files = list(Path(delta_dir).glob("weight_v*/*.safetensors"))
+        assert delta_files, f"No disk delta safetensors were written under {delta_dir}"
+
+
+if __name__ == "__main__":
+    prepare()
+    for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
+        os.environ.pop(proxy_var, None)
+    execute()