-
Notifications
You must be signed in to change notification settings - Fork 2.4k
[model] feat: support parameter generator for model engine #3529
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.
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
[model] feat: support parameter generator for model engine #3529
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code Review
This pull request introduces a parameter generator for model engines and adds corresponding tests. The implementation for FSDP and Megatron engines has been added. My review focuses on correctness and consistency in these new engine implementations.
I've identified a critical issue in the FSDP engine implementation where state_dict() is called without the necessary context, which could lead to incorrect behavior with sharded models. Additionally, there's a high-severity issue in the Megatron engine implementation regarding inconsistent parameter offloading logic compared to the FSDP engine, which could cause GPU memory issues. The refactoring in the actor and critic roles looks good.
What does this PR do?
Checklist Before Starting
[{modules}] {type}: {description}(This will be checked by the CI){modules}includefsdp,megatron,sglang,vllm,rollout,trainer,ci,training_utils,recipe,hardware,deployment,ray,worker,single_controller,misc,perf,model,algo,env,tool,ckpt,doc,data,like[megatron, fsdp, doc]{type}is infeat,fix,refactor,chore,test[BREAKING]to the beginning of the title.[BREAKING][fsdp, megatron] feat: dynamic batchingTest
API and Usage Example
# Add code snippet or script demonstrating how to use thisDesign & Code Changes
Checklist Before Submitting
Important
Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.
pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=alwaysci-requestchannel in theverlSlack workspace. (If not accessible, please try the Feishu group (飞书群).)