« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025 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"""memory estimation API"""
16# pylint: disable=protected-access
17from __future__ import annotations
18from typing import TYPE_CHECKING
19import argparse
20import os
21import copy
23# import pprint
24import json
25import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as Hard
26from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
28from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
29from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._utils import _Utils
30from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._hook_manager import _HookManager
31from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory
33if TYPE_CHECKING:
34 from typing import Any, Callable, Dict, Optional, Sequence, Union
35 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.hook_base import MemEvalHook
38class EvaluatorV2(_Utils, _HookManager):
39 """Memory evaluator class"""
41 def __init__(self, *args, **kwargs):
42 self.ppb = None
43 super().__init__(*args, **kwargs)
44 self._child_cls = self
46 def reset_config(self) -> None:
47 """reset current config"""
48 self.update_config(self.config_path)
50 def load_hook_cls(self, hook_cls: MemEvalHook) -> None:
51 """load an instance of MemEvalHook"""
52 self.hook_cls = hook_cls
53 self.reset_config()
54 if not self._ccfg.multimodal:
55 hook = list(self._ccfg.hooks_dict.values())[0]
56 hook(self)
58 def _estimate_on_copy(self, *args) -> Any:
59 """Run the backbone on a copy of the config, leaving the caller's intact.
61 The backbone applies each layer's hook in place, one layer after the
62 next. Run on the caller's config, it would leave the last layer's
63 state behind for whatever reads that config next, which in a search
64 is the next candidate.
65 """
66 caller = self._ccfg
67 self._ccfg = copy.deepcopy(caller)
68 self._overhead_obj._ccfg = self._ccfg
69 try:
70 return self._estimate_backbone(*args)
71 finally:
72 self._ccfg = caller
73 self._overhead_obj._ccfg = caller
75 def estimate_peak(
76 self,
77 stages: list = None,
78 verbose: bool = False,
79 spec_stage_id: int = -1,
80 plot: bool = False,
81 ) -> float:
82 """Peak stage memory estimation"""
83 res = self._estimate_on_copy(
84 stages, verbose, False, spec_stage_id, plot
85 )
86 insights, _ = res
87 stage_mems = [i["Static"] + i["Dynamic"] for i in insights]
88 peak_mem = max(stage_mems)
89 peak_stage = stage_mems.index(peak_mem)
91 logger.output(
92 "model_name: %s, peak memory : \033[1m%s MB\033[0m (stage _%s)",
93 self._ccfg.model_name,
94 peak_mem,
95 peak_stage,
96 )
97 return peak_mem
99 def estimate_peak_insight(self, stages: list = None) -> list:
100 """subcomponents' proportion estimation"""
101 insights, _ = self._estimate_on_copy(stages, False, False, -1, False)
102 return insights
104 def estimate_layer_memory(
105 self,
106 stages: list = None,
107 device_type=Hard.Device_A2,
108 layer_times: Optional[Callable] = None,
109 ) -> Dict:
110 """PPB's input
112 Args:
113 layer_times: Prices a layer as ``(forward, backward)`` from its
114 config, kind, layer type and recompute switches, such as
115 ``perf_estimation.estimate.LayerTimes``. With it, every layer
116 description carries its times, in units of the first body's
117 forward time, instead of a placeholder and offers the
118 selective recompute options worth pricing, and the result is
119 not cached.
120 """
121 logger.info(device_type)
122 if layer_times is not None:
123 with self._ppb_obj.timing(layer_times):
124 return self._estimate_on_copy(stages, False, True, -1, False)[1]
125 if self.ppb:
126 return self.ppb
127 res = self._estimate_on_copy(stages, False, True, -1, False)
128 _, ppb = res
129 self.ppb = ppb
130 return ppb
132 def estimate_switch_profiles(
133 self,
134 layer_times: Callable,
135 stages: list = None,
136 most_in_flight: Optional[int] = None,
137 each_switch: bool = True,
138 in_flight: Sequence[int] = (),
139 ) -> Dict:
140 """What each recompute switch saves and costs, per model and layer kind.
142 Walks the stages as :meth:`estimate_layer_memory` does and measures
143 the first body layer of each kind: plain, with each op the switches
144 name recomputed alone, and fully recomputed. The config is left as it
145 was found.
147 Args:
148 layer_times: Prices a layer, as for :meth:`estimate_layer_memory`.
149 stages: The partition; the config's own when omitted.
150 most_in_flight: The most micro-batches any stage keeps in flight.
151 A buffer is split into what grows with the micro-batches in
152 flight and what is kept once by its size at one micro-batch
153 and at this many, so the split is exact at both. Under 1F1B
154 without interleaving, the least of the stages and the
155 micro-batches when omitted.
156 each_switch: Whether to measure each switch alone. Without, a
157 profile prices only the plain and the fully recomputed
158 layer, which is all one mode for every layer needs unless it
159 is selective.
160 in_flight: The counts of micro-batches in flight the stages
161 keep. At each between one and the most, a profile also
162 states what the split charges beyond the buffers kept, so
163 that an option's memory is exact at every count.
165 Returns:
166 ``{(model name, layer kind): SwitchProfile}``, in model order.
167 """
168 with self._ppb_obj.profiling(layer_times, most_in_flight, each_switch, in_flight) as profiles:
169 self._estimate_on_copy(stages, False, True, -1, False)
170 return profiles
172 # Specific estimation
174 def static_mem_stage(self, stage_id: int) -> float:
175 """stage estimation proportions"""
176 insights = self.estimate_peak_insight()
177 return insights[stage_id]["Static"]
179 def dynamic_mem_stage(self, stage_id: int) -> float:
180 """stage estimation proportions"""
181 insights = self.estimate_peak_insight()
182 return insights[stage_id]["Dynamic"]
184 def logs_mem_stage(self, stage_id: int) -> list:
185 """stage estimation hook calls trace"""
186 insights = self.estimate_peak_insight()
187 return insights[stage_id]["Node Log"]
189 def static_mem_layer(
190 self, node: Union[str, LayerType], stage_id: int
191 ) -> Union[float, list]:
192 """accounting all hooks for this node"""
193 logs = self.logs_mem_stage(stage_id)
194 extracted_log = set()
195 for k in logs.keys():
196 if k[3] == node.name[0]:
197 extracted_log.add(logs[k]["_param"])
198 extracted_log = list(extracted_log)
199 if len(extracted_log) == 1:
200 return extracted_log[0]
201 return extracted_log
203 def dynamic_mem_layer(
204 self, node: Union[str, LayerType], stage_id: int
205 ) -> Union[float, list]:
206 """accounting all hooks for this node"""
207 logs = self.logs_mem_stage(stage_id)
208 extracted_log = set()
209 for k in logs.keys():
210 if k[3] == node.name[0]:
211 extracted_log.add(logs[k]["_activ"] + logs[k]["_comm"])
212 extracted_log = list(extracted_log)
213 if len(extracted_log) == 1:
214 return extracted_log[0]
215 return extracted_log
217 def mem_fit(
218 self, mem: float, tolerance: float = 0, margin: float = 0
219 ) -> bool:
220 """check if input memory fits in device"""
221 # Expect arguments to be in MB
222 memory = Memory.from_mb(mem)
223 tolerance_mem = Memory.from_mb(tolerance)
224 margin_mem = Memory.from_mb(margin)
225 cap = self._ccfg.device_capacity - margin_mem
226 diff = abs(memory - cap)
227 is_close = diff <= tolerance_mem
228 is_fit = memory <= cap
230 if Memory.zero() < tolerance_mem and is_close:
231 logger.info(
232 "Prediction is CLOSE to memory device (%s of diff)", diff
233 )
234 return True
235 if is_fit:
236 logger.output(
237 "estimation FITS into device memory (%s<=%s-%s)",
238 mem,
239 self._ccfg.device_capacity,
240 margin_mem,
241 )
242 else:
243 logger.output(
244 "estimation DOES NOT FIT into device memory (%s>%s-%s)",
245 mem,
246 self._ccfg.device_capacity,
247 margin_mem,
248 )
249 return is_fit
252def estimate_memory(config: Any) -> bool:
253 """fast usage for estimation/fit for given input config"""
254 e = EvaluatorV2(config)
255 peak = e.estimate_peak()
256 return e.mem_fit(peak)
259def main():
260 """commandline"""
261 parser = argparse.ArgumentParser(
262 description="Command line usage: Estimate peak stage memory"
263 )
265 parser.add_argument(
266 "model_config_path",
267 nargs=1,
268 help="Model config file (MindFormer YAML or MindSpeed JSON)",
269 )
271 parser.add_argument(
272 "--framework",
273 default=None,
274 type=str,
275 help="Specify a framework name",
276 )
278 parser.add_argument(
279 "--code-path",
280 default=None,
281 type=str,
282 help="Specify a source code path (Additional parsing)",
283 )
285 parser.add_argument(
286 "--verbose", action="store_true", help="Show estimation trace"
287 )
288 parser.add_argument("--plot", action="store_true", help="Plot estimation")
289 parser.add_argument(
290 "--fit",
291 action="store_true",
292 help="Check if estimation fits in device memory",
293 )
294 parser.add_argument(
295 "--stage", default=-1, type=int, help="Specify pipeline stage ID"
296 )
297 parser.add_argument(
298 "--hook",
299 default=None,
300 type=str,
301 help="Specify hook class (defined in hooks/)",
302 )
303 parser.add_argument(
304 "--trace-fun",
305 default=None,
306 type=str,
307 help="Specify a formula function name to get it traced",
308 )
309 parser.add_argument(
310 "--ppb",
311 action="store_true",
312 help="Generate pipeline balancing layers description",
313 )
314 parser.add_argument(
315 "--ctx", action="store_true", help="Show ctx variables"
316 )
317 parser.add_argument(
318 "--ccfg", action="store_true", help="Show ccfg variables"
319 )
320 parser.add_argument(
321 "--warnings", action="store_true", help="Show warnings"
322 )
323 args = parser.parse_args()
325 path = args.model_config_path[0]
326 if not os.path.exists(path):
327 raise argparse.ArgumentTypeError(f"`{path}` was not found")
328 if not path.endswith((".yaml", ".json",".toml")):
329 raise argparse.ArgumentTypeError(f"`{path}` has invalid file type")
331 e = EvaluatorV2(
332 path,
333 framework=args.framework,
334 source_code=args.code_path,
335 log_level=args.warnings,
336 hook_cls=args.hook,
337 trace_fun=args.trace_fun,
338 )
339 if args.ctx:
340 e.print_ctx()
341 if args.ccfg:
342 e.print_ccfg()
343 if args.ppb:
344 ppb = e.estimate_layer_memory()
345 print(json.dumps(ppb, indent=2))
346 else:
347 peak_mem = e.estimate_peak(
348 verbose=args.verbose, spec_stage_id=args.stage, plot=args.plot
349 )
351 if args.fit:
352 e.mem_fit(peak_mem)
355if __name__ == "__main__":
356 main()