Skip to content

feat(fsdp): patch Megatron torch_FSDP2 with Primus implementation #35

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 3 commits into from
Apr 24, 2025
Merged

feat(fsdp): patch Megatron torch_FSDP2 with Primus implementation #35

merged 3 commits into from
Apr 24, 2025

Conversation

@Xiaoming-AMD
Copy link
Collaborator

@Xiaoming-AMD Xiaoming-AMD commented Apr 24, 2025

Feat(fsdp): patch Megatron torch_FSDP2 with Primus implementation

  • Monkey-patched TorchTorchFullyShardedDataParallel in Megatron with custom PrimusTorchFullyShardedDataParallel
  • Ensured robust patching with try/except and clearer logging
  • Wrapped PrimusTorchFullyShardedDataParallel to use safe default for sub_modules_to_wrap

wenxie-amd
wenxie-amd approved these changes Apr 24, 2025
View reviewed changes
Copy link
Contributor

@wenxie-amd wenxie-amd left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@wenxie-amd wenxie-amd merged commit 734caea into main Apr 24, 2025
2 checks passed
@Xiaoming-AMD Xiaoming-AMD deleted the dev/xiaoming/fix_torch-fsdp_init branch April 28, 2025 05:27
@Xiaoming-AMD Xiaoming-AMD changed the title [Feat(fsdp)] patch Megatron torch_FSDP2 with Primus implementation feature(fsdp): patch Megatron torch_FSDP2 with Primus implementation Jun 4, 2025
@Xiaoming-AMD Xiaoming-AMD changed the title feature(fsdp): patch Megatron torch_FSDP2 with Primus implementation feat(fsdp): patch Megatron torch_FSDP2 with Primus implementation Jun 25, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@wenxie-amd wenxie-amd wenxie-amd approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@Xiaoming-AMD @wenxie-amd