feat(refit): add digest verification for vLLM IPC weight transfers - #3961
Open
NolenLiang wants to merge 6 commits into
Open
feat(refit): add digest verification for vLLM IPC weight transfers#3961NolenLiang wants to merge 6 commits into
NolenLiang wants to merge 6 commits into
Conversation
Adds opt-in verification that the bytes loaded by vLLM workers during a colocated CUDA-IPC refit are exactly the bytes the policy workers sent, configured via policy.generation.refit_cfg.verify.mode (off|log|enforce, default off with zero overhead). Each side hashes its view of the transfer with a polynomial rolling hash over the raw bytes computed in int64 with wraparound (mod 2^64). Integer modular arithmetic is associative and commutative, so the digest is bit-identical regardless of reduction order, chunking, or device -- unlike a floating-point checksum, which would be subject to the very nondeterminism this check is meant to catch. Digests stay on-device during streaming (no per-tensor sync) and the receiver returns them with the final COMPLETE ACK as a pyobj, so worker return values and the synchronizer's success handling are unchanged. The sender compares and warns (log) or raises with the mismatched parameter names (enforce). Because the receiver slices the staged buffer from its own prepare_refit_info metadata, digest equality also implicitly validates that the refit metadata still matches what the exporter streams. Scope: colocated IPC/ZMQ transport only; the NCCL collective path is a follow-up. Sparse-delta refit already has its own payload checksums. Signed-off-by: Nolen Liang <nliang@nvidia.com>
- Replace the linear polynomial digest with position-salted SplitMix64 mixing: the linear form cancels deterministically when two int64 lanes both flip their top bit, since 2^63 * (R^a + R^b) is 0 mod 2^64 for odd R. - Fold dtype and shape into the digest seed so equal-size metadata drift is detected, not just byte corruption. - Group digest materialization by device to support mixed CPU/CUDA parameter streams (e.g. scalar KV scales). - Reject verify_digests on generation backends without support (Dynamo, TensorRT-LLM, SGLang) and fail at synchronizer construction for non-colocated transports instead of silently skipping verification. - Drop repeated "off" parameter defaults along the call chain; the mode is resolved once from the config model and passed explicitly. - Add digest.py to the Pyrefly check list and extend the signature-contract tests to the IPC update path of every generation backend. Signed-off-by: Nolen Liang <nliang@nvidia.com>
- Widen the digest to two independently-parameterized 64-bit channels with positions injected through the nonlinear mixer. A linear position salt is absorbable: corrupted lanes can soak up the salt difference between two positions and permute the salted values, which a commutative sum cannot see -- and for any single bijective per-lane transform such a permutation is constructible in closed form. With two channels the same corrupted lanes must simultaneously preserve both nonlinearly coupled sums, for which no closed-form construction exists. Adds the two-lane construction as a regression test; digests serialize as 32-hex strings. - Validate refit verification support at VllmGeneration construction: collective/sparse/PPO/distillation paths construct synchronizers directly (bypassing the weight-sync factory), which previously let verify.mode=enforce be silently ignored on unsupported topologies. - Make VllmRefitVerifyConfig extra="forbid" and stop coercing explicit invalid values (verify: false, misspelled keys) into the "off" default; only absent fields default. Adds validation tests. - Sync the verify block into the PPO and distillation exemplar YAMLs and their reference configs. Signed-off-by: Nolen Liang <nliang@nvidia.com>
…rify config - Replace commutative-sum digest channels with an ordered binary tree fold. k commutative channels impose only k multiset constraints while n lanes provide n degrees of freedom, so lane permutations satisfying every channel simultaneously exist and were constructed for the one- and two-channel variants; in a tree a lane's position is bound to its path, leaving no permutation freedom. Both children pass through the mixer (distinct tweaks) before the linear combination -- a bare linear combination leaves 2^63 * (odd + odd) == 0 mod 2^64 open, which the paired-top-bit regression caught. Each level is elementwise over disjoint pairs, so the fold stays parallel and device-independent; chunk roots are order-chained. The chunk size is now an algorithm constant. Adds the six-lane dual-channel collision pair (int64 and bf16-reinterpreted) as regressions. - Only a truly absent verify field defaults to off: explicit "verify: null" now fails in Pydantic, and unknown top-level refit_cfg keys (e.g. "verfiy") are rejected at VllmGeneration construction, scoped to allow custom checkpoint-engine selector keys. Signed-off-by: Nolen Liang <nliang@nvidia.com>
…atter - resolve_refit_verify_config now feeds any explicit non-null refit_cfg to Pydantic as-is, so non-mapping values ([] / "" / false) fail loudly instead of silently defaulting verification to off; only a truly absent refit_cfg defaults. - The unknown-key check now also covers VllmRefitConfig.model_extra: paths that call normalize_vllm_refit_config before constructing VllmGeneration (sparse/NIXL/plugin transports) validate and write the model back, so a typo like "verfiy" previously slipped into model_extra and bypassed the dict-based check. Adds regressions in the real normalize -> construct order, including the plugin-selector allowance. - normalize_vllm_refit_config no longer coerces falsy refit_cfg values into an empty config. - Reformat with the repository-pinned ruff (0.9.9); an unpinned newer ruff had reformatted one test file differently. Signed-off-by: Nolen Liang <nliang@nvidia.com>
github-actions
Bot
added
the
Documentation
Contributor
Author
|
/ok to test 18adf3d |
Signed-off-by: Nolen Liang <nliang@nvidia.com>
Contributor
Author
|
/ok to test 3740298 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Adds opt-in sender/receiver digest verification for colocated vLLM CUDA-IPC weight refits.
Policy and vLLM workers independently digest each transferred parameter's bytes, dtype, and shape. The receiver returns its digests with the final acknowledgment, and the sender either logs mismatches or fails the refit. The digest is a deterministic, ordered two-channel tree fold that runs on the tensor's device and remains stable across CPU/CUDA devices and chunk boundaries.
This PR also validates verification configuration at setup, rejects unsupported transports and generation backends instead of silently ignoring the setting, preserves the existing IPC wire format when verification is disabled, and updates the refit guide and example configurations.
Issues
N/A — this is a proactive feature and is not linked to an existing issue.
Usage
"off"is the default: no digest computation and no IPC wire-format change."log"reports mismatched parameters and continues."enforce"raises on a mismatch.Verification is currently supported only for colocated vLLM CUDA-IPC refits (
colocated.enabled=trueandrefit_transport=null). Enabling it for collective, NCCL reshard, sparse/NIXL/plugin transports, or another generation backend fails during setup.Before your PR is "Ready for review"
Pre checks:
Focused unit tests were run locally; GPU functional coverage is left to CI.
Additional Information
Validation performed:
git diff --checkpasses.Risk and scope: