[vllm, checkpoint] feat: add rollout weight sync debug check

docs/advance/checkpoint.rst

@@ -286,6 +286,33 @@ Recommended Configurations
   e.g. ``save_contents=['model', 'hf_model', 'optimizer', 'extra']`` to write both
   ``model/dist_ckpt/`` and ``model/huggingface/`` in one step.
 
+Check Trainer-to-Rollout Weight Sync
+------------------------------------
+
+When adapting a model or checkpoint-engine backend, enable the weight-sync check
+on a PPO/GRPO command. The check runs after the initial
+``CheckpointEngineManager.update_weights()`` call and after the rollout backend
+has loaded the weights from the source HuggingFace checkpoint:
+
+.. code:: bash
+
+    actor_rollout_ref.rollout.name=vllm \
+    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+    actor_rollout_ref.rollout.data_parallel_size=1 \
+    actor_rollout_ref.rollout.checkpoint_engine.backend=nccl \
+    actor_rollout_ref.rollout.checkpoint_engine.check_weight_sync=True
+
+The check compares the initial rollout backend-loaded weights with the source
+HuggingFace checkpoint and reports missing, unexpected, and mismatched keys.
+Strict equality currently supports vLLM with rollout tensor and data parallel
+sizes 1, after casting the reference checkpoint tensors to the loaded rollout
+tensor dtype.
+
+This check validates the checkpoint-engine weight conversion and transfer path,
+including vLLM ``load_weights`` for tensor and data parallel sizes 1. It does
+not validate rollout generation quality or backend-specific inference kernel
+behavior.
+
 Convert FSDP and Megatron Checkpoints to HuggingFace Format Model
 -----------------------------------------------------------------
 

tests/checkpoint_engine/test_weight_sync_on_cpu.py

