Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / runtime / memory_profiler.py: 96%
180 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"""Process-local allocator memory profiler used by Trainer runtimes."""
17from __future__ import annotations
19import logging
20import os
21import stat
22import sys
23import time
24from typing import Any, Optional
26from hyper_parallel.trainer.config.training import MemoryConfig
27from hyper_parallel.trainer.runtime.device import get_device_type, get_torch_device
30logger = logging.getLogger(__name__)
33class MemoryProfiler:
34 """Manage one allocator-history profiling session in the current process.
36 Note:
37 The module owns one singleton, so only one Trainer profiling session
38 may be active in a process at a time.
39 """
41 def __init__(self) -> None:
42 """Initialize an inactive profiler."""
43 self.enable = False
44 self.mem_info = False
45 self.current_step = 0
46 self.start_step = 0
47 self.end_step = 0
48 self.save_path = ""
49 self.dump_ranks: list[int] = []
50 self.stacks = "all"
51 self.global_rank = 0
52 self.tp_rank = 0
53 self.dp_rank = 0
54 self._session_active = False
55 self._history_started = False
56 self._device_api: Any = None
57 self._memory_api: Any = None
59 def reset(
60 self,
61 config: Optional[MemoryConfig],
62 *,
63 global_rank: int,
64 tp_rank: int,
65 dp_rank: int,
66 ) -> None:
67 """Replace the process-local session with a Trainer configuration.
69 Args:
70 config: Memory profiler configuration, or ``None`` to disable it.
71 global_rank: Rank in the global process group.
72 tp_rank: Rank in the tensor-parallel group.
73 dp_rank: Rank in the data-parallel group.
74 """
75 self._discard_previous_session()
76 if config is None:
77 return
79 self._validate(config)
80 self.enable = config.enable
81 self.mem_info = config.mem_info
82 self.current_step = 0
83 self.start_step = config.start_step
84 self.end_step = config.end_step
85 self.save_path = config.save_path
86 self.dump_ranks = list(config.dump_ranks)
87 self.stacks = config.stacks
88 self.global_rank = global_rank
89 self.tp_rank = tp_rank
90 self.dp_rank = dp_rank
91 self._session_active = self.enable or self.mem_info
93 if self.enable:
94 self._resolve_device_apis(require_history=True)
95 # Match MindSpeed's session-relative control point: step zero runs
96 # after distributed setup and before model construction.
97 self.step()
99 def step(self) -> None:
100 """Advance the state machine before a Trainer training step starts."""
101 if not self._session_active:
102 return
104 try:
105 if self.enable:
106 if self.current_step == self.start_step:
107 self._record()
108 # The control point precedes training, so dump after the inclusive end step.
109 if self.current_step == self.end_step + 1:
110 try:
111 self._dump_and_stop_history()
112 finally:
113 self.enable = False
114 if not self.mem_info and not self._history_started:
115 self._deactivate_session()
117 if self.mem_info:
118 self._log_memory_info()
119 except BaseException:
120 self.abort()
121 raise
122 finally:
123 self.current_step += 1
125 if not self.enable and not self.mem_info:
126 self._deactivate_session()
128 def stop(self) -> None:
129 """Dump an active partial window and end a normal training session."""
130 if not self._session_active and not self._history_started:
131 return
133 try:
134 try:
135 if self._history_started:
136 self._dump_and_stop_history()
137 finally:
138 self.enable = False
139 self.mem_info = False
140 if not self._history_started:
141 self._deactivate_session()
142 except BaseException:
143 self.abort()
144 raise
146 def abort(self) -> None:
147 """Stop allocator history without dumping or masking a Trainer error."""
148 self.enable = False
149 self.mem_info = False
150 try:
151 self._stop_history()
152 except Exception: # pylint: disable=broad-exception-caught
153 logger.exception("Failed to stop allocator memory history while aborting training.")
154 return
156 self._deactivate_session()
158 def _discard_previous_session(self) -> None:
159 if self._history_started:
160 self._stop_history()
161 self._deactivate_session()
163 def _deactivate_session(self) -> None:
164 self.enable = False
165 self.mem_info = False
166 self._session_active = False
167 self._device_api = None
168 self._memory_api = None
170 @staticmethod
171 def _validate(config: MemoryConfig) -> None:
172 MemoryProfiler._validate_bool(config.enable, "enable")
173 MemoryProfiler._validate_bool(config.mem_info, "mem_info")
174 MemoryProfiler._validate_steps(config.start_step, config.end_step)
175 MemoryProfiler._validate_save_path(config.save_path)
176 MemoryProfiler._validate_dump_ranks(config.dump_ranks)
177 MemoryProfiler._validate_stacks(config.stacks)
179 @staticmethod
180 def _validate_bool(value: bool, name: str) -> None:
181 if not isinstance(value, bool):
182 raise TypeError(f"memory.{name} must be a bool")
184 @staticmethod
185 def _validate_steps(start_step: int, end_step: int) -> None:
186 for name, value in (("start_step", start_step), ("end_step", end_step)):
187 if isinstance(value, bool) or not isinstance(value, int):
188 raise TypeError(f"memory.{name} must be an int")
189 if start_step < 0:
190 raise ValueError("memory.start_step must be non-negative")
191 if end_step < start_step:
192 raise ValueError("memory.end_step must be greater than or equal to memory.start_step")
194 @staticmethod
195 def _validate_save_path(save_path: str) -> None:
196 if not isinstance(save_path, str) or not save_path:
197 raise ValueError("memory.save_path must be a non-empty string")
199 @staticmethod
200 def _validate_dump_ranks(dump_ranks: list[int]) -> None:
201 if not isinstance(dump_ranks, list):
202 raise ValueError("memory.dump_ranks must be a list of non-negative integers")
203 if any(
204 isinstance(rank, bool) or not isinstance(rank, int) or rank < 0
205 for rank in dump_ranks
206 ):
207 raise ValueError("memory.dump_ranks must be a list of non-negative integers")
209 @staticmethod
210 def _validate_stacks(stacks: str) -> None:
211 if stacks not in ("python", "all"):
212 raise ValueError("memory.stacks must be 'python' or 'all'")
214 def _record(self) -> None:
215 """Start allocator memory history recording."""
216 self._history_started = True
217 try:
218 # Allocator history has no public API in Torch or torch-npu.
219 self._memory_api._record_memory_history( # pylint: disable=protected-access
220 stacks=self.stacks,
221 max_entries=sys.maxsize,
222 )
223 except Exception:
224 # The private allocator API does not expose whether a failed call
225 # enabled history, so make a best-effort stop before propagating.
226 self.abort()
227 raise
229 def _dump_and_stop_history(self) -> None:
230 try:
231 self._dump()
232 finally:
233 self._stop_history()
235 def _dump(self) -> None:
236 """Dump the allocator snapshot for a configured rank."""
237 if self.global_rank not in self.dump_ranks:
238 return
239 os.makedirs(self.save_path, exist_ok=True)
240 timestamp = time.strftime("%Y-%m-%d-%H-%M")
241 file_path = os.path.join(
242 self.save_path,
243 f"snapshot_{timestamp}_{self.global_rank}.pickle",
244 )
245 # Allocator snapshots have no public API in Torch or torch-npu.
246 self._memory_api._dump_snapshot(file_path) # pylint: disable=protected-access
247 # Some torch-npu releases create snapshots with owner-write permission only.
248 os.chmod(file_path, stat.S_IRUSR | stat.S_IWUSR)
249 logger.info("Memory snapshot dumped to %s", file_path)
251 def _stop_history(self) -> None:
252 if not self._history_started:
253 return
254 self._memory_api._record_memory_history(enabled=None) # pylint: disable=protected-access
255 self._history_started = False
257 def _log_memory_info(self) -> None:
258 """Log and reset peak allocator statistics for the current step."""
259 self._resolve_device_apis(require_history=False)
260 logger.info(
261 "Memory usage step=%s global_rank=%s tp_rank=%s dp_rank=%s "
262 "max_memory_reserved=%s max_memory_allocated=%s",
263 self.current_step,
264 self.global_rank,
265 self.tp_rank,
266 self.dp_rank,
267 self._device_api.max_memory_reserved(),
268 self._device_api.max_memory_allocated(),
269 )
270 self._device_api.reset_peak_memory_stats()
272 def _resolve_device_apis(self, *, require_history: bool) -> None:
273 """Resolve and validate the active device memory APIs."""
274 if self._device_api is None:
275 device_type = get_device_type()
276 if device_type not in ("cuda", "npu"):
277 raise RuntimeError(
278 "memory profiling requires a CUDA or NPU device, "
279 f"but the active device type is {device_type!r}"
280 )
281 self._device_api = get_torch_device()
283 if not require_history or self._memory_api is not None:
284 return
285 self._memory_api = getattr(self._device_api, "memory", None)
286 required_methods = ("_record_memory_history", "_dump_snapshot")
287 missing = [
288 method
289 for method in required_methods
290 if self._memory_api is None
291 or not callable(getattr(self._memory_api, method, None))
292 ]
293 if missing:
294 raise RuntimeError(
295 "memory profiling backend is missing allocator API(s): "
296 + ", ".join(missing)
297 )
300memory_profiler = MemoryProfiler()
303__all__ = ["memory_profiler"]