Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / memory_estimation / _backbone.py: 60%

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"""backbone memory estimation module""" 

16from __future__ import annotations 

17from typing import TYPE_CHECKING 

18import os 

19import ast 

20import math 

21import logging 

22import inspect 

23import pprint 

24import importlib 

25import matplotlib.pyplot as plt 

26from PIL import Image 

27from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config 

28from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger as nd_logger 

29from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig 

30from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger 

31from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context, MemType 

32from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.utils import EvalUtils 

33from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._bwd_overhead import _BackwardOverhead 

34from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._ppb import _PPB 

35from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.hook_base import MemEvalHook 

36 

37if TYPE_CHECKING: 

38 from typing import Any, Dict, Tuple 

39 

40current_dir = os.path.dirname(os.path.abspath(__file__)) 

41EVAL_YML = os.path.join(current_dir, "configs_eval/default.yaml") 

42 

43 

44class _Backbone: 

45 """backbone class""" 

46 

47 def __init__(self, config: Any, **kwargs): 

48 self._child_cls = self 

49 self.mb = EvalUtils.mb 

50 self.eval_cfg = Config(kwargs.get("eval_yml", EVAL_YML)) 

51 self._ctx = kwargs.get("ctx", Context()) 

52 self._ccfg = kwargs.get("ccfg", None) 

53 self.framework = kwargs.get("framework", None) 

54 self.source_code = kwargs.get("source_code", None) 

55 self.hook_cls, self.config_path = None, None 

56 if not self._ccfg: 

57 self.hook_cls = kwargs.get("hook_cls", self.eval_cfg.hook_class) 

58 if isinstance(self.hook_cls, str): 

59 self._load_eval_yaml_hook_cls(self.hook_cls) 

60 if self.hook_cls and not isinstance(self.hook_cls, MemEvalHook): 

61 raise AttributeError( 

62 f"'{self.hook_cls}' is not a MemEvalHook instance" 

63 ) 

64 if config: 

65 if kwargs.get("log_level", 1) == 0: 

66 # logger.setLevel(logging.CRITICAL) 

67 nd_logger.setLevel(logging.CRITICAL) 

68 self.update_config(config) 

69 else: 

70 raise AttributeError("missing config") 

71 self.evaluator_instances = None 

72 self.ppb = None 

73 self._overhead_obj = _BackwardOverhead( 

74 self, self._ccfg, self._ctx, self._inner_dynamic_mem 

75 ) 

76 self._ppb_obj = _PPB(self.eval_cfg, self._inner_dynamic_mem) 

77 

78 @property 

79 def ccfg(self) -> CostModelConfig: 

80 """read-only""" 

81 return self._ccfg 

82 

83 @property 

84 def ctx(self) -> Context: 

85 """read-only""" 

86 return self._ctx 

87 

88 def _load_eval_yaml_hook_cls(self, hook_cls): 

89 """hook_class in eval yaml""" 

90 target_mod_path = None 

91 try: 

92 # search in folder 'hooks' 

93 hooks_dir = os.path.join(current_dir, "hooks") 

94 for f in os.listdir(hooks_dir): 

95 if f.endswith(".py"): 

96 mod_path = f"hyper_parallel.auto_parallel.sapp_nd.memory_estimation.hooks.{f.split('.')[0]}" 

97 spec = importlib.util.find_spec(mod_path) 

98 if spec is None or spec.origin is None: 

99 continue 

100 with open(spec.origin, "r", encoding="utf-8") as mf: 

101 source = mf.read() 

102 tree = ast.parse(source) 

103 mod_cls = None 

104 for node in ast.walk(tree): 

105 if isinstance(node, ast.ClassDef) and node.name == hook_cls: 

106 mod_cls = node 

107 break 

108 if mod_cls: 

109 target_mod_path = mod_path 

110 break 

111 if target_mod_path: 

112 module = importlib.import_module(target_mod_path) 

113 self.hook_cls = getattr(module, hook_cls)() 

114 except (ModuleNotFoundError, ImportError) as e: 

115 print(e) 

116 

117 # Peak memory estimation 

118 

119 def update_config(self, new_config: Any) -> None: 

120 """processing input config""" 

121 if self._ccfg is not None: 

122 self._ccfg.update_config(new_config, self.hook_cls, self.framework, self.source_code) 

123 else: 

