Skip to content

Updating support of Megatron-LM - #3842

Merged
SunMarc merged 31 commits into
huggingface:mainfrom
pengdurice:peng-fix-using-megatron-oss-v1
Dec 3, 2025
Merged

Updating support of Megatron-LM#3842
SunMarc merged 31 commits into
huggingface:mainfrom
pengdurice:peng-fix-using-megatron-oss-v1

Conversation

@pengdurice

@pengdurice pengdurice commented Nov 21, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

What does this PR do?

This PR is to make sure that accelerate continues support using Megatron-LM as backend for training large scale LLMs.
In details:

  1. FP8 quantization of the model is by-passed if Megatron is used since the model will be re-initialized later on.
  2. Added a set of useful parameters necessary for enabling Megatron-LM: e.g. gradient checkpointing, cpu optimizer offloading.
  3. Made changes on calling Megatron's functions since the func signatures have been changed.
  4. Removed megatron_generate since it is no longer supported.
  5. Updated the user guide doc.
  6. The change has been tested with fine-tuning GLM4.6 model using 64 H200 GPUs with 70K sequence length.

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a Github issue or the forum? Please add a link
    to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Peng-Writer and others added 29 commits October 13, 2025 16:05
@Peng-Writer
now the training of gpt model runs e2e
080fd5d
@Peng-Writer
add qwen_moe config
05950aa
@Peng-Writer
clean up the code
58981c8
@Peng-Writer
clean up the code
25d3566
@Peng-Writer
default build tokenizer to False
5e29abe
@Peng-Writer
add qwen and glm models
05051e9
@Peng-Writer
add deepseek v3 config
4fdfdec
@Peng-Writer
glm4.5 air running, with fp8, loading model from saved pytorch dist
40c111e
@Peng-Writer
glm4.5 (355b-a32b) running, with fp8, loading model from saved pytorc…
51b8b31 
…h dist
@Peng-Writer
save for now
015d954
@Peng-Writer
11/13/12pm, 5092, error comes back again: kernel didn't clean all nans
4eb8640
@Peng-Writer
glm4.5air train with 70k seq length OK by hacking the transformer eng…
72d1190 
…inee's code
@pengdurice
remove logging
97a36ed
@pengdurice
remove logging
3492e2a
@pengdurice
save before updating to a new branch
6dbe2c4
@pengdurice
fix an issue when newest Megatron-LM is used
22494f1
@pengdurice
add the changes and get ready to revert back
df9c8d8
@pengdurice
removing the assignments in the model_provider_func func of megatron_…
a7554fd 
…lm.py and found out that it is self.expert_tensor_parallel_size that causes the OOM isue
@pengdurice
move last pp # of layers to the config
029b2c1
@pengdurice
fix last layer non error
46110c3
@pengdurice
allow cpu offload optimizer and NOT saving optimizer states
eda478c
@pengdurice
add h2d and d2h overlapp for offload optimizer to cpu
2395323
@pengdurice
clean up the code for submitting the PR
09983b2
@pengdurice
Merge branch 'main' into peng-fix-using-megatron-oss-v1
3fab2f8
@pengdurice
move back to in
2c28681
@pengdurice
other small changes
7404398
@pengdurice
other small changes
f66de72
@pengdurice
update readme doc
560146e
@pengdurice
update readme doc
8b78137
@pengdurice

pengdurice commented Nov 21, 2025

Copy link
Copy Markdown
Contributor Author

@SunMarc This is Peng. I am working on using accelerate and megatron-LM for fine tuning GLM4.6 models. This PR includes some updates of accelerate to better support new features as well as configurations from Megatron. Would you please help take a look? Thank you!

@SunMarc

SunMarc commented Nov 24, 2025

Copy link
Copy Markdown
Member

Of course, I will have a look this week !

@SunMarc
SunMarc self-requested a review November 24, 2025 15:34
@pengdurice

pengdurice commented Nov 24, 2025

Copy link
Copy Markdown
Contributor Author

Of course, I will have a look this week !

Thank you! Take your time and have a nice holiday ahead if you are in the US!

SunMarc
SunMarc approved these changes Dec 2, 2025
View reviewed changes

@SunMarc SunMarc left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Thanks a lot, we don't maintain that much megatron as it is quite complex for users but happy to have this PR !

Comment thread src/accelerate/utils/dataclasses.py Outdated
"attention_dropout": self.attention_dropout,
"hidden_dropout": self.hidden_dropout,
"attention_softmax_in_fp32": self.attention_softmax_in_fp32,
# "expert_tensor_parallel_size": self.expert_tensor_parallel_size,

@SunMarc SunMarc Dec 2, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

don't need to pass it ?

@pengdurice pengdurice Dec 2, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I don't quite find this config useful, so left it commented out. Let me uncomment it. thanks for the comment!

@pengdurice pengdurice Dec 2, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

@SunMarc , fixed it with a new commit. Would you please take another look? Thanks!

@SunMarc

SunMarc commented Dec 3, 2025

Copy link
Copy Markdown
Member

@bot /style

@github-actions

github-actions Bot commented Dec 3, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Style bot fixed some files and pushed the changes.

@HuggingFaceDocBuilderDev

HuggingFaceDocBuilderDev commented Dec 3, 2025

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@SunMarc
SunMarc merged commit 7e38469 into huggingface:main Dec 3, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@SunMarc SunMarc SunMarc 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.

4 participants

@pengdurice @SunMarc @HuggingFaceDocBuilderDev @Peng-Writer