Skip to content

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

Closed
le-czs wants to merge 1 commit into
verl-project:mainfrom
le-czs:feat-checkpoint-weight-sync-check
Closed

le-czs wants to merge 1 commit into
verl-project:mainfrom
le-czs:feat-checkpoint-weight-sync-check

Conversation

@le-czs

@le-czs le-czs commented Jun 9, 2026 •

Copy link
Copy Markdown
Contributor

What does this PR do?

Addresses #6414.

Adds a checkpoint engine debug check that verifies whether vLLM loaded the same weights sent by the trainer during the initial trainer-to-rollout sync. This PR is limited to vLLM with tensor and data parallel sizes 1.

The check compares the source Hugging Face checkpoint state dict with the rollout backend's loaded state dict. It handles vLLM fused qkv_proj and gate_up_proj weights, tied lm_head.weight, and reports missing, unexpected, and mismatched keys. Value comparison is strict after casting the source checkpoint tensor to the loaded rollout tensor dtype/device.

This PR intentionally does not include check_weight_sync_only or fully async shutdown cleanup. Follow-up stacked draft PR: le-czs#1. It will be rebased and opened against verl-project/verl after this PR lands.

This is not duplicating an existing PR. I checked these searches before opening the PR:

Related open PRs are about MTP sleep, FP8 transfer, generic remote backends, or fully async transfer queue. I did not find an open PR that adds a vLLM loaded-weight comparison for #6414.

AI assistance was used for implementation and PR preparation. I reviewed the changed lines before submission.

Checklist Before Starting

  • I have searched for similar PRs to avoid duplicate work. Query links are included above.
  • I have formatted the PR title as [vllm, checkpoint] feat: add rollout weight sync debug check.

Test

  • python3 -m pytest tests/checkpoint_engine/test_weight_sync_on_cpu.py -q
    • Result: 11 passed, 1 warning
  • scripts/generate_trainer_config.sh
    • Result: All good
  • git diff --check
    • Result: passed
  • pre-commit run --all-files --show-diff-on-failure --color=always
    • Result: passed

API and Usage Example

Example overrides:

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

check_weight_sync=True runs the loaded-weight comparison after the initial trainer-to-rollout sync.

Design & Code Changes

  • Adds checkpoint engine config field check_weight_sync.
  • Adds verl.checkpoint_engine.weight_sync for strict state dict comparison and vLLM fused-weight mapping.
  • Adds a vLLM server RPC path that returns loaded-weight comparison results from the rollout backend.
  • Restricts the current vLLM check to TP=1 and DP=1; DP>1 is rejected because the current collective_rpc path only covers one DP shard.
  • Runs the check after the initial weight sync, because the comparison source is the original HF checkpoint.
  • Adds CPU tests for manager check selection, strict mismatch reporting, vLLM fused comparison, unsupported parallelism rejection, and missing/unexpected keys.
  • Adds docs in docs/advance/checkpoint.rst and verl/checkpoint_engine/README.md.

Checklist Before Submitting

  • I have read the Contribute Guide and followed the PR process.
  • I have run pre-commit checks: pre-commit run --all-files --show-diff-on-failure --color=always.
  • I have added or updated documentation.
  • I have added tests or explained why tests are not feasible.
  • I have requested CI in the required channel, if needed.
  • If this PR changes recipes, I have updated the recipe submodule.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a trainer-to-rollout weight synchronization check to verify that backend-loaded rollout weights match the source HuggingFace checkpoint, adding configuration options, documentation, unit tests, and early-exit support. A critical issue was identified in verl/checkpoint_engine/base.py where rollout.check_loaded_weights_equal is called without ray.get, which would cause verification failures on the workers to be silently ignored. Wrapping this call in ray.get is recommended to ensure errors are properly propagated.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread verl/checkpoint_engine/base.py Resolved
@le-czs

le-czs commented Jun 9, 2026

Copy link
Copy Markdown
Contributor Author

Verification Result

image

@Luosuu Luosuu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Left two inline comments on correctness issues in the new debug path.

Comment thread verl/workers/rollout/llm_server.py Resolved

@Luosuu Luosuu left a comment •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this PR touches too much for the weight sync debug check and actually this PR can be broken down into two parts:

  1. a minimal version that check weight sync for vLLM (DP/TP=1)
  2. Another PR for check_weight_sync_only and fully async clean shutdown

please try to make each PR minimal and easy to review. Thank you!

@le-czs

le-czs commented Jun 15, 2026

Copy link
Copy Markdown
Contributor Author

I think this PR touches too much for the weight sync debug check and actually this PR can be broken down into two parts:

  1. a minimal version that check weight sync for vLLM (DP/TP=1)
  2. Another PR for check_weight_sync_only and fully async clean shutdown