124 self._ccfg = CostModelConfig(new_config, self.hook_cls, self.framework, self.source_code) 

125 if isinstance(new_config, str): 

126 if not self.config_path: 

127 logger.info( 

128 "%s Process config file: %s", 

129 "=" * 30, 

130 new_config.split("/")[-1], 

131 ) 

132 self.config_path = new_config 

133 self.evaluator_instances = None 

134 self.ppb = None 

135 

136 def _inner_static_mem(self) -> float: 

137 """static memory evaluation for backbone estimation""" 

138 if self._ctx.current_node in self._ctx.node_eval: 

139 p = self._ctx.eval.stat.p(self._ccfg, self._ctx) 

140 ost = self._ctx.eval.stat.os(self._ccfg, self._ctx) 

141 grad = self._ctx.eval.stat.grad(self._ccfg, self._ctx) 

142 res = p + ost + grad 

143 self._ctx.save2log("_param", res) 

144 # Log routed/shared expert param breakdown for MoE layers 

145 if self._ccfg.n_exp > 1 and self.is_regular_layer(self._ctx.current_node): 

146 result = self._ctx.eval.num_p(self._ccfg, self._ctx) 

147 if isinstance(result, tuple) and len(result) == 3: 

148 _, routed_exp_p, shared_exp_p = result 

149 routed_exp_mem = ( 

150 routed_exp_p / self._ccfg.ep 

151 * self._ccfg.bytes_p / self._ccfg.shard_p_os_exp 

152 ) 

153 shared_exp_mem = ( 

154 shared_exp_p 

155 * self._ccfg.bytes_p / self._ccfg.shard_p_os_exp_partial 

156 ) 

157 self._ctx.save2log("routed_exp_param", routed_exp_mem) 

158 self._ctx.save2log("shared_exp_param", shared_exp_mem) 

159 return p + ost + grad 

160 return 0 

161 

162 def _inner_dynamic_mem( 

163 self, ppb=False, default_micro_factor=None 

164 ) -> tuple[float, float]: 

165 """dynamic memory evaluation for backbone estimation""" 

166 if self._ctx.current_node in self._ctx.node_eval: 

167 if ppb: 

168 micro_factor = 1 

169 elif default_micro_factor: 

170 micro_factor = default_micro_factor 

171 else: 

172 sched = self._ccfg.pp_sched 

173 micro_factor = self._ctx.pp_micro_eval[sched]( 

174 self._ccfg, self._ctx 

175 ) 

176 self._ctx.micro_factor = max(1, micro_factor) 

177 activation = self._ctx.eval.dyn.activation(self._ccfg, self._ctx) 

178 self._ctx.save2log("_activ", activation) 

179 comm = self._inner_comm_mem(micro_factor) 

180 return activation, comm 

181 return 0, 0 

182 

183 def _inner_comm_mem(self, micro_factor) -> float: 

184 """dynamic memory evaluation for backbone estimation""" 

185 if self._ctx.current_node in self._ctx.node_eval: 

186 comm_eval_field = self._ctx.eval.dyn.comm 

187 comm_cat = { 

188 "dp": MemType.AG_COMM, 

189 "tp": MemType.AG_COMM, 

190 "cp": MemType.AG_COMM, 

191 "ep": MemType.A2A_COMM, 

192 } 

193 comm_mem = {} 

194 for k, fun in vars(comm_eval_field).items(): 

195 if fun is None: 

196 continue 

197 sig_param = inspect.signature(fun).parameters.values() 

198 if ( 

199 # any( 

200 # p.kind == inspect.Parameter.VAR_KEYWORD 

201 # for p in sig_param 

202 # ) 

203 # and len(sig_param) > 2 

204 len(sig_param) > 2 

205 ): 

206 comm = fun(self._ccfg, self._ctx, micro_factor) 

207 else: 

208 comm = fun(self._ccfg, self._ctx) 

209 comm_mem[k] = comm 

210 res = EvalUtils.eval_expr_insight( 

211 expr=self._ctx.comm_expr, 

212 ctx=self._ctx, 

213 mem_val=comm_mem, 

214 mem_cat=comm_cat, 

215 ) 

216 self._ctx.save2log("_comm", res) 

217 return res 

218 return 0 

219 

220 def __subplot_stages(self, stage_i, stat_mem_i, dyn_mem_i, i) -> None: 

