[megatron] improving weight syncing - bucketed param gather + cuda ipc flattening #487
Conversation
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
erictang000
changed the title
…into weight_syncing
SumanthRH
requested changes
Nov 21, 2025
Comment on lines
34
to
36
| sizes: List[int] | ||
| extras: Optional[List[Dict[str, Any]]] | ||
| packed: bool |
Member
There was a problem hiding this comment.
The current PR will break compatibility for FSDP and DeepSpeed since they don't send these arguments.
You should probably update the "sizes" to be a NotRequired argument. And you should update Fsdp and deepspeed worker files to send packed=False.
Collaborator
Author
There was a problem hiding this comment.
good catch, updated and tested both
| request["names"], | ||
| request["dtypes"], | ||
| request["shapes"], | ||
| request["sizes"], |
Member
There was a problem hiding this comment.
This should be an optional entry. Should use safer .get
SumanthRH
reviewed
Nov 21, 2025
| names: List[str] | ||
| dtypes: List[str] | ||
| shapes: List[List[int]] | ||
| sizes: Optional[List[int]] |
Member
There was a problem hiding this comment.
Nit: this is NotRequired . Optional means that hte field exists but is None
from typing import NotRequired
SumanthRH
reviewed
Nov 21, 2025
| shapes: List[List[int]] | ||
| sizes: Optional[List[int]] | ||
| extras: Optional[List[Dict[str, Any]]] | ||
| packed: Optional[bool] |
SumanthRH
reviewed
Nov 21, 2025
|
|
||
| offset = 0 | ||
| for name, shape, size in zip(names, shapes, sizes): | ||
| weight_list.append((name, packed_tensor[offset : offset + size].view(*shape))) |
SumanthRH
approved these changes
Nov 21, 2025
Member
SumanthRH
left a comment
SumanthRH
left a comment
There was a problem hiding this comment.
Have left some minor comments. Please take a look
li-boxuan
pushed a commit
to li-boxuan/SkyRL
that referenced
this pull request
Nov 23, 2025
…c flattening (NovaSky-AI#487) Blocked on NovaSky-AI#453. Post Megatron-Bridge migration, we can now do a bucketed gathering of parameters, with precomputed size metadata, instead of iterating through parameters one by one. We then flatten the bucket into a single tensor, send over the metadata, and recover the weights/shapes on the inference engine side. For Qwen3-30B-A3B, tp=2, ep=8, etp=1, 8xh100: Before (48s): <img width="739" height="14" alt="image" src="https://github.com/user-attachments/assets/d90a2091-f36b-445e-898d-81c5019ea8a4" /> After bucketing gather (45s): <img width="754" height="24" alt="image" src="https://github.com/user-attachments/assets/5c0c725f-aafc-4989-9f77-c8a06f95e942" /> After flattening cuda ipc (and removing redundant torch.device() calls) - 6s: <img width="711" height="19" alt="image" src="https://github.com/user-attachments/assets/11a0be18-08cb-460a-9c3c-06ab06745b45" />
2 tasks
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Blocked on #453.
Post Megatron-Bridge migration, we can now do a bucketed gathering of parameters, with precomputed size metadata, instead of iterating through parameters one by one. We then flatten the bucket into a single tensor, send over the metadata, and recover the weights/shapes on the inference engine side.
For Qwen3-30B-A3B, tp=2, ep=8, etp=1, 8xh100:
Before (48s):
After bucketing gather (45s):
After flattening cuda ipc (and removing redundant torch.device() calls) - 6s:
