feat: support sparse HCCL weight sync with TP - #408
Conversation
Documentation build overview
48 files changed ·
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for Ascend Delta Weight Sync and Sparse Weight Sync strategies, allowing efficient weight synchronization across non-colocated rollout engines. It adds new classes (UpdateWeightFromDiskDelta and UpdateWeightFromSparseDistributed), delta encoding/decoding utilities, checksum verification, and integrates these strategies into the Megatron-Bridge actor and vLLM engine. Feedback on the changes highlights a critical issue in the PyTorch distributed communication logic within sparse_gather.py, where global ranks are incorrectly used instead of group-local ranks in dist.P2POp calls, which could lead to runtime errors during tensor parallel training.
| """ | ||
| rank = dist.get_rank(group) | ||
| world = dist.get_world_size(group) | ||
| destination = dist.get_global_rank(group, 0) |
There was a problem hiding this comment.
In PyTorch distributed, when a process group group is specified in dist.P2POp, the peer (or destination) rank must be the rank within that process group (i.e., the group rank), not the global rank. Passing the global rank will cause runtime errors or incorrect communication when group is not the world group (e.g., under tensor parallel training where tp_size > 1).
Since the destination is always the rank 0 of the process group, it should be set to 0 directly.
| destination = dist.get_global_rank(group, 0) | |
| destination = 0 |
| ) | ||
| index_list.append(index_buffer) | ||
| value_list.append(value_buffer) | ||
| peer = dist.get_global_rank(group, group_rank) |
There was a problem hiding this comment.
Similarly, the peer rank passed to dist.P2POp must be the group rank (group_rank) rather than the global rank when group is specified. Using the global rank here will lead to out-of-bounds errors or incorrect communication when the process group is not the world group.
| peer = dist.get_global_rank(group, group_rank) | |
| peer = group_rank |
1 participant
Summary
Validation
Notes
This PR is based on
validated-delta, so it currently includes those prerequisite commits in the diff. The matching vLLM-Ascend branch isfeature/sparse-hccl-weight-sync.