From e903d5b8b9d31827c98c4cffaaa2c1bf1df451fd Mon Sep 17 00:00:00 2001 From: caozs2 Date: Wed, 21 Oct 2026 15:45:07 +0800 Subject: [PATCH] [vllm, checkpoint] feat: add rollout weight sync debug check Rebase the vLLM TP=1/DP=1 loaded-weight check onto current main. Co-authored-by: Grok 4.7 Co-authored-by: Cursor --- docs/advance/checkpoint.rst | 27 ++++ .../test_weight_sync_on_cpu.py | 150 ++++++++++++++++++ verl/checkpoint_engine/README.md | 22 +++ verl/checkpoint_engine/base.py | 45 +++++- verl/checkpoint_engine/weight_sync.py | 136 ++++++++++++++++ .../_generated_ppo_megatron_trainer.yaml | 1 + .../_generated_ppo_torchtitan_trainer.yaml | 1 + .../config/_generated_ppo_trainer.yaml | 1 + .../config/_generated_ppo_veomni_trainer.yaml | 1 + verl/trainer/config/rollout/rollout.yaml | 5 + verl/workers/config/rollout.py | 2 + verl/workers/rollout/vllm_rollout/utils.py | 42 +++++ .../rollout/vllm_rollout/vllm_rollout.py | 14 ++ 13 files changed, 445 insertions(+), 2 deletions(-) create mode 100644 tests/checkpoint_engine/test_weight_sync_on_cpu.py create mode 100644 verl/checkpoint_engine/weight_sync.py diff --git a/docs/advance/checkpoint.rst b/docs/advance/checkpoint.rst index 4b31446e1d1..472987ecce9 100644 --- a/docs/advance/checkpoint.rst +++ b/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 ----------------------------------------------------------------- diff --git a/tests/checkpoint_engine/test_weight_sync_on_cpu.py b/tests/checkpoint_engine/test_weight_sync_on_cpu.py new file mode 100644 index 00000000000..7a5834dae16 --- /dev/null +++ b/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, + ) diff --git a/verl/checkpoint_engine/README.md b/verl/checkpoint_engine/README.md index 305642c200e..36c4f6f002b 100644 --- a/verl/checkpoint_engine/README.md +++ b/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) diff --git a/verl/checkpoint_engine/base.py b/verl/checkpoint_engine/base.py index c47e5184171..0ead8886285 100644 --- a/verl/checkpoint_engine/base.py +++ b/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 diff --git a/verl/checkpoint_engine/weight_sync.py b/verl/checkpoint_engine/weight_sync.py new file mode 100644 index 00000000000..65ad53ef366 --- /dev/null +++ b/verl/checkpoint_engine/weight_sync.py @@ -0,0 +1,136 @@ +# 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 typing import Any + +import torch + + +def assert_vllm_weight_sync_supported_parallelism(parallel_config: Any) -> None: + """Validate the vLLM parallelism modes covered by loaded-weight checks.""" + tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1) or 1) + dp_size = int(getattr(parallel_config, "data_parallel_size", 1) or 1) + dp_local_size = int(getattr(parallel_config, "data_parallel_size_local", 1) or 1) + if tp_size != 1 or dp_size != 1 or dp_local_size != 1: + raise NotImplementedError( + "vLLM loaded-weight equality check currently supports " + "tensor_parallel_size=1 and data_parallel_size=1 only. " + f"Got tensor_parallel_size={tp_size}, data_parallel_size={dp_size}, " + f"data_parallel_size_local={dp_local_size}." + ) + + +def assert_weight_sync_equal( + expected_state_dict: dict[str, torch.Tensor], + actual_state_dict: dict[str, torch.Tensor], + num_hidden_layers: int, + tie_word_embeddings: bool = False, + max_mismatches: int = 10, +) -> dict[str, Any]: + """Compare source HF weights against backend-loaded weights. + + The comparison handles common vLLM fused weights for QKV and SwiGLU gate/up + projections. It raises AssertionError with missing, unexpected, and + mismatched samples when the loaded weights do not match the source weights. + """ + + expected_names = set(expected_state_dict) + actual_names = set(actual_state_dict) + compared_expected_names = set() + compared_actual_names = set() + mismatched = [] + + def check_tensor(name: str, expected_weight: torch.Tensor, actual_weight: torch.Tensor) -> None: + actual_weight = actual_weight.detach() + if expected_weight.shape != actual_weight.shape: + mismatched.append( + f"{name}: shape expected={tuple(expected_weight.shape)} actual={tuple(actual_weight.shape)}" + ) + return + + expected_weight = expected_weight.to(device=actual_weight.device, dtype=actual_weight.dtype) + if not torch.equal(expected_weight, actual_weight): + max_abs_diff = (expected_weight - actual_weight).abs().max().item() + mismatched.append(f"{name}: max_abs_diff={max_abs_diff}") + + for name, expected_weight in expected_state_dict.items(): + if name not in actual_state_dict: + continue + check_tensor(name, expected_weight, actual_state_dict[name]) + compared_expected_names.add(name) + compared_actual_names.add(name) + + def check_fused_tensor(actual_name: str, expected_component_names: list[str]) -> None: + if actual_name not in actual_state_dict: + return + if any(name not in expected_state_dict for name in expected_component_names): + return + expected_weight = torch.cat([expected_state_dict[name] for name in expected_component_names], dim=0) + check_tensor(actual_name, expected_weight, actual_state_dict[actual_name]) + compared_expected_names.update(expected_component_names) + compared_actual_names.add(actual_name) + + for layer_idx in range(num_hidden_layers): + prefix = f"model.layers.{layer_idx}" + check_fused_tensor( + actual_name=f"{prefix}.self_attn.qkv_proj.weight", + expected_component_names=[ + f"{prefix}.self_attn.q_proj.weight", + f"{prefix}.self_attn.k_proj.weight", + f"{prefix}.self_attn.v_proj.weight", + ], + ) + check_fused_tensor( + actual_name=f"{prefix}.self_attn.qkv_proj.bias", + expected_component_names=[ + f"{prefix}.self_attn.q_proj.bias", + f"{prefix}.self_attn.k_proj.bias", + f"{prefix}.self_attn.v_proj.bias", + ], + ) + check_fused_tensor( + actual_name=f"{prefix}.mlp.gate_up_proj.weight", + expected_component_names=[ + f"{prefix}.mlp.gate_proj.weight", + f"{prefix}.mlp.up_proj.weight", + ], + ) + + if ( + tie_word_embeddings + and "lm_head.weight" in expected_names + and "lm_head.weight" not in actual_names + and "model.embed_tokens.weight" in compared_expected_names + ): + compared_expected_names.add("lm_head.weight") + + missing = sorted(expected_names - compared_expected_names) + unexpected = sorted(actual_names - compared_actual_names) + + if missing or unexpected or mismatched: + details = [] + if missing: + details.append(f"missing={missing[:max_mismatches]}") + if unexpected: + details.append(f"unexpected={unexpected[:max_mismatches]}") + if mismatched: + details.append(f"mismatched={mismatched[:max_mismatches]}") + raise AssertionError("; ".join(details)) + + return { + "checked": len(compared_expected_names), + "missing": 0, + "unexpected": 0, + "mismatched": 0, + } diff --git a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml index e8ab900fdcb..b5199d6d6c0 100644 --- a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml @@ -386,6 +386,7 @@ actor_rollout_ref: update_weights_bucket_megabytes: 2048 engine_kwargs: {} custom_backend_module: null + check_weight_sync: false trace: _target_: verl.workers.config.TraceConfig project_name: ${oc.select:trainer.project_name,null} diff --git a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml index b8605dc18df..c2b94a685c6 100644 --- a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml @@ -337,6 +337,7 @@ actor_rollout_ref: update_weights_bucket_megabytes: 2048 engine_kwargs: {} custom_backend_module: null + check_weight_sync: false trace: _target_: verl.workers.config.TraceConfig project_name: ${oc.select:trainer.project_name,null} diff --git a/verl/trainer/config/_generated_ppo_trainer.yaml b/verl/trainer/config/_generated_ppo_trainer.yaml index 58af8e1bec8..63595e4bd37 100644 --- a/verl/trainer/config/_generated_ppo_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_trainer.yaml @@ -374,6 +374,7 @@ actor_rollout_ref: update_weights_bucket_megabytes: 2048 engine_kwargs: {} custom_backend_module: null + check_weight_sync: false trace: _target_: verl.workers.config.TraceConfig project_name: ${oc.select:trainer.project_name,null} diff --git a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml index e46389c3520..aab6d9e3584 100644 --- a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml @@ -355,6 +355,7 @@ actor_rollout_ref: update_weights_bucket_megabytes: 2048 engine_kwargs: {} custom_backend_module: null + check_weight_sync: false trace: _target_: verl.workers.config.TraceConfig project_name: ${oc.select:trainer.project_name,null} diff --git a/verl/trainer/config/rollout/rollout.yaml b/verl/trainer/config/rollout/rollout.yaml index d503665c4ba..e5c27a6784a 100644 --- a/verl/trainer/config/rollout/rollout.yaml +++ b/verl/trainer/config/rollout/rollout.yaml @@ -315,6 +315,11 @@ checkpoint_engine: # in CheckpointEngineRegistry. custom_backend_module: null + # Compare initial rollout backend-loaded weights with the source HF checkpoint. + # Currently implemented for vLLM backend checks with tensor_model_parallel_size=1 + # and data_parallel_size=1. + check_weight_sync: False + # trace rollout data trace: diff --git a/verl/workers/config/rollout.py b/verl/workers/config/rollout.py index 4bfef74d226..5578a485528 100644 --- a/verl/workers/config/rollout.py +++ b/verl/workers/config/rollout.py @@ -140,6 +140,8 @@ class CheckpointEngineConfig(BaseConfig): # backend is instantiated, allowing custom backends to register themselves # in CheckpointEngineRegistry. custom_backend_module: Optional[str] = None + # Debug option: compare initial rollout backend-loaded weights with the source HF checkpoint. + check_weight_sync: bool = False @dataclass diff --git a/verl/workers/rollout/vllm_rollout/utils.py b/verl/workers/rollout/vllm_rollout/utils.py index 8112b7e4069..e6e315c4d2a 100644 --- a/verl/workers/rollout/vllm_rollout/utils.py +++ b/verl/workers/rollout/vllm_rollout/utils.py @@ -238,6 +238,48 @@ def _iter_all_models_with_config(self): if draft_cfg is not None: yield self._get_drafter_model(), draft_cfg + @torch.no_grad() + def check_loaded_weights_equal( + self, + model_path: str, + trust_remote_code: bool = False, + ) -> dict[str, Any]: + """Compare vLLM-loaded weights against a HuggingFace checkpoint. + + This debug helper is intended for TP=1 exact checks. With tensor + parallelism, vLLM stores sharded tensors and direct state_dict comparison + is not meaningful without applying each layer's sharding rule. + """ + from transformers import AutoModelForCausalLM + + from verl.checkpoint_engine.weight_sync import ( + assert_vllm_weight_sync_supported_parallelism, + assert_weight_sync_equal, + ) + from verl.utils.fs import copy_to_local + + parallel_config = self.model_runner.vllm_config.parallel_config + # vLLM collective_rpc reaches only one DP shard, so DP>1 is not covered yet. + assert_vllm_weight_sync_supported_parallelism(parallel_config) + + local_path = copy_to_local(model_path) + expected_model = AutoModelForCausalLM.from_pretrained( + local_path, + device_map="cpu", + trust_remote_code=trust_remote_code, + ) + expected_state_dict = expected_model.state_dict() + actual_state_dict = self.model_runner.model.state_dict() + result = assert_weight_sync_equal( + expected_state_dict=expected_state_dict, + actual_state_dict=actual_state_dict, + num_hidden_layers=getattr(expected_model.config, "num_hidden_layers", 0), + tie_word_embeddings=getattr(expected_model.config, "tie_word_embeddings", False), + ) + + del expected_model + return result + def monkey_patch_model(self, vocab_size: int, banned_token_ids: Optional[list[int]] = None): for model in self._iter_all_models(): # patch compute_logits to avoid sampling OOV and other illegal tokens diff --git a/verl/workers/rollout/vllm_rollout/vllm_rollout.py b/verl/workers/rollout/vllm_rollout/vllm_rollout.py index 79a08d5d433..24bf0242367 100644 --- a/verl/workers/rollout/vllm_rollout/vllm_rollout.py +++ b/verl/workers/rollout/vllm_rollout/vllm_rollout.py @@ -252,6 +252,20 @@ async def update_weights( if self.replica_rank == 0 and self.rollout_rank == 0: logger.info(f"update_weights done, time cost: {time.time() - start_time:.2f}s") + async def check_loaded_weights_equal( + self, + model_path: str, + trust_remote_code: bool = False, + ): + """Check vLLM worker weights against a HuggingFace checkpoint.""" + return await self._execute_method( + "check_loaded_weights_equal", + kwargs={ + "model_path": model_path, + "trust_remote_code": trust_remote_code, + }, + ) + async def _update_delta_weights( self, weights,