[perf] feat: mstx profiler support schedule to mini batch#7105
[perf] feat: mstx profiler support schedule to mini batch#7105mengchengTang wants to merge 1 commit into
Conversation
![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 adds support for scheduled profiling in the NPU profiler using torch_npu.profiler.schedule configurations (such as wait, warmup, active, repeat, and skip_first). This allows users to capture a subset of mini-batches during training update loops to reduce profiling data volume. The feedback suggests tracking the active schedule configuration globally at the class level (_active_schedule_kwargs) within NPUProfiler. This prevents potential AttributeError and state mismatches when different profiler instances (e.g., rollout and training workers) start and stop the process-wide NPU profiler.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| _define_count = 0 | ||
| # Process-global handle to the currently running NPU profiler. torch_npu.profiler is | ||
| # process-wide, so a step() issued by one NPUProfiler instance (e.g. the inner actor | ||
| # TrainingWorker running the mini-batch loop) must advance the profiler that another | ||
| # instance started (e.g. the outer ActorRolloutRefWorker). | ||
| _active_prof = None |
There was a problem hiding this comment.
To support safe start/stop operations across multiple NPUProfiler instances in the same process (e.g., ActorRolloutRefWorker and TrainingWorker), we should track the active schedule globally at the class level. This prevents instance-specific state mismatches when one instance stops a profiler started by another.
| _define_count = 0 | |
| # Process-global handle to the currently running NPU profiler. torch_npu.profiler is | |
| # process-wide, so a step() issued by one NPUProfiler instance (e.g. the inner actor | |
| # TrainingWorker running the mini-batch loop) must advance the profiler that another | |
| # instance started (e.g. the outer ActorRolloutRefWorker). | |
| _active_prof = None | |
| _define_count = 0 | |
| # Process-global handle to the currently running NPU profiler. torch_npu.profiler is | |
| # process-wide, so a step() issued by one NPUProfiler instance (e.g. the inner actor | |
| # TrainingWorker running the mini-batch loop) must advance the profiler that another | |
| # instance started (e.g. the outer ActorRolloutRefWorker). | |
| _active_prof = None | |
| _active_schedule_kwargs = None |
| def start(self, **kwargs): | ||
| role = kwargs.get("role", None) | ||
| if not self.discrete and NPUProfiler._define_count == 0: | ||
| self._schedule_kwargs = self._resolve_schedule_kwargs() | ||
| self.profile_npu = get_npu_profiler( | ||
| contents=self.profile_contents, | ||
| profile_level=self.profile_level, | ||
| profile_save_path=self.profile_save_path, | ||
| analysis=self.analysis, | ||
| role=role, | ||
| schedule=self._schedule_kwargs, | ||
| ) | ||
| self.profile_npu.start() | ||
| NPUProfiler._active_prof = self.profile_npu | ||
| NPUProfiler._define_count += 1 |
There was a problem hiding this comment.
Store the active schedule at the class level (NPUProfiler._active_schedule_kwargs) when starting the profiler, so that any instance stopping the profiler can correctly determine if a schedule is active.
| def start(self, **kwargs): | |
| role = kwargs.get("role", None) | |
| if not self.discrete and NPUProfiler._define_count == 0: | |
| self._schedule_kwargs = self._resolve_schedule_kwargs() | |
| self.profile_npu = get_npu_profiler( | |
| contents=self.profile_contents, | |
| profile_level=self.profile_level, | |
| profile_save_path=self.profile_save_path, | |
| analysis=self.analysis, | |
| role=role, | |
| schedule=self._schedule_kwargs, | |
| ) | |
| self.profile_npu.start() | |
| NPUProfiler._active_prof = self.profile_npu | |
| NPUProfiler._define_count += 1 | |
| def start(self, **kwargs): | |
| role = kwargs.get("role", None) | |
| if not self.discrete and NPUProfiler._define_count == 0: | |
| self._schedule_kwargs = self._resolve_schedule_kwargs() | |
| self.profile_npu = get_npu_profiler( | |
| contents=self.profile_contents, | |
| profile_level=self.profile_level, | |
| profile_save_path=self.profile_save_path, | |
| analysis=self.analysis, | |
| role=role, | |
| schedule=self._schedule_kwargs, | |
| ) | |
| self.profile_npu.start() | |
| NPUProfiler._active_prof = self.profile_npu | |
| NPUProfiler._active_schedule_kwargs = self._schedule_kwargs | |
| NPUProfiler._define_count += 1 |
| def stop(self): | ||
| if not self.discrete and NPUProfiler._define_count == 1: | ||
| self.profile_npu.step() | ||
| # Continuous mode emits a trailing step to flush the final window; when a | ||
| # schedule is configured, stepping is driven per mini-batch instead. | ||
| if not self._schedule_kwargs: | ||
| self.step() | ||
| self.profile_npu.stop() | ||
| NPUProfiler._active_prof = None | ||
| self._schedule_kwargs = None | ||
| NPUProfiler._define_count -= 1 |
There was a problem hiding this comment.
When stopping the profiler, use the class-level NPUProfiler._active_prof and NPUProfiler._active_schedule_kwargs instead of instance-specific variables. This avoids AttributeError when stop() is called on an instance that did not start the profiler, and ensures the correct schedule check is performed.
| def stop(self): | |
| if not self.discrete and NPUProfiler._define_count == 1: | |
| self.profile_npu.step() | |
| # Continuous mode emits a trailing step to flush the final window; when a | |
| # schedule is configured, stepping is driven per mini-batch instead. | |
| if not self._schedule_kwargs: | |
| self.step() | |
| self.profile_npu.stop() | |
| NPUProfiler._active_prof = None | |
| self._schedule_kwargs = None | |
| NPUProfiler._define_count -= 1 | |
| def stop(self): | |
| if not self.discrete and NPUProfiler._define_count == 1: | |
| # Continuous mode emits a trailing step to flush the final window; when a | |
| # schedule is configured, stepping is driven per mini-batch instead. | |
| if not NPUProfiler._active_schedule_kwargs: | |
| self.step() | |
| if NPUProfiler._active_prof is not None: | |
| NPUProfiler._active_prof.stop() | |
| NPUProfiler._active_prof = None | |
| NPUProfiler._active_schedule_kwargs = None | |
| self._schedule_kwargs = None | |
| NPUProfiler._define_count -= 1 |
mengchengTang
changed the title
6d5b7d2 to
b35678c
Compare
1 participant
What does this PR do?
flow #7099
Checklist Before Starting
[{modules}] {type}: {description}(This will be checked by the CI){modules}includefsdp,megatron,veomni,sglang,vllm,rollout,trainer,ci,training_utils,recipe,hardware,deployment,ray,worker,single_controller,misc,perf,model,algo,env,tool,ckpt,doc,data,cfg,reward,fully_async,one_step_off,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 (飞书群).)recipesubmodule, please also update the reference to the submodule commit viagit submodule update --remoteorcd recipe && git pull origin main.