221 """Plot static and dynamic memory for one pipeline stage.""" 

222 _, ax = plt.subplots(figsize=(5, 8)) 

223 bottoms = {"stat": 0, "dyn": 0} 

224 color_stat = plt.get_cmap("cividis") 

225 color_dyn = plt.get_cmap("plasma") 

226 total_lay = sum(len(chunk) for chunk in stage_i) 

227 color_denominator = max(1, total_lay - 1) 

228 for chunk_id, chunk in enumerate(stage_i): 

229 for lay_id, lay_type in enumerate(chunk): 

230 real_id = self._ctx.real_lay_ids[chunk_id][i][lay_id] 

231 name = f"Lay_{real_id}_{lay_type.name[0]}" 

232 idx = chunk_id * len(dyn_mem_i) + lay_id 

233 stat = self.mb(stat_mem_i[chunk_id][lay_id]) 

234 dyn = self.mb(dyn_mem_i[chunk_id][lay_id]) 

235 ax.bar( 

236 0.1, 

237 [dyn], 

238 bottom=bottoms["dyn"], 

239 label=name, 

240 color=color_dyn(idx / color_denominator), 

241 width=0.15, 

242 linewidth=0.5, 

243 edgecolor="black", 

244 ) 

245 ax.text( 

246 0.2, 

247 bottoms["dyn"] + dyn / 2, 

248 f"DYN_{name}", 

249 va="center", 

250 fontsize=5, 

251 ) 

252 ax.bar( 

253 0.5, 

254 [stat], 

255 bottom=bottoms["stat"], 

256 label=name, 

257 color=color_stat(idx / color_denominator), 

258 width=0.15, 

259 linewidth=0.5, 

260 edgecolor="black", 

261 ) 

262 ax.text( 

263 0.6, 

264 bottoms["stat"] + stat / 2, 

265 f"STAT_{name}", 

266 va="center", 

267 fontsize=5, 

268 ) 

269 bottoms["dyn"] += dyn 

270 bottoms["stat"] += stat 

271 ax.axhline( 

272 self.get_max_device_memory(), 

273 color="red", 

274 linewidth=1, 

275 ls="dotted", 

276 ) 

277 ax.axhline( 

278 bottoms["dyn"] + bottoms["stat"], 

279 color="blue", 

280 linewidth=1, 

281 ls="dashed", 

282 ) 

283 ax.text( 

284 0, 

285 self.get_max_device_memory(), 

286 "Device Memory", 

287 color="red", 

288 va="bottom", 

289 fontsize=8, 

290 ) 

291 ax.text( 

292 0.62, 

293 bottoms["dyn"] + bottoms["stat"], 

294 "Prediction total", 

295 color="blue", 

296 va="bottom", 

297 fontsize=8, 

298 ) 

299 ax.set_xticks([]) 

300 ax.set_ylabel("Size (MB)") 

301 ax.set_xlim([0, 0.8]) 

302 ax.set_xlabel(f"Stage_{i}") 

303 

304 def __plot_stages(self, stages, stat_mems, dyn_mems) -> None: 

305 """plot bars for estimations""" 

306 

307 logger.info("Plotting predictions in plots/") 

308 if not os.path.exists("plots"): 

309 os.makedirs("plots") 

310 imgs = [] 

311 for stage_id in range(self._ccfg.p): 

312 self.__subplot_stages( 

313 stages[stage_id], 

314 stat_mems[stage_id], 

315 dyn_mems[stage_id], 

316 stage_id, 

317 ) 

318 img = f"plots/MemPlot_stage_{stage_id}.png" 

319 imgs += [(stage_id, img)] 

320 logger.info("save plot: %s", img) 

321 plt.savefig(img, dpi=300, bbox_inches="tight") 

322 plt.clf() 

323 plt.close() 

324 # Concatenate every stage plots 

325 if imgs: 

326 imgs = sorted(imgs) 

327 canvas = [Image.open(i) for _, i in imgs] 

328 stage_canva = Image.new( 

329 "RGB", 

330 ( 

331 canvas[0].size[0] 

332 * 2 ** math.ceil(math.log2(len(canvas)) / 2), 

333 canvas[0].size[1] 

334 * 2 ** math.floor(math.log2(len(canvas)) / 2), 

335 ), 

336 ) 

337 offset_x, offset_y = 0, 0 

338 for i in canvas: 

339 stage_canva.paste(i, (offset_x, offset_y)) 

340 offset_x += i.size[0] 

341 if offset_x >= stage_canva.size[0]: 

342 offset_x = 0 

343 offset_y += i.size[1] 

344 stage_canva.save("plots/MemPlot_all_stages.png") 

345 logger.info("save stage plot: plots/MemPlot_all_stages.png") 

346 

347 def __update_stage_logs(self, stage_logs: list, stage_id: int) -> None: 

348 """update stage's temporary buffers from ctx's buffers""" 

349 if not stage_logs[stage_id].node_compute_log: 

350 stage_logs[stage_id].node_compute_log = {} 

351 stage_logs[stage_id].node_compute_log.update( 

352 self._ctx.node_compute_log 

353 ) 

354 if not stage_logs[stage_id].accu_mem_type: 

355 stage_logs[stage_id].accu_mem_type = { 

356 mt: 0 for mt in list(MemType) 

357 } 

358 for mem_type in list(MemType): 

359 val = self._ctx.accu_mem_type[mem_type] 

360 stage_logs[stage_id].accu_mem_type[mem_type] += val 

361 

362 def __preprocess_layer_custom_config_list(self, stages: list) -> list: 

363 """flatten layer_custom_config for backbone estimation""" 

364 flatten = sum( 

365 [[f[1]] * f[0] for f in self._ccfg.layer_custom_config], [] 

366 ) 

367 total_n_lay = self._ccfg.n_lay + self._ccfg.n_mtp 

368 total_n_lay_stages = self._ccfg.count_layers(stages) 

369 if not self._ccfg.multimodal and total_n_lay != total_n_lay_stages: 

370 raise(AttributeError( 

371 f"Mismatch of num_layers between parsed value ({total_n_lay})" 

372 f" and generated partitions ({total_n_lay_stages})" 

373 f" => offset may be incorrect ({self._ccfg.offset})" 

374 )) 

375 if self._ccfg.n_lay > 0 and len(flatten) != total_n_lay: 

376 raise AttributeError( 

377 f"layer_custom_config occurrences ({len(flatten)})" 

378 f" != num_layers ({total_n_lay})" 

379 ) 

380 if self._ccfg.pp_sched == "zero_bubble_v": 

381 n_layer_first_chunk = ( 

382 sum(len(s[0]) for s in stages) - 1 

383 ) # Except embedding layer 

384 flatten = ( 

385 flatten[:n_layer_first_chunk] 

386 + flatten[n_layer_first_chunk:][::-1] 

387 ) 

388 return flatten 

389 

390 def _estimate_backbone(self, *args) -> Tuple[list, Dict]: 

391 """Evaluator's main function for estimation""" 

392 stages = args[0] 

393 spec_stage_id = args[3] 

394 

395 if spec_stage_id >= self._ccfg.p or spec_stage_id < 0: 

396 spec_stage_id = -1 

397 # Process partition generation 

398 if not stages: 

399 stages = self._ccfg.generate_partitions_vpp() 

400 # multimodal=self._ccfg.multimodal 

401 # ) 

402 if not self._ccfg.multimodal: 

403 return self.__estimate_stages_backbone( 

404 stages, args[1], args[2], spec_stage_id, args[4] 

405 ) 

406 res = [] 

407 original_ccfg = self._ccfg 

408 common_lc = [] 

409 self.evaluator_instances = [] 

410 # Build common layer_custom_config + Build temporary evaluators 

411 for m in self._ccfg.mm_order: 

412 self._ccfg.mm_ccfgs[m].config = original_ccfg.config 

413 if not self._child_cls: 

414 raise AttributeError("expected non null _child_cls") 

415 tmp_evaluator = type(self._child_cls)( 

416 None, 

417 ccfg=self._ccfg.mm_ccfgs[m], 

418 trace_fun=self.toggle_func_trace 

419 ) 

420 tmp_evaluator.import_eval_yaml() 

421 num_layer = tmp_evaluator.get_num_layers() 

422 # assert not isinstance(num_layer, tuple) 

423 strategy = tmp_evaluator.get_strategy() 

424 full_rec = strategy["full_rec"] 

425 offset = strategy["offset"] 

426 self._ccfg.hooks_dict[m](tmp_evaluator) 

427 strategy = tmp_evaluator.get_strategy() 

428 if ( 

429 tmp_evaluator.get_num_layers() != num_layer 

430 or strategy["full_rec"] != full_rec 

431 or strategy["offset"] != offset 

432 ): 

433 # Reverify num layers, recomp, offset after hook 

434 stages[m] = ( 

435 tmp_evaluator.ccfg.generate_partitions_vpp_unimodal() 

436 ) 

437 tmp_evaluator.set_layer_custom(None) 

438 common_lc += self._ccfg.mm_ccfgs[m].layer_custom_config 

439 self.evaluator_instances += [tmp_evaluator] 

440 if args[1]: 

441 logger.info("Submodule %s", self._ccfg.mm_ccfgs[m].model_name) 

442 tmp_evaluator.print_ctx() 

443 self.print_stages(stages[m]) 

444 self.set_layer_custom(common_lc) 

445 if args[1]: 

446 logger.info( 

447 "Combined layer_custom_config for %s\n%s", 

448 self._ccfg.model_name, 

449 pprint.pformat(self._ccfg.layer_custom_config, compact=True), 

450 ) 

451 logger.info( 

452 "Sub evaluator instances for %s\n%s", 

453 self._ccfg.model_name, 

454 pprint.pformat(self.evaluator_instances, compact=True), 

455 ) 

456 

457 res = self.__estimate_stages_backbone( 

458 self._ccfg.combine_partition_multimodal(stages), 

459 args[1], 

460 args[2], 

461 spec_stage_id, 

462 args[4], 

463 ) 

464 self._ccfg = original_ccfg 

465 return res 

466 

467 def __estimate_stages_backbone(self, *args) -> Tuple[list, Dict]: 

468 """Evaluator's main function for stage estimation""" 

469 stages = args[0] 

470 verbose = args[1] 

471 compute_ppb = args[2] 

472 spec_stage_id = args[3] 

473 

474 if verbose: 

475 logger.info("Partition of layers :") 

476 self._ccfg.print_stages(stages, spec_stage_id) 

477 insights = [] 

478 # Compute peak memory 

479 flatten = self.__preprocess_layer_custom_config_list(stages) 

480 if verbose: 

481 logger.info( 

482 "Flatten layer_custom_config\n%s", 

483 pprint.pformat( 

484 [f if not f else f.__name__ for f in flatten], compact=True 

485 ), 

486 ) 

487 

488 stage_misc = { 

489 "stat": [[[0 for _ in c] for c in s] for s in stages], 

490 "dyn": [[[0 for _ in c] for c in s] for s in stages], 

491 "logs": [Config({}) for _ in range(self._ccfg.p)], 

492 } 

493 # tmp_ppb_lay_desc = [] # PPB purpose 

494 ppb_lay_desc = [] 

495 record_lay_types = {} 

496 self.__chunk_stage_lay_loops( 

497 flatten, 

498 stages, 

499 record_lay_types, 

500 stage_misc, 

501 verbose, 

502 compute_ppb, 

503 ppb_lay_desc, # tmp_ppb_lay_desc, 

504 ) 

505 self.__postprocess_stages( 

506 stages, 

507 record_lay_types, 

508 stage_misc, 

509 verbose, 

510 spec_stage_id, 

511 insights, 

512 ) 

513 # PPB Input 

514 ppb_input = None 

515 if compute_ppb == 1: 

516 self._ppb_obj.ppb_combine_bodies(ppb_lay_desc) 

517 ppb_input = {"layers_description": ppb_lay_desc} 

518 elif compute_ppb == 2: 

519 self._ppb_obj.ppb_combine_bodies_new(ppb_lay_desc) 

520 ppb_input = {"layers_description_new": ppb_lay_desc} 

521 if args[4]: # Plot 

522 self.__plot_stages(stages, stage_misc["stat"], stage_misc["dyn"]) 

523 return insights, ppb_input 

524 

525 def __update_evaluator(self, node, verbose): 

526 if self.evaluator_instances and node == self._ctx.head_node: 

527 tmp_eval = self.evaluator_instances.pop(0) 

528 self._ccfg = tmp_eval.ccfg 

529 self._ctx.copy_tmp_buff(tmp_eval.ctx) 

530 self._ctx = tmp_eval.ctx 

531 if verbose: 

532 logger.info( 

533 "Update ccfg and ctx, module : %s", 

534 self._ccfg.model_name, 

535 ) 

536 

