Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / callbacks / profiler_callback.py: 90%
161 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""MindFormers-derived training profiler backed by PyTorch."""
17from __future__ import annotations
19import json
20import logging
21from pathlib import Path
22from typing import TYPE_CHECKING, Any, Optional
24# Trainer is Torch-owned; match the other callbacks under the legacy backend-import checker.
25import torch # pylint: disable=forbidden-backend-import
27from hyper_parallel.trainer.callbacks.base import Callback
28from hyper_parallel.trainer.config.training import ProfilingAnalysisConfig, ProfilingConfig
29from hyper_parallel.trainer.runtime.device import get_device_type
30from hyper_parallel.trainer.state import TrainerState
33if TYPE_CHECKING:
34 from hyper_parallel.trainer.base import BaseTrainer
37logger = logging.getLogger(__name__)
40class ProfilerCallback(Callback):
41 """Collect a bounded CPU, CUDA, or Ascend NPU profile.
43 The profiler schedule, rank selection, output layout, metadata, and NPU
44 options follow MindFormers ``ProfileMonitor`` semantics. PyTorch profiler
45 activities replace the corresponding MindSpore profiler backend.
46 """
48 def __init__(self, trainer: "BaseTrainer") -> None:
49 """Build the profiler selected by ``TrainerConfig.profiler``.
51 Args:
52 trainer: Trainer that owns the callback lifecycle.
53 """
54 super().__init__(trainer)
55 self.config: ProfilingConfig = trainer.config.profiler
56 self.profiler: Any = None
57 self.output_path: Optional[Path] = None
58 self._profiler_module: Any = None
59 self._device_type: Optional[str] = None
60 self._is_profiler_started = False
61 self._metadata_recorded = False
62 self._training_step = 0
64 if not self.config.enable:
65 return
67 self._validate_types()
68 self.start_step, self.stop_step = self._normalize_steps()
69 self.level = self._normalize_level()
71 rank = trainer.global_rank
72 if not self._is_profile_required(rank):
73 return
75 self._device_type, self._profiler_module = self._resolve_profiler_module()
77 output_root = Path(self.config.output_path or "./output").expanduser()
78 self.output_path = output_root / "profile" / f"rank_{rank}"
79 logger.info("Profile save path: %s", self.output_path)
81 schedule_config = self._get_schedule()
82 trace_handler_options = {}
83 if self._device_type == "npu":
84 trace_handler_options = {
85 "analyse_flag": self.config.analysis.enable,
86 "async_mode": self.config.analysis.async_mode,
87 }
88 profile_kwargs = {
89 "activities": self._get_activities(),
90 "profile_memory": self.config.memory,
91 "with_stack": self.config.with_stack,
92 "schedule": schedule_config,
93 "on_trace_ready": self._profiler_module.tensorboard_trace_handler(
94 str(self.output_path), **trace_handler_options
95 ),
96 }
97 if self._device_type == "npu":
98 profile_kwargs["experimental_config"] = self._get_npu_experimental_config()
99 self.profiler = self._profiler_module.profile(**profile_kwargs)
101 def on_step_begin(
102 self,
103 state: TrainerState,
104 micro_batches: Optional[list[dict[str, Any]]] = None,
105 **kwargs: Any,
106 ) -> None:
107 """Start the profiler and align its schedule with the first training step.
109 Args:
110 state: Current trainer state.
111 micro_batches: Optional micro-batches prepared for the current step.
112 **kwargs: Additional callback context supplied by the trainer.
113 """
114 del state, micro_batches, kwargs
115 if self.profiler is None:
116 return
118 self._training_step += 1
119 if not self._is_profiler_started:
120 self._start()
121 if self._training_step == self.start_step:
122 self._record_metadata()
124 def on_step_end(
125 self,
126 state: TrainerState,
127 loss: Optional[float] = None,
128 loss_dict: Optional[dict[str, float]] = None,
129 grad_norm: Optional[float] = None,
130 **kwargs: Any,
131 ) -> None:
132 """Advance the profiler schedule after one complete optimizer step.
134 Args:
135 state: Current trainer state.
136 loss: Optional scalar loss for the completed step.
137 loss_dict: Optional named losses for the completed step.
138 grad_norm: Optional gradient norm for the completed step.
139 **kwargs: Additional callback context supplied by the trainer.
140 """
141 del state, loss, loss_dict, grad_norm, kwargs
142 if self.profiler is None or not self._is_profiler_started:
143 return
145 self.profiler.step()
146 if self._training_step == self.stop_step:
147 logger.info("End of profiling. Analyze the trace under %s", self.output_path)
148 self._stop()
150 def on_train_end(self, state: TrainerState, **kwargs: Any) -> None:
151 """Flush a partial profile when training ends before the stop step.
153 Args:
154 state: Final trainer state.
155 **kwargs: Additional callback context supplied by the trainer.
156 """
157 del state, kwargs
158 self.close()
160 def close(self) -> None:
161 """Stop collection and release a profiler after completion or failure."""
162 self._stop()
164 def _validate_types(self) -> None:
165 """Validate profiler fields that dataclass construction cannot constrain."""
166 boolean_fields = (
167 "enable",
168 "memory",
169 "with_stack",
170 "data_simplification",
171 )
172 for name in boolean_fields:
173 if not isinstance(getattr(self.config, name), bool):
174 raise TypeError(f"profiler.{name} must be a bool")
175 if not isinstance(self.config.analysis, ProfilingAnalysisConfig):
176 raise TypeError("profiler.analysis must be a ProfilingAnalysisConfig")
177 for name in ("enable", "async_mode"):
178 if not isinstance(getattr(self.config.analysis, name), bool):
179 raise TypeError(f"profiler.analysis.{name} must be a bool")
180 for name in ("start_step", "stop_step", "warmup_steps"):
181 value = getattr(self.config, name)
182 if isinstance(value, bool) or not isinstance(value, int):
183 raise TypeError(f"profiler.{name} must be an int")
184 if self.config.output_path is not None and not isinstance(self.config.output_path, str):
185 raise TypeError("profiler.output_path must be a string or null")
186 ranks = self.config.rank_ids
187 if ranks is not None:
188 if not isinstance(ranks, list):
189 raise TypeError("profiler.rank_ids must be a list or null")
190 if any(isinstance(rank, bool) or not isinstance(rank, int) or not 0 <= rank < self.trainer.world_size
191 for rank in ranks):
192 raise ValueError(f"profiler.rank_ids must contain only integers in [0, {self.trainer.world_size})")
194 def _normalize_steps(self) -> tuple[int, int]:
195 """Normalize the inclusive profiling window to valid positive steps."""
196 start_step = self.config.start_step
197 stop_step = self.config.stop_step
198 if start_step < 1:
199 logger.warning("profiler.start_step must be greater than 0; reset it to 1")
200 start_step = 1
201 if stop_step < 1:
202 logger.warning("profiler.stop_step must be greater than 0; reset it to 10")
203 stop_step = 10
204 if start_step > stop_step:
205 logger.warning(
206 "profiler.stop_step must be greater than or equal to profiler.start_step; reset both to 1 and 10"
207 )
208 start_step, stop_step = 1, 10
209 if not 0 <= self.config.warmup_steps < start_step:
210 raise ValueError("profiler.warmup_steps must be non-negative and less than profiler.start_step")
211 return start_step, stop_step
213 def _normalize_level(self) -> int:
214 level = self.config.level
215 if level is None:
216 return 0
217 if isinstance(level, bool) or not isinstance(level, int):
218 raise TypeError("profiler.level must be an int or null")
219 if level not in (0, 1, 2):
220 logger.warning("Invalid profiler.level %s; reset it to 0", level)
221 return 0
222 return level
224 def _get_schedule(self) -> Any:
225 return self._profiler_module.schedule(
226 wait=0,
227 warmup=self.config.warmup_steps,
228 active=self.stop_step - self.start_step + 1,
229 repeat=1,
230 skip_first=self.start_step - self.config.warmup_steps,
231 )
233 def _is_profile_required(self, rank: int) -> bool:
234 """Return whether the current rank is selected for profiling."""
235 ranks = self.config.rank_ids or []
236 return not ranks or rank in ranks
238 def _resolve_profiler_module(self) -> tuple[str, Any]:
239 """Resolve the profiler backend for the active device type."""
240 device_type = get_device_type()
241 if device_type == "npu":
242 # torch-npu is optional and must not become a trainer import-time dependency.
243 import torch_npu # pylint: disable=import-outside-toplevel
245 return device_type, torch_npu.profiler
246 if device_type in ("cpu", "cuda"):
247 return device_type, torch.profiler
248 raise RuntimeError(f"profiler does not support device type {device_type!r}")
250 def _get_activities(self) -> list[Any]:
251 activities = [self._profiler_module.ProfilerActivity.CPU]
252 if self._device_type == "cuda":
253 activities.append(self._profiler_module.ProfilerActivity.CUDA)
254 elif self._device_type == "npu":
255 activities.append(self._profiler_module.ProfilerActivity.NPU)
256 return activities
258 def _get_npu_experimental_config(self) -> Any:
259 """Build the torch-npu experimental profiler configuration."""
260 profiler_level = getattr(self._profiler_module.ProfilerLevel, f"Level{self.level}")
261 options = {
262 "profiler_level": profiler_level,
263 "data_simplification": self.config.data_simplification,
264 "mstx": False,
265 }
266 export_type = getattr(getattr(self._profiler_module, "ExportType", None), "Text", None)
267 if export_type is not None:
268 options["export_type"] = [export_type]
269 # torch-npu exposes this version-dependent profiler entry point only under its private name.
270 experimental_config = getattr(self._profiler_module, "_ExperimentalConfig")
271 return experimental_config(**options)
273 def _start(self) -> None:
274 self.profiler.start()
275 self.profiler.step()
276 self._is_profiler_started = True
278 def _record_metadata(self) -> None:
279 """Record distributed topology metadata once per profiling session."""
280 if self._metadata_recorded:
281 return
282 parallel = self.trainer.config.accelerator
283 metadata = {
284 "tensor_model_parallel_size": parallel.tp_size,
285 "pipeline_model_parallel_size": parallel.pp_size,
286 "data_parallel_size": getattr(self.trainer.mesh, "dp_size", 1),
287 "expert_model_parallel_size": parallel.ep_size,
288 "sequence_parallel": parallel.sequence_parallel,
289 "parallel_mode": "distributed" if self.trainer.world_size > 1 else "stand_alone",
290 "world_size": self.trainer.world_size,
291 }
292 try:
293 self.profiler.add_metadata_json("distributed_args", json.dumps(metadata))
294 self._metadata_recorded = True
295 except AttributeError as error:
296 logger.warning("Profiler failed to record distributed args: %s", error)
298 def _stop(self) -> None:
299 """Stop the profiler session and release its resources."""
300 profiler = self.profiler
301 self.profiler = None
302 try:
303 if profiler is not None and self._is_profiler_started:
304 profiler.stop()
305 finally:
306 self._is_profiler_started = False
309__all__ = ["ProfilerCallback"]