@@ -0,0 +1,150 @@
+# Copyright 2024 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.
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from verl.checkpoint_engine.base import CheckpointEngineManager
+from verl.checkpoint_engine.weight_sync import assert_vllm_weight_sync_supported_parallelism, assert_weight_sync_equal
+from verl.workers.config import CheckpointEngineConfig
+
+
+def test_checkpoint_engine_manager_checks_initial_sync_only():
+    manager = object.__new__(CheckpointEngineManager)
+    manager.config = CheckpointEngineConfig(check_weight_sync=True)
+
+    assert manager._should_check_loaded_weights_equal(None)
+    assert manager._should_check_loaded_weights_equal(0)
+    assert not manager._should_check_loaded_weights_equal(1)
+
+
+def test_vllm_weight_sync_parallelism_check_accepts_tp1_dp1():
+    parallel_config = SimpleNamespace(tensor_parallel_size=1, data_parallel_size=1, data_parallel_size_local=1)
+
+    assert_vllm_weight_sync_supported_parallelism(parallel_config)
+
+
+@pytest.mark.parametrize(
+    "parallel_config",
+    [
+        SimpleNamespace(tensor_parallel_size=2, data_parallel_size=1, data_parallel_size_local=1),
+        SimpleNamespace(tensor_parallel_size=1, data_parallel_size=2, data_parallel_size_local=1),
+        SimpleNamespace(tensor_parallel_size=1, data_parallel_size=2, data_parallel_size_local=2),
+        SimpleNamespace(tensor_parallel_size=1, data_parallel_size=1, data_parallel_size_local=2),
+    ],
+)
+def test_vllm_weight_sync_parallelism_check_rejects_uncovered_parallelism(parallel_config):
+    with pytest.raises(NotImplementedError, match="tensor_parallel_size=1 and data_parallel_size=1"):
+        assert_vllm_weight_sync_supported_parallelism(parallel_config)
+
+
+def test_assert_weight_sync_equal_accepts_direct_state_dict():
+    expected_state_dict = {"model.embed_tokens.weight": torch.arange(6, dtype=torch.float32).view(2, 3)}
+    actual_state_dict = {"model.embed_tokens.weight": expected_state_dict["model.embed_tokens.weight"].clone()}
+
+    result = assert_weight_sync_equal(
+        expected_state_dict=expected_state_dict,
+        actual_state_dict=actual_state_dict,
+        num_hidden_layers=0,
+    )
+
+    assert result == {"checked": 1, "missing": 0, "unexpected": 0, "mismatched": 0}
+
+
+def test_assert_weight_sync_equal_rejects_any_value_mismatch():
+    expected_state_dict = {"model.embed_tokens.weight": torch.ones(2, 3)}
+    actual_state_dict = {"model.embed_tokens.weight": torch.ones(2, 3)}
+    actual_state_dict["model.embed_tokens.weight"][0, 0] += 1e-6
+
+    with pytest.raises(AssertionError, match="mismatched"):
+        assert_weight_sync_equal(
+            expected_state_dict=expected_state_dict,
+            actual_state_dict=actual_state_dict,
+            num_hidden_layers=0,
+        )
+
+
+def test_assert_weight_sync_equal_accepts_vllm_fused_state_dict():
+    prefix = "model.layers.0"
+    q_proj = torch.arange(8, dtype=torch.float32).view(2, 4)
+    k_proj = q_proj + 10
+    v_proj = q_proj + 20
+    gate_proj = torch.arange(12, dtype=torch.float32).view(3, 4)
+    up_proj = gate_proj + 30
+    q_bias = torch.arange(2, dtype=torch.float32)
+    k_bias = q_bias + 10
+    v_bias = q_bias + 20
+    embed_tokens = torch.arange(20, dtype=torch.float32).view(5, 4)
+
+    expected_state_dict = {
+        f"{prefix}.self_attn.q_proj.weight": q_proj,
+        f"{prefix}.self_attn.k_proj.weight": k_proj,
+        f"{prefix}.self_attn.v_proj.weight": v_proj,
+        f"{prefix}.self_attn.q_proj.bias": q_bias,
+        f"{prefix}.self_attn.k_proj.bias": k_bias,
+        f"{prefix}.self_attn.v_proj.bias": v_bias,
+        f"{prefix}.mlp.gate_proj.weight": gate_proj,
+        f"{prefix}.mlp.up_proj.weight": up_proj,
+        "model.embed_tokens.weight": embed_tokens,
+        "lm_head.weight": embed_tokens,
+    }
+    actual_state_dict = {
+        f"{prefix}.self_attn.qkv_proj.weight": torch.cat([q_proj, k_proj, v_proj], dim=0),
+        f"{prefix}.self_attn.qkv_proj.bias": torch.cat([q_bias, k_bias, v_bias], dim=0),
+        f"{prefix}.mlp.gate_up_proj.weight": torch.cat([gate_proj, up_proj], dim=0),
+        "model.embed_tokens.weight": embed_tokens.clone(),
+    }
+
+    result = assert_weight_sync_equal(
+        expected_state_dict=expected_state_dict,
+        actual_state_dict=actual_state_dict,
+        num_hidden_layers=1,
+        tie_word_embeddings=True,
+    )
+
+    assert result == {"checked": 10, "missing": 0, "unexpected": 0, "mismatched": 0}
+
+
+def test_assert_weight_sync_equal_reports_fused_value_mismatch():
+    prefix = "model.layers.0"
+    q_proj = torch.ones(2, 4)
+    k_proj = torch.ones(2, 4) * 2
+    v_proj = torch.ones(2, 4) * 3
+    actual_qkv = torch.cat([q_proj, k_proj, v_proj], dim=0)
+    actual_qkv[0, 0] += 1
+
+    expected_state_dict = {
+        f"{prefix}.self_attn.q_proj.weight": q_proj,
+        f"{prefix}.self_attn.k_proj.weight": k_proj,
+        f"{prefix}.self_attn.v_proj.weight": v_proj,
+    }
+    actual_state_dict = {f"{prefix}.self_attn.qkv_proj.weight": actual_qkv}
+
+    with pytest.raises(AssertionError, match="mismatched"):
+        assert_weight_sync_equal(
+            expected_state_dict=expected_state_dict,
+            actual_state_dict=actual_state_dict,
+            num_hidden_layers=1,
+        )
+
+
+def test_assert_weight_sync_equal_reports_missing_and_unexpected_keys():
+    with pytest.raises(AssertionError, match="missing=.*unexpected="):
+        assert_weight_sync_equal(
+            expected_state_dict={"expected.weight": torch.ones(1)},
+            actual_state_dict={"actual.weight": torch.ones(1)},
+            num_hidden_layers=0,
+        )

verl/checkpoint_engine/README.md

