-
Notifications
You must be signed in to change notification settings - Fork 2.9k
[megatron] chore: clean legacy code path part 1, make engine use mbridge by default #4528
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
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 is the first part of a series to clean up legacy code paths in Megatron and make mbridge the default. The changes enforce the use of mbridge in the Megatron engine and remove the old code paths that did not use it. The logic for selecting forward functions has also been refactored.
My review found two critical issues in verl/models/mcore/registry.py where checking for Vision Language Model (VLM) architectures is implemented incorrectly. The current implementation will fail to identify VLMs, causing incorrect forward functions to be used. I have provided suggestions to fix this.
| Get the forward function for given model architecture. | ||
| """ | ||
| assert len(hf_config.architectures) == 1, "Only one architecture is supported for now" | ||
| if hf_config.architectures[0] in SupportedVLM: |
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.
The check hf_config.architectures[0] in SupportedVLM is incorrect. The in operator on an Enum checks for member identity, not the string value of the members. Since hf_config.architectures[0] is a string, this condition will always evaluate to False, leading to the wrong forward function being selected for Vision Language Models (VLMs).
| if hf_config.architectures[0] in SupportedVLM: | |
| if hf_config.architectures[0] in {item.value for item in SupportedVLM}: |
| Get the forward function for given model architecture. | ||
| """ | ||
| assert len(hf_config.architectures) == 1, "Only one architecture is supported for now" | ||
| if hf_config.architectures[0] in SupportedVLM: |
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.
Similar to the issue in get_mcore_forward_fn, the check hf_config.architectures[0] in SupportedVLM is incorrect. The in operator on an Enum checks for member identity, not the string value of the members. This will always evaluate to False, causing the wrong fused forward function to be selected for Vision Language Models (VLMs).
| if hf_config.architectures[0] in SupportedVLM: | |
| if hf_config.architectures[0] in {item.value for item in SupportedVLM}: |
| engine.virtual_pipeline_model_parallel_size=${VPP_SIZE} \ | ||
| engine.context_parallel_size=${CP_SIZE} \ | ||
| engine.use_mbridge=False" | ||
| engine.use_mbridge=True" |
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.
We should also set use_mbridge to True? https://github.com/volcengine/verl/blob/main/verl/trainer/config/engine/megatron.yaml#L77-L78
What does this PR do?
this is one of a series PRs to clean the legacy megatron code path and make bridge the default path for megatron. #4496
This PR make megatron engine use only support bridge, and mark the legacy code in
mcore/registry.pyChecklist 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 (飞书群).)