Repository navigation
feat: add delta-compressed collective refit - #2444
Merged
Merged
Conversation
|
Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds an optional delta-compressed weight transfer protocol for non-colocated vLLM collective refit, enabling the trainer source rank to send full weights or additive deltas (dense / sparse_indices / sparse_bitmask) and apply deltas additively through existing vLLM weight loaders.
Changes:
- Introduces a delta-aware packed weight transfer protocol (
full/delta/done) with sparse delta encodings and a trainer-sideDeltaCompressionTrackerbaseline. - Integrates the new transfer path into DTensor v1/v2 and Megatron policy workers via a shared
dispatch_packed_weight_transfer(...)helper. - Updates vLLM collective refit to optionally consume the new full/delta protocol and adds unit tests + example configs.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
tests/unit/utils/test_weight_transfer.py |
Adds unit coverage for delta tracker behavior, sparse transports, additive load context, and producer/consumer roundtrips. |
nemo_rl/utils/weight_transfer.py |
Implements delta tracking, sparse encodings, packed full/delta broadcast protocol, and additive load context. |
nemo_rl/utils/weight_transfer_types.py |
Defines shared literal types/constants for delta compression and transfer kinds. |
nemo_rl/utils/torch_dtypes.py |
Centralizes dtype string→torch.dtype mappings (canonical + aliases). |
nemo_rl/models/policy/workers/megatron_policy_worker.py |
Switches collective weight broadcast to the delta-aware dispatcher when enabled. |
nemo_rl/models/policy/workers/dtensor_policy_worker.py |
Switches collective weight broadcast to the delta-aware dispatcher when enabled. |
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py |
Switches collective weight broadcast to the delta-aware dispatcher when enabled. |
nemo_rl/models/generation/vllm/vllm_worker.py |
Determines whether to use delta transfer and forwards that flag to the vLLM worker extension. |
nemo_rl/models/generation/vllm/vllm_worker_async.py |
Forwards the delta-transfer enablement flag in the async prepare_refit_info path. |
nemo_rl/models/generation/vllm/vllm_backend.py |
Adds delta-aware collective consumer path and additive-delta loading through existing loaders. |
nemo_rl/models/generation/vllm/config.py |
Extends vLLM generation config typing with delta_compression settings. |
nemo_rl/models/automodel/setup.py |
Reuses canonical dtype mapping from torch_dtypes instead of duplicating it. |
examples/configs/grpo_math_1B.yaml |
Documents/introduces the new delta_compression config block (disabled by default). |
examples/configs/distillation_math.yaml |
Documents/introduces the new delta_compression config block (disabled by default). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Contributor
|
Awesome @HollowMan6 I found delta weight transfer has its own weight transfer function, which seems a duplicated one compared with the full weight transfer. It is out of the scope of this PR, but is there any block to have delt and full weight transfer shared the same communication function while have their independent protocol to pack, unpack the model weights? |
Contributor
Author
|
Thank you @ZhiyuLi-Nvidia for pointing this out, I just did some refactoring according to your suggestion, and it looks fine. |
github-actions
Bot
added
the
Documentation
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Contributor
Author
|
/ok to test cf6f490 |
terrykong
approved these changes
Jul 19, 2026
This was referenced Jul 20, 2026
This was referenced Aug 2, 2026
2 of 4 tasks
github-actions Bot
pushed a commit
to gitsocial-org/arxiv-ai-daily
that referenced
this pull request
Oct 7, 2026
… Trillion-Parameter Scale arXiv: 2610.08430 Authors: Songlin Jiang, Zhiyu Li, Terry Kong, Yu Yao, Youngeun Kwon, et al. Published: 2026-10-06 Agentic reinforcement learning (RL) disaggregates training from rollout, so each policy update must reach the rollout clusters before the next batch. Transferring a full 1T checkpoint for such weight synchronization (refit) takes 87.5 min between two AWS regions. Measurements of BF16 training show that about 1% of weights change their stored values per step. Recent systems exploit this sparsity but fall short on placement, exactness, or efficiency: they reimplement placement rules, assemble full tensors, rebuild values arithmetically, or use a cross-cluster collective, and none fully recovers from mid-refit failures. We present NeMo-DCR (Delta-Compressed Refit), which sends only changes yet is bit-exact: receivers obtain the same parameter and buffer bits as a dense refit. For placement, fixed affine mappings project changes from training shards into the checkpoint's canonical coordinates, residual conversion covers the other changes, and the serving runtime's native loader places all changes in receiver storage. For exactness, compressible XOR masks carry affine changes whose projection and loader preserve stored bits, and overwrites carry the others. Receivers apply both in place, retries overwrite partial writes, and a joint commit binds the policy to the baseline for the next delta. For efficiency, object storage or a relay tree streams payloads during delta construction, without a cross-cluster collective. Even at 3% and 5% change rates, NeMo-DCR refits of 30B-1T models are 12-40$\times$ faster than a transport-only full-checkpoint reference. A 1T relay-tree refit at 3% takes 150 s instead of 87.5 min, making refits practical for cross-cluster agentic RL at trillion-parameter scale. Links: - Abstract: https://arxiv.org/abs/2610.08430 - PDF: https://arxiv.org/pdf/2610.08430 Categories: cs.DC, cs.AI Note: The code is open-sourced in NVIDIA NeMo RL PR #2444 at NVIDIA-NeMo/RL#2444
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
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 optional delta-compressed weight transfer with s3 or zeromq backend for non-colocated vLLM collective refit.
Please see docs/design-docs/sparse-delta-refit.md for docs.
Issues
N/A
Usage
Before your PR is "Ready for review"
Pre checks:
Additional Information