please try to make each PR minimal and easy to review. Thank you!

I reduced #6666 to the minimal vLLM TP=1/DP=1 weight-sync check.
I also prepared the follow-up for check_weight_sync_only and fully async cleanup as a stacked draft PR: le-czs#1. Once #6666 lands, I’ll rebase it and open it against verl-project/verl.

@le-czs

le-czs commented Jun 16, 2026

Copy link
Copy Markdown
Contributor Author

I think this PR touches too much for the weight sync debug check and actually this PR can be broken down into two parts:

  1. a minimal version that check weight sync for vLLM (DP/TP=1)
  2. Another PR for check_weight_sync_only and fully async clean shutdown

please try to make each PR minimal and easy to review. Thank you!

Thanks. I’ve reduced this PR to the minimal vLLM TP=1/DP=1 check and moved the follow-up scope out.

Could you re-review when you have time? The previous change request should be addressed now.

@le-czs
le-czs force-pushed the feat-checkpoint-weight-sync-check branch from 5298bc9 to c370a58 Compare June 18, 2026 02:02
@le-czs

le-czs commented Jun 22, 2026

Copy link
Copy Markdown
Contributor Author

Hi @Luosuu, I rebased the branch and the previous inline threads are now resolved/outdated. The PR is still blocked by the old CHANGES_REQUESTED review.

Could you re-review when you have time? If the scope looks OK now, I’ll request CI.

Rebase the vLLM TP=1/DP=1 loaded-weight check onto current main.

Co-authored-by: Grok 4.7 <noreply@x.ai>
Co-authored-by: Cursor <cursoragent@cursor.com>
@le-czs
le-czs force-pushed the feat-checkpoint-weight-sync-check branch from c370a58 to e903d5b Compare October 21, 2026 07:45
@le-czs

le-czs commented Oct 21, 2026

Copy link
Copy Markdown
Contributor Author

Rebased onto current main (cd58f550) and resolved the conflicts. Scope is unchanged: the vLLM loaded-weight check for tensor_parallel_size=1 and data_parallel_size=1.

@Luosuu the June 12 CHANGES_REQUESTED review is still blocking this PR. The inline threads were already addressed, and check_weight_sync_only / fully async shutdown stay out of this PR. Could you re-review?

@HollowMan6 this review was reassigned to you on Sept 21. Please take a look when you have time.

I will request CI in the Feishu group.

@coderabbitai

coderabbitai Bot commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration
  • Configuration used: Organization UI
  • Review profile: CHILL
  • Plan: Advanced
  • Run ID: 1771ebd5-d313-49ae-afb6-9ebd4c453783
📥 Commits

Reviewing files that changed from the base of the PR and between cd58f55 and e903d5b.

📒 Files selected for processing: 13
  • docs/advance/checkpoint.rst
  • tests/checkpoint_engine/test_weight_sync_on_cpu.py
  • verl/checkpoint_engine/README.md
  • verl/checkpoint_engine/base.py
  • verl/checkpoint_engine/weight_sync.py
  • verl/trainer/config/_generated_ppo_megatron_trainer.yaml
  • verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml
  • verl/trainer/config/_generated_ppo_trainer.yaml
  • verl/trainer/config/_generated_ppo_veomni_trainer.yaml
  • verl/trainer/config/rollout/rollout.yaml
  • verl/workers/config/rollout.py
  • verl/workers/rollout/vllm_rollout/utils.py
  • verl/workers/rollout/vllm_rollout/vllm_rollout.py

Included review availability: This review used your included allowance. Your plan provides up to 10 included reviews per hour; 9 remain after this review.

📝 Summary

Summary by CodeRabbit

  • New Features

    • Added an optional checkpoint weight synchronization check to validate that rollout weights match the source model checkpoint.
    • Reports missing, unexpected, and mismatched weights.
    • Supports strict validation for vLLM with tensor and data parallelism set to 1.
    • The check runs after initial weight loading and is disabled by default.
  • Documentation

    • Added configuration guidance, supported backend limitations, and usage examples.
  • Tests

    • Added coverage for matching weights, mismatches, missing or unexpected keys, fused weights, and unsupported parallel configurations.

The change adds an opt-in check that compares rollout-loaded weights with a source Hugging Face checkpoint after initial synchronization. The comparison supports vLLM with tensor and data parallel sizes of 1, and reports missing, unexpected, and mismatched weights.

Changes: Trainer-to-rollout weight sync

Layer / File(s) Summary
Configure and compare weights
verl/workers/config/rollout.py, verl/trainer/config/rollout/rollout.yaml, verl/trainer/config/_generated_ppo_*_trainer.yaml, verl/checkpoint_engine/weight_sync.py, tests/checkpoint_engine/test_weight_sync_on_cpu.py
Adds the disabled-by-default configuration option and comparison logic for direct and supported fused vLLM tensors. Tests cover supported parallelism, matching tensors, mismatches, and missing or unexpected keys.
Inspect weights loaded by vLLM
verl/workers/rollout/vllm_rollout/utils.py, verl/workers/rollout/vllm_rollout/vllm_rollout.py
Adds a vLLM helper that loads reference weights on CPU and compares them with loaded model weights. The server adapter forwards the checkpoint path and trust setting to that helper.
Run the check after initial synchronization
verl/checkpoint_engine/base.py, tests/checkpoint_engine/test_weight_sync_on_cpu.py, verl/checkpoint_engine/README.md, docs/advance/checkpoint.rst
The checkpoint manager runs the check only when enabled and the load count is None or 0. Documentation describes the timing, configuration, comparison results, and supported constraints.

Sequence Diagram(s)

sequenceDiagram
  participant CheckpointEngineManager
  participant CheckpointEngineWorker
  participant ServerAdapter
  participant vLLMColocateWorkerExtension
  participant HuggingFaceCheckpoint
  CheckpointEngineManager->>CheckpointEngineWorker: Request initial loaded-weight check
  CheckpointEngineWorker->>ServerAdapter: Forward model path and trust setting
  ServerAdapter->>vLLMColocateWorkerExtension: Execute check_loaded_weights_equal
  vLLMColocateWorkerExtension->>HuggingFaceCheckpoint: Load reference weights on CPU
  HuggingFaceCheckpoint-->>vLLMColocateWorkerExtension: Return reference state dict
  vLLMColocateWorkerExtension-->>ServerAdapter: Return comparison result
  ServerAdapter-->>CheckpointEngineWorker: Return comparison result
  CheckpointEngineWorker-->>CheckpointEngineManager: Return check summary
Loading

Priority: ⬇️ Low

Merge Risk: ⚪ Minimal · up to e903d

No merge-blocking issue was established for the opt-in weight-sync check; it is ready for normal merge checks.

Security architecture risk: 🔵 Low · up to e903d

The check is disabled by default and normally reuses the configured model source and remote-code trust setting. Failed validation prevents generation from resuming. No new attacker-accessible boundary was established, but control-plane exposure and complete failure recovery remain unverified.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

  • inferred — The demonstrated normal execution scope is the rollout workers selected by the checkpoint manager. The added operation loads model files and may execute model code when remote-code trust is enabled. Runtime tenant isolation, network exposure, and credential scope were not established, so a wider attacker-accessible blast radius cannot be determined.

Trust Boundaries and Controls

  • observed — The normal manager call supplies no path or trust overrides: the worker uses the configured model path and trust_remote_code value. vLLM startup already consumes that trust decision. The new worker API nevertheless accepts explicit overrides, making control-plane authorization relevant to arbitrary callers.
  • observed — The repository wrapper forwards collective RPCs to the engine, while the HTTP application is constructed by dependency-provided build_app. The inspected repository source does not resolve whether external HTTP callers can reach the new worker method or what deployment authorization protects the Ray control plane.

Resilience and Maintainability Implications

  • inferred — The validation failure path withholds generation rather than serving weights that failed comparison. The principal trainer startup caller does not catch this failure locally. Worker termination, resource reclamation, retry serialization, and recovery after interrupted validation remain unverified; absence of automatic resumption is not itself evidence of an insecure state.

Hardening Proposals

  • proposed — Consider preflighting supported configurations and defining a bounded validation operation with explicit failed-state cleanup and retry identity. Recovery should preserve the fail-closed gate rather than unconditionally resuming generation after a mismatch.
Pre-merge checks: Passed 4 · Failed 1

Failed Failed checks: 1 warning

Check name Status Explanation Resolution
Docstring Coverage Warning Docstring coverage is 56.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 25 functions across 6 files. (7 skipped: … Write docstrings for the functions missing them to satisfy the coverage threshold.
Passed Passed checks: 4 passed
Check name Status Explanation
Title check Passed The title clearly identifies the vLLM checkpoint weight-sync debug check, which is the main change in the pull request.
Description check Passed The description directly explains the vLLM weight-sync check, its scope, implementation, tests, documentation, and excluded follow-up work.
Linked Issues check Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check Passed Check skipped because no linked issues were found for this pull request.
Full details: Docstring Coverage

Explanation

Docstring coverage is 56.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 25 functions across 6 files. (7 skipped: 7 unsupported.)

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Comment @coderabbitai help to get the list of available commands.

@wuxibin89

Copy link
Copy Markdown
Collaborator

You can debug rollout weight sync with scripts/rollout_logprob_diff.py, please check https://verl.readthedocs.io/en/latest/low_precision/qat.html

@le-czs le-czs closed this Oct 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants