« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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, Dict, 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_peak(
59 self,
60 stages: list = None,
61 verbose: bool = False,
62 spec_stage_id: int = -1,
63 plot: bool = False,
64 ) -> float:
65 """Peak stage memory estimation"""
66 original_ccfg = copy.deepcopy(self._ccfg)
67 res = self._estimate_backbone(
68 stages, verbose, False, spec_stage_id, plot
69 )
70 self._ccfg = original_ccfg
71 self._ccfg.parser.ccfg = self._ccfg
72 self._overhead_obj._ccfg = self._ccfg
73 insights, _ = res
74 stage_mems = [i["Static"] + i["Dynamic"] for i in insights]
75 peak_mem = max(stage_mems)
76 peak_stage = stage_mems.index(peak_mem)
78 logger.output(
79 "model_name: %s, peak memory : \033[1m%s MB\033[0m (stage _%s)",
80 self._ccfg.model_name,
81 peak_mem,
82 peak_stage,
83 )
84 return peak_mem
86 def estimate_peak_insight(self, stages: list = None) -> list:
87 """subcomponents' proportion estimation"""
88 original_ccfg = copy.deepcopy(self._ccfg)
89 insights, _ = self._estimate_backbone(stages, False, False, -1, False)
90 self._ccfg = original_ccfg
91 self._ccfg.parser.ccfg = self._ccfg
92 self._overhead_obj._ccfg = self._ccfg
93 return insights
95 def estimate_layer_memory(
96 self, stages: list = None, ppb_format=1, device_type=Hard.Device_A2
97 ) -> Dict:
98 """PPB's input"""
99 logger.info(device_type)
100 if self.ppb:
101 return self.ppb
102 original_ccfg = copy.deepcopy(self._ccfg)
103 res = self._estimate_backbone(stages, False, ppb_format, -1, False)
104 self._ccfg = original_ccfg
105 self._ccfg.parser.ccfg = self._ccfg
106 self._overhead_obj._ccfg = self._ccfg
107 _, ppb = res
108 self.ppb = ppb
109 return ppb
111 # Specific estimation
113 def static_mem_stage(self, stage_id: int) -> float:
114 """stage estimation proportions"""
115 insights = self.estimate_peak_insight()
116 return insights[stage_id]["Static"]
118 def dynamic_mem_stage(self, stage_id: int) -> float:
119 """stage estimation proportions"""
120 insights = self.estimate_peak_insight()
121 return insights[stage_id]["Dynamic"]
123 def logs_mem_stage(self, stage_id: int) -> list:
124 """stage estimation hook calls trace"""
125 insights = self.estimate_peak_insight()
126 return insights[stage_id]["Node Log"]
128 def static_mem_layer(
129 self, node: Union[str, LayerType], stage_id: int
130 ) -> Union[float, list]:
131 """accounting all hooks for this node"""
132 logs = self.logs_mem_stage(stage_id)
133 extracted_log = set()
134 for k in logs.keys():
135 if k[3] == node.name[0]:
136 extracted_log.add(logs[k]["_param"])
137 extracted_log = list(extracted_log)
138 if len(extracted_log) == 1:
139 return extracted_log[0]
140 return extracted_log
142 def dynamic_mem_layer(
143 self, node: Union[str, LayerType], stage_id: int
144 ) -> Union[float, list]:
145 """accounting all hooks for this node"""
146 logs = self.logs_mem_stage(stage_id)
147 extracted_log = set()
148 for k in logs.keys():
149 if k[3] == node.name[0]:
150 extracted_log.add(logs[k]["_activ"] + logs[k]["_comm"])
151 extracted_log = list(extracted_log)
152 if len(extracted_log) == 1:
153 return extracted_log[0]
154 return extracted_log
156 def mem_fit(
157 self, mem: float, tolerance: float = 0, margin: float = 0
158 ) -> bool:
159 """check if input memory fits in device"""
160 # Expect arguments to be in MB
161 memory = Memory.from_mb(mem)
162 tolerance_mem = Memory.from_mb(tolerance)
163 margin_mem = Memory.from_mb(margin)
164 cap = self._ccfg.device_capacity - margin_mem
165 diff = abs(memory - cap)
166 is_close = diff <= tolerance_mem
167 is_fit = memory <= cap
169 if Memory.zero() < tolerance_mem and is_close:
170 logger.info(
171 "Prediction is CLOSE to memory device (%s of diff)", diff
172 )
173 return True
174 if is_fit:
175 logger.output(
176 "estimation FITS into device memory (%s<=%s-%s)",
177 mem,
178 self._ccfg.device_capacity,
179 margin_mem,
180 )
181 else:
182 logger.output(
183 "estimation DOES NOT FIT into device memory (%s>%s-%s)",
184 mem,
185 self._ccfg.device_capacity,
186 margin_mem,
187 )
188 return is_fit
191def estimate_memory(config: Any) -> bool:
192 """fast usage for estimation/fit for given input config"""
193 e = EvaluatorV2(config)
194 peak = e.estimate_peak()
195 return e.mem_fit(peak)
198def main():
199 """commandline"""
200 parser = argparse.ArgumentParser(
201 description="Command line usage: Estimate peak stage memory"
202 )
204 parser.add_argument(
205 "model_config_path",
206 nargs=1,
207 help="Model config file (MindFormer YAML or MindSpeed JSON)",
208 )
210 parser.add_argument(
211 "--framework",
212 default=None,
213 type=str,
214 help="Specify a framework name",
215 )
217 parser.add_argument(
218 "--code-path",
219 default=None,
220 type=str,
221 help="Specify a source code path (Additional parsing)",
222 )
224 parser.add_argument(
225 "--verbose", action="store_true", help="Show estimation trace"
226 )
227 parser.add_argument("--plot", action="store_true", help="Plot estimation")
228 parser.add_argument(
229 "--fit",
230 action="store_true",
231 help="Check if estimation fits in device memory",
232 )
233 parser.add_argument(
234 "--stage", default=-1, type=int, help="Specify pipeline stage ID"
235 )
236 parser.add_argument(
237 "--hook",
238 default=None,
239 type=str,
240 help="Specify hook class (defined in hooks/)",
241 )
242 parser.add_argument(
243 "--trace-fun",
244 default=None,
245 type=str,
246 help="Specify a formula function name to get it traced",
247 )
248 parser.add_argument(
249 "--ppb",
250 action="store_true",
251 help="Generate pipeline balancing layers description",
252 )
253 parser.add_argument(
254 "--ppb-new",
255 action="store_true",
256 help="Generate pipeline balancing layers description (New format)",
257 )
258 parser.add_argument(
259 "--ctx", action="store_true", help="Show ctx variables"
260 )
261 parser.add_argument(
262 "--ccfg", action="store_true", help="Show ccfg variables"
263 )
264 parser.add_argument(
265 "--warnings", action="store_true", help="Show warnings"
266 )
267 args = parser.parse_args()
269 path = args.model_config_path[0]
270 if not os.path.exists(path):
271 raise argparse.ArgumentTypeError(f"`{path}` was not found")
272 if not path.endswith((".yaml", ".json",".toml")):
273 raise argparse.ArgumentTypeError(f"`{path}` has invalid file type")
275 e = EvaluatorV2(
276 path,
277 framework=args.framework,
278 source_code=args.code_path,
279 log_level=args.warnings,
280 hook_cls=args.hook,
281 trace_fun=args.trace_fun,
282 )
283 if args.ctx:
284 e.print_ctx()
285 if args.ccfg:
286 e.print_ccfg()
287 if args.ppb:
288 ppb = e.estimate_layer_memory()
289 print(json.dumps(ppb, indent=2))
290 elif args.ppb_new:
291 ppb = e.estimate_layer_memory(ppb_format=2)
292 print(json.dumps(ppb, indent=2))
293 else:
294 peak_mem = e.estimate_peak(
295 verbose=args.verbose, spec_stage_id=args.stage, plot=args.plot
296 )
298 if args.fit:
299 e.mem_fit(peak_mem)
302if __name__ == "__main__":
303 main()