537 def __update_next_layer_custom_function(self, *args): 

538 """Apply the next layer custom hook before evaluating a layer.""" 

539 flatten, verbose = args[0], args[1] 

540 record_lay_types = args[2] 

541 stage_id, chunk_id, lay_id = args[3], args[4], args[5] 

542 node = args[6] 

543 if self.is_regular_layer(node) and flatten: 

544 hook = flatten.pop(0) 

545 if hook: 

546 if verbose: 

547 logger.info("Apply hook %s", hook.__name__) 

548 record_lay_types[(stage_id, chunk_id, lay_id)] = ( 

549 self._ccfg, 

550 self._ctx, 

551 hook, 

552 ) 

553 hook(self) 

554 else: 

555 record_lay_types[(stage_id, chunk_id, lay_id)] = ( 

556 self._ccfg, 

557 self._ctx, 

558 lambda _: None, 

559 ) 

560 else: 

561 record_lay_types[(stage_id, chunk_id, lay_id)] = ( 

562 self._ccfg, 

563 self._ctx, 

564 lambda _: None, 

565 ) 

566 if verbose: 

567 logger.info( 

568 "stage_id=%s, chunk_id=%s, lay_id=%s, node=%s", 

569 stage_id, 

570 chunk_id, 

571 lay_id, 

572 node, 

573 ) 

574 self._ccfg.print_parallelism() 

575 

576 def __chunk_stage_lay_loops(self, *args): 

577 """Evaluate every layer in every stage and chunk.""" 

578 flatten, stages, record_lay_types = args[0], args[1], args[2] 

579 sm = args[3] 

580 verbose, compute_ppb, ppb_lay_desc = args[4], args[5], args[6] 

581 self._ctx.real_lay_ids = [] 

582 count = 0 

583 for chunk_id in range(self._ccfg.vp): 

584 self._ctx.real_lay_ids += [[]] 

585 for stage_id in range(self._ccfg.p): 

586 self._ctx.real_lay_ids[chunk_id] += [[]] 

587 self._ctx.init_tmp_buff() 

588 for lay_id in range(len(stages[stage_id][chunk_id])): 

589 node = stages[stage_id][chunk_id][lay_id] 

590 if self.is_regular_layer(node): 

591 self._ctx.real_lay_ids[chunk_id][stage_id] += [count] 

592 count += 1 

593 else: 

594 self._ctx.real_lay_ids[chunk_id][stage_id] += [""] 

595 # Update evaluator (multimodal) 

596 self.__update_evaluator(node, verbose) 

597 # Update next layer custom function 

598 self.__update_next_layer_custom_function( 

599 flatten, 

600 verbose, 

601 record_lay_types, 

602 stage_id, 

603 chunk_id, 

604 lay_id, 

605 node, 

606 ) 

607 self._ctx.current_stage_id = stage_id 

608 self._ctx.current_chunk_id = chunk_id 

609 self._ctx.current_lay_id = lay_id 

610 self._ctx.current_node = node 

611 static_mem = self._inner_static_mem() 

612 sm["stat"][stage_id][chunk_id][lay_id] = static_mem 

613 sm["dyn"][stage_id][chunk_id][lay_id] = sum( 

614 self._inner_dynamic_mem() 

615 ) 

616 if verbose: 

617 logger.info("pp micro factor for dynamic: %s",self._ctx.micro_factor) 

618 # PPB Purpose 

619 if compute_ppb == 1: 

620 desc = self._ppb_obj.lay_ppb( 

621 self._ccfg, 

622 self._ctx, 

623 sm["stat"][stage_id][chunk_id][lay_id], 

624 ) 

625 self._ppb_obj.add_to_ppb_list(ppb_lay_desc, desc) 

626 elif compute_ppb == 2: 

627 desc = self._ppb_obj.lay_ppb_new( 

628 self._ccfg, 

629 self._ctx, 

630 sm["stat"][stage_id][chunk_id][lay_id], 

631 ) 

632 self._ppb_obj.add_to_ppb_list(ppb_lay_desc, desc) 

633 self.__update_stage_logs(sm["logs"], stage_id) 

634 

635 def __postprocess_stages(self, *args): 

636 """Build memory insights from raw stage evaluation buffers.""" 

637 stages, record_lay_types = args[0], args[1] 

638 sm = args[2] 

639 verbose, spec_stage_id = args[3], args[4] 