@@ -52,7 +52,29 @@ In whitelist json file, you need to add IP addresses using the following format.
 export HCCL_INTRA_ROCE_ENABLE=1
 ```
 
+### Weight sync check
+
+Enable `check_weight_sync` on a PPO/GRPO command to verify the initial
+trainer-to-rollout weight synchronization after `CheckpointEngineManager.update_weights()`
+and the rollout backend's `load_weights()` path. Add these overrides to an existing
+training command:
+
+```bash
+actor_rollout_ref.rollout.name=vllm \
+actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+actor_rollout_ref.rollout.data_parallel_size=1 \
+actor_rollout_ref.rollout.checkpoint_engine.backend=nccl \
+actor_rollout_ref.rollout.checkpoint_engine.check_weight_sync=True
+```
+
+The check currently supports vLLM strict equality with rollout tensor and data parallel sizes 1.
+It reports missing, unexpected, and mismatched keys when the initial backend-loaded
+weights do not match the source HuggingFace checkpoint.
+The comparison is strict after casting the reference checkpoint tensors to the
+loaded rollout tensor dtype.
+
 ### Benchmark
+
 1. benchmark setup
 - model: Qwen/Qwen3-30B-A3B-Base
 - actor: fsdp world_size=2 (since Ascend 910C has 64GB of HBM, we set world_size=4)

verl/checkpoint_engine/base.py

@@ -373,6 +373,23 @@ def is_leader_rank(self) -> bool:
         """Get leader rank flag from the underlying rollout server adapter."""
         return self.server_adapter.is_leader_rank
 
+    @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+    async def check_loaded_weights_equal(
+        self,
+        model_path: str | None = None,
+        trust_remote_code: bool | None = None,
+    ):
+        """Check backend-loaded rollout weights against a HuggingFace checkpoint."""
+        check_fn = getattr(self.server_adapter, "check_loaded_weights_equal", None)
+        if check_fn is None:
+            raise NotImplementedError(
+                f"Rollout backend {self.rollout_config.name!r} does not support loaded-weight equality checks."
+            )
+        return await check_fn(
+            model_path=model_path or self.model_config.path,
+            trust_remote_code=self.model_config.trust_remote_code if trust_remote_code is None else trust_remote_code,
+        )
+
 
 _worker_cls = ray.remote(CheckpointEngineWorker)
 
@@ -463,6 +480,22 @@ def remove_replicas(self, replicas: list[RolloutReplica]):
         replicas_set = set(replicas)
         self.replicas = [r for r in self.replicas if r not in replicas_set]
 
+    def check_loaded_weights_equal(self, rollout: RayWorkerGroup):
+        """Check rollout backend-loaded weights when debug checking is enabled."""
+        if not self.config.check_weight_sync:
+            return None
+
+        summaries = rollout.check_loaded_weights_equal()
+        summaries = [summary for summary in summaries if summary is not None]
+        print(f"[CheckpointEngineManager] weight sync check passed: {summaries}")
+        return summaries
+
+    def _should_check_loaded_weights_equal(self, global_steps: int | None) -> bool:
+        """Whether this update should be compared with the source HF checkpoint."""
+        if not self.config.check_weight_sync:
+            return False
+        return global_steps in (None, 0)
+
     @auto_await
     async def sleep_replicas(self):
         """Sleep all rollout replicas: free weight and kv_cache device memory."""
@@ -517,6 +550,10 @@ 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))
+            if self._should_check_loaded_weights_equal(global_steps):
+                raise NotImplementedError(
+                    "Weight sync checking is not supported for checkpoint engine backend='naive'."
+                )
             return {}
 
         # 1. abort and save all unfinished requests for partial rollout
@@ -553,10 +590,14 @@ async def update_weights(self, global_steps: int = None):
             + rollout.execute_checkpoint_engine(["finalize"] * rollout.world_size)
         )
 
-        # 7. restore kv_cache after weight sync
+        # 7. optionally validate the initial backend-loaded weights before generation resumes
+        if self._should_check_loaded_weights_equal(global_steps):
+            self.check_loaded_weights_equal(rollout)
+
+        # 8. restore kv_cache after weight sync
         await self.resume_kv_cache_replicas()
 
-        # 8. resume all unfinished requests for partial rollout
+        # 9. resume all unfinished requests for partial rollout
         await self.resume_generation_replicas()
 
         return sync_metrics