Skip to content

Distributed improvement of Muon implementation #7808

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
Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Distributed improvement of Muon implementation #7808

wants to merge 1 commit into from
+182 −56

Conversation

tianshijing
Copy link
Contributor

What does this PR do?

The distributed training of Muon was carefully considered.

  1. ​​Distributed Training Support​​: Added gradient synchronization via reduce_scatter_tensor and parameter updates via all_gather_into_tensor for proper distributed training.
  2. ​​Performance Optimization​​: Implemented communication-computation overlap with asynchronous operations when enabled (overlap_comm=True).
    ​​3. Memory Efficiency​​: Only allocates communication buffers in distributed mode and uses gradient sharding to minimize memory usage.
    ​​4. Robustness​​: Enhanced error handling with assertions for 2D parameters and better None gradient management.
    ​​5. Backward Compatibility​​: Maintains original functionality for non-distributed cases while adding distributed capabilities.
    Fixes # (issue)

Before submitting

@hiyouga hiyouga mentioned this pull request May 13, 2025
Open
1 task
tpoisonooo
tpoisonooo reviewed May 13, 2025
View reviewed changes
src/llamafactory/train/trainer_utils.py
logger.info_rank0(
f"Using Muon optimizer with {len(muon_params)} Muon params and {len(adamw_params)} AdamW params."
f"Using Muon optimizer with {len(muon_params)} Muon params and {len(adamw_params)} AdamW params. "
Copy link
Contributor

Choose a reason for hiding this comment

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

看起来重复了。

@hiyouga hiyouga added the pending This problem is yet to be addressed label Jun 16, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
1 more reviewer

@tpoisonooo tpoisonooo tpoisonooo left review comments

Reviewers whose approvals may not affect merge requirements
Assignees
No one assigned
Labels
pending This problem is yet to be addressed
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@tianshijing @tpoisonooo @hiyouga