Repository navigation
Refit speedup for non-colocated (NCCL) #817
Description
Currently the initial implementation #613 is gathering weights on rank 0 train worker then broadcast to all vllm workers, which is the similar to https://github.com/vllm-project/vllm/blob/v0.9.0/examples/offline_inference/rlhf.py.
Two paths that can currently be thought of:
- from @SahilJain314: not broadcasted everything from rank 0.
broadcasting from many nodes with the full weights is much more performant (less traffic on a single node's network). Or better yet, this lets us avoid doing a full gather on rank 0.
- cosmos-rl solution
Activity
guyueh1 commentedon Aug 26, 2025
@youngeunkwon0405 fyi
euronymous-aithal commentedon Sep 17, 2025
@youngeunkwon0405 what is the update on this ?
youngeunkwon0405 commentedon Sep 17, 2025
@euronymous-aithal I am currently working on enhancing performance-related visibility (e.g., nsys-related, more perf metrics in general, etc.). Probably after that, I will work on refit performance unless I find more severe perf issues.
youngeunkwon0405 commentedon Oct 3, 2025
#1264 Might improve the non-colocated perf a lot. Need to test the perf improvement in larger models (e.g., DSV3 or LLaMa3 405B).
youngeunkwon0405 commentedon Oct 8, 2025
Hi @terrykong, there is another optimization (#1313) for non-colocated refit.
#1264 pushes the upper bound broadcast bandwidth we could achieve.
#1313 do the tensor packing to achieve higher communication bandwidth optimization. And also do the batched weight update to amortize overheads. This optimization is crucial for fine-grain MoE models.
qiaochuz-nv commentedon May 5, 2026
QA RCCA Analysis
RCCA Status: Done
RCCA Category: covered
Test Action: none
Fix PR: #613, #1264, #1313, #1379
Root Cause: fix PR includes upstream test coverage.
Trigger Config: Refit speedup for non-colocated (NCCL)
Coverage Assessment:
Conclusion: RCCA classification completed for this closed NeMo-RL GitHub issue.