[megatron] feat: support discrete profiling for mindspeed #4271
Conversation
tardis-key
requested review from
FightingZhen,
ISEEKYAN and
vermouth1992
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for discrete profiling for MindSpeed on Ascend devices by adding a role parameter to @DistProfiler.annotate decorators in verl/workers/megatron_workers.py. This allows for more granular profiling data collection. The documentation has also been updated to reflect this new capability.
My main feedback is that the changes for discrete profiling seem incomplete. While the role parameter has been added to methods in ActorRolloutRefWorker, it's missing from similar annotations in CriticWorker and RewardModelWorker within the same file. Applying this change consistently would make the feature complete and more robust.
tardis-key
changed the title
| This is a tutorial for data collection using the GRPO or DAPO algorithm | ||
| based on FSDP on Ascend devices. | ||
| This is a tutorial for performance data collection using the GRPO or DAPO algorithm | ||
| based on FSDP on Ascend devices, which is also applicable on MindSpeed(Megatron) |
There was a problem hiding this comment.
Just describe like based on FSDP or MindSpeed (Megatron) on Ascend devices., same as chinese version.
tardis-key
force-pushed
the
main
branch
2 times, most recently
from
892d57f to
9af5016
Compare
|
After configuring profiling for tests/speical_npu/run_qwen2_5_05b_grpo.sh, the profiling directory is as follows,
And the result of tests/speical_npu/run_qwen2_5_05b_grpo_mindspeed.sh are also consistent. Due to an exception in collecting "generate" under AgentLoop, the "generate_sequences" folder was not obtained. This will be fixed in another PR. |
tardis-key
force-pushed
the
main
branch
2 times, most recently
from
b890bf5 to
bd7868c
Compare
| @@ -1,10 +1,10 @@ | |||
| Data collection based on FSDP backend on Ascend devices(en) | |||
| Data collection based on FSDP or MindSpeed(Megatron) on Ascend devices(en) | |||
There was a problem hiding this comment.
Change title to Performance Data collection based on ..., same as the chinese version.
verl/workers/megatron_workers.py
Outdated
|
|
||
| @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="critic")) | ||
| @DistProfiler.annotate(color="pink") | ||
| @DistProfiler.annotate(color="pink", role="update_critic") |
There was a problem hiding this comment.
Align with actor, here role should be critic_update
verl/workers/fsdp_workers.py
Outdated
|
|
||
| @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="critic")) | ||
| @DistProfiler.annotate(color="pink") | ||
| @DistProfiler.annotate(color="pink", role="update_critic") |
There was a problem hiding this comment.
same here, change the role name to critic_update
Suggest fix the rollout period data unavailable error in this PR @tardis-key |
|
All review comments have been revised, and the pr has been rebased to latest code. |
#4320 is working on this. And that fix involves profiling of async inference engine and will take some time. As VeRL is to be released soon, it is recommended that this pr should be merged to enable the discrete profiling on MindSpeed. |
tardis-key
changed the title
Co-authored-by: mengchengTang <[email protected]> Co-authored-by: Shangwei-Li <[email protected]>
2 participants
What does this PR do?
#4206
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 (飞书群).)