640 insights = args[5] 

641 for stage_id in range(self._ccfg.p): 

642 ins = {} # Mem Insights purpose 

643 ins["Static"] = sum( 

644 sum(mem for mem in c) for c in sm["stat"][stage_id] 

645 ) 

646 ins["Dynamic"] = sum( 

647 sum(mem for mem in c) for c in sm["dyn"][stage_id] 

648 ) 

649 self._ctx.init_tmp_buff() 

650 if not self._ccfg.freeze: 

651 ins["Dynamic"] += self._overhead_obj.estimate( 

652 stages, stage_id, record_lay_types 

653 ) 

654 self.__update_stage_logs(sm["logs"], stage_id) 

655 safety_buffer = 1024 * 1024 * 1024 # 1 GB 

656 if ins["Dynamic"] > 0: 

657 ins["Dynamic"] += safety_buffer 

658 stage_accu = sm["logs"][stage_id].accu_mem_type 

659 ins["ModelParameters"] = self.mb(stage_accu[MemType.MODEL_PARAM]) 

660 ins["OptimizerStates"] = self.mb(stage_accu[MemType.OPTIM_STATE]) 

661 ins["AccumulGradients"] = self.mb(stage_accu[MemType.ACCU_GRAD]) 

662 ins["Attn"] = self.mb(stage_accu[MemType.ATTN_ACTIV]) 

663 ins["FFn"] = self.mb(stage_accu[MemType.FFN_ACTIV]) 

664 ins["Norm"] = self.mb(stage_accu[MemType.NORM_ACTIV]) 

665 ins["AllGather Comm"] = self.mb(stage_accu[MemType.AG_COMM]) 

666 ins["All2All Comm"] = self.mb(stage_accu[MemType.A2A_COMM]) 

667 ins["Node Log"] = sm["logs"][stage_id].node_compute_log 

668 # VERBOSE 

669 if verbose and spec_stage_id in (-1, stage_id): 

670 self.__verbose_insights(sm, stage_id, ins) 

671 ins["Static"] = self.mb(ins["Static"]) 

672 ins["Dynamic"] = self.mb(ins["Dynamic"]) 

673 insights += [ins] 

674 

675 def __verbose_insights(self, *args): 

676 """logging purpose""" 

677 sm = args[0] 

678 stage_id = args[1] 

679 ins = args[2] 

680 stat_i = max(1, ins["Static"]) 

681 dyn_i = max(1, ins["Dynamic"]) 

682 # logs_i = sm["logs"][stage_id] 

683 accu_i = sm["logs"][stage_id].accu_mem_type 

684 logger.info( 

685 "stage _%s : %s MB", 

686 stage_id, 

687 self.mb(ins["Static"] + ins["Dynamic"]), 

688 ) 

689 logger.info( 

690 "\tStatic\t%s\t" 

691 "ModelParam %s (%s%%), " 

692 "OptimStates %s (%s%%), " 

693 "Gradients %s (%s%%)", 

694 self.mb(ins["Static"]), 

695 ins["ModelParameters"], 

696 round(accu_i[MemType.MODEL_PARAM] / stat_i * 100), 

697 ins["OptimizerStates"], 

698 round(accu_i[MemType.OPTIM_STATE] / stat_i * 100), 

699 ins["AccumulGradients"], 

700 round(accu_i[MemType.ACCU_GRAD] / stat_i * 100), 

701 ) 

702 logger.info( 

703 "\tDynamic\t%s\t" 

704 "Attn %s (%d%%), " 

705 "FFn %s (%d%%), " 

706 "Norm %s (%d%%), " 

707 "AllGather Comm %s (%d%%), " 

708 "All2All Comm %s (%d%%), ", 

709 self.mb(ins["Dynamic"]), 

710 ins["Attn"], 

711 round(accu_i[MemType.ATTN_ACTIV] / dyn_i * 100), 

712 ins["FFn"], 

713 round(accu_i[MemType.FFN_ACTIV] / dyn_i * 100), 

714 ins["Norm"], 

715 round(accu_i[MemType.NORM_ACTIV] / dyn_i * 100), 

716 ins["AllGather Comm"], 

717 round(accu_i[MemType.AG_COMM] / dyn_i * 100), 

718 ins["All2All Comm"], 

719 round(accu_i[MemType.A2A_COMM] / dyn_i * 100), 

720 )