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

1# Copyright 2025-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"""backbone memory estimation module""" 

16from __future__ import annotations 

17from typing import TYPE_CHECKING 

18import os 

19import ast 

20import math 

21import logging 

22import inspect 

23import functools 

24import pprint 

25import importlib 

26import matplotlib.pyplot as plt 

27from PIL import Image 

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

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

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

31from hyper_parallel.auto_parallel.sapp_nd.nd.common.arch_hooks import ( 

32 apply_layer_kind, 

33 check_and_apply_custom_hook, 

34 layer_kinds, 

35) 

36from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_order import layer_switches 

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

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

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

40from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.body import EvalBody 

41from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.tail import EvalTailSingle 

42from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType 

43from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.comm import EvalLayerComm 

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

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

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

47 

48if TYPE_CHECKING: 

49 from typing import Any, Dict, List, Optional, Tuple 

50 

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

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

53 

54 

55@functools.lru_cache(maxsize=4096) 

56def _parameters(fun: Any) -> Tuple[inspect.Parameter, ...]: 

57 """A communication formula's parameters, read once a formula rather than once a layer and candidate.""" 

58 return tuple(inspect.signature(fun).parameters.values()) 

59 

60 

61class _Backbone: 

62 """backbone class""" 

63 

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

65 self._child_cls = self 

66 self.mb = EvalUtils.mb 

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

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

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

70 # During a walk, the regular layers' own recompute switches still to 

71 # visit, where the config states several (derive's layer_switches). 

72 self._switch_order = None 

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

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

75 self.hook_cls, self.config_path = None, None 

76 # The machine's device count, as run_nd -d states it, for a config 

77 # that cannot state its own: an AutoModels yaml has no world size. 

78 self.devices = getattr(kwargs.get("machine"), "number", None) or 0 

79 if not self._ccfg: 

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

81 if isinstance(self.hook_cls, str): 

82 self._load_eval_yaml_hook_cls(self.hook_cls) 

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

84 raise AttributeError( 

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

86 ) 

87 if config: 

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

89 # logger.setLevel(logging.CRITICAL) 

90 nd_logger.setLevel(logging.CRITICAL) 

91 self.update_config(config) 

92 else: 

93 raise AttributeError("missing config") 

94 self.evaluator_instances = None 

95 # The submodule a multimodal model's stages are fitted on, by name, 

96 # or None for a model of one module. 

97 self.main_module = None 

98 self.ppb = None 

99 # Each stage's dynamic memory, in MB, at the two points of the last 

100 # stage estimate where its peak can come: as warm-up ends, and as a 

101 # later micro-batch's backward ends (0 where that is no candidate). 

102 # The stage's Dynamic insight is the larger. 

103 self.peak_points: List[Tuple[int, int]] = [] 

104 self._overhead_obj = _BackwardOverhead( 

105 self, self._ccfg, self._ctx, self._inner_dynamic_mem 

106 ) 

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

108 

109 @property 

110 def ccfg(self) -> CostModelConfig: 

111 """read-only""" 

112 return self._ccfg 

113 

114 @property 

115 def ctx(self) -> Context: 

116 """read-only""" 

117 return self._ctx 

118 

119 def _load_eval_yaml_hook_cls(self, hook_cls): 

120 """hook_class in eval yaml""" 

121 target_mod_path = None 

122 try: 

123 # search in folder 'hooks' 

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

125 for f in os.listdir(hooks_dir): 

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

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

128 spec = importlib.util.find_spec(mod_path) 

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

130 continue 

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

132 source = mf.read() 

133 tree = ast.parse(source) 

134 mod_cls = None 

135 for node in ast.walk(tree): 

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

137 mod_cls = node 

138 break 

139 if mod_cls: 

140 target_mod_path = mod_path 

141 break 

142 if target_mod_path: 

143 module = importlib.import_module(target_mod_path) 

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

145 except (ModuleNotFoundError, ImportError) as e: 

146 print(e) 

147 

148 # Peak memory estimation 

149 

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

151 """processing input config""" 

152 if self._ccfg is not None: 

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

154 else: 

155 self._ccfg = CostModelConfig( 

156 new_config, self.hook_cls, self.framework, self.source_code, devices=getattr(self, "devices", 0) 

157 ) 

158 if isinstance(new_config, str) and not self.config_path: 

159 logger.info( 

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

161 "=" * 30, 

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

163 ) 

164 # Kept whatever its form, so reset_config() can re-parse a config 

165 # handed over in memory rather than read from a path. 

166 self.config_path = new_config 

167 self.evaluator_instances = None 

168 self.ppb = None 

169 

170 def _inner_static_mem(self) -> float: 

171 """static memory evaluation for backbone estimation""" 

172 self._ctx.node_grad = 0 

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

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

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

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

177 self._ctx.node_grad = grad 

178 res = p + ost + grad 

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

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

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

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

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

184 _, routed_exp_p, shared_exp_p = result 

185 routed_exp_mem = ( 

186 routed_exp_p / self._ccfg.ep 

187 * self._ccfg.bytes_p / self._ccfg.shard_p_os_exp 

188 ) 

189 shared_exp_mem = ( 

190 shared_exp_p 

191 * self._ccfg.bytes_p / self._ccfg.shard_p_os_exp_partial 

192 ) 

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

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

195 return p + ost + grad 

196 return 0 

197 

198 def _inner_dynamic_mem( 

199 self, ppb=False, default_micro_factor=None 

200 ) -> tuple[float, float]: 

201 """dynamic memory evaluation for backbone estimation""" 

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

203 self._ctx.seq_chunks = 1 

204 if ppb: 

205 micro_factor = 1 

206 elif default_micro_factor: 

207 micro_factor = default_micro_factor 

208 else: 

209 sched = self._ccfg.pp_sched 

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

211 self._ccfg, self._ctx 

212 ) 

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

214 # A schedule that splits the sequence evaluates the node at one 

215 # chunk's length, and leaves the config's as it found it. 

216 full_s = self._ccfg.s 

217 if self._ctx.seq_chunks > 1: 

218 self._ccfg.s = full_s / self._ctx.seq_chunks 

219 try: 

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

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

222 comm = self._inner_comm_mem(micro_factor) 

223 finally: 

224 self._ccfg.s = full_s 

225 return activation, comm 

226 return 0, 0 

227 

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

229 """dynamic memory evaluation for backbone estimation""" 

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

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

232 comm_cat = { 

233 "dp": MemType.AG_COMM, 

234 "tp": MemType.AG_COMM, 

235 "cp": MemType.AG_COMM, 

236 "ep": MemType.A2A_COMM, 

237 } 

238 comm_mem = {} 

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

240 if fun is None: 

241 continue 

242 sig_param = _parameters(fun) 

243 if ( 

244 # any( 

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

246 # for p in sig_param 

247 # ) 

248 # and len(sig_param) > 2 

249 len(sig_param) > 2 

250 ): 

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

252 else: 

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

254 comm_mem[k] = comm 

255 # FSDP that reshards frees a layer's gathered parameters once it 

256 # has run: a layer keeps none between its passes, and its 

257 # backward holds its own and the next layer's, prefetched, as 

258 # its working set counts them. The root's, the embedding and 

259 # output tables, stay gathered from the forward on: their 

260 # nodes keep them, and a working set holds no second copy. 

261 if "dp" in comm_mem and self._ccfg.reshards: 

262 if self.is_regular_layer(self._ctx.current_node): 

263 comm_mem["dp"] *= self._ctx.working_set 

264 elif self._ctx.working_set: 

265 comm_mem["dp"] = 0 

266 res = EvalUtils.eval_expr_insight( 

267 expr=self._ctx.comm_expr, 

268 ctx=self._ctx, 

269 mem_val=comm_mem, 

270 mem_cat=comm_cat, 

271 ) 

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

273 return res 

274 return 0 

275 

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

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

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

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

280 color_stat = plt.get_cmap("cividis") 

281 color_dyn = plt.get_cmap("plasma") 

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

283 color_denominator = max(1, total_lay - 1) 

284 for chunk_id, chunk in enumerate(stage_i): 

285 for lay_id, lay_type in enumerate(chunk): 

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

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

288 idx = chunk_id * len(dyn_mem_i) + lay_id 

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

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

291 ax.bar( 

292 0.1, 

293 [dyn], 

294 bottom=bottoms["dyn"], 

295 label=name, 

296 color=color_dyn(idx / color_denominator), 

297 width=0.15, 

298 linewidth=0.5, 

299 edgecolor="black", 

300 ) 

301 ax.text( 

302 0.2, 

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

304 f"DYN_{name}", 

305 va="center", 

306 fontsize=5, 

307 ) 

308 ax.bar( 

309 0.5, 

310 [stat], 

311 bottom=bottoms["stat"], 

312 label=name, 

313 color=color_stat(idx / color_denominator), 

314 width=0.15, 

315 linewidth=0.5, 

316 edgecolor="black", 

317 ) 

318 ax.text( 

319 0.6, 

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

321 f"STAT_{name}", 

322 va="center", 

323 fontsize=5, 

324 ) 

325 bottoms["dyn"] += dyn 

326 bottoms["stat"] += stat 

327 ax.axhline( 

328 self.get_max_device_memory(), 

329 color="red", 

330 linewidth=1, 

331 ls="dotted", 

332 ) 

333 ax.axhline( 

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

335 color="blue", 

336 linewidth=1, 

337 ls="dashed", 

338 ) 

339 ax.text( 

340 0, 

341 self.get_max_device_memory(), 

342 "Device Memory", 

343 color="red", 

344 va="bottom", 

345 fontsize=8, 

346 ) 

347 ax.text( 

348 0.62, 

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

350 "Prediction total", 

351 color="blue", 

352 va="bottom", 

353 fontsize=8, 

354 ) 

355 ax.set_xticks([]) 

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

357 ax.set_xlim([0, 0.8]) 

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

359 

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

361 """plot bars for estimations""" 

362 

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

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

365 os.makedirs("plots") 

366 imgs = [] 

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

368 self.__subplot_stages( 

369 stages[stage_id], 

370 stat_mems[stage_id], 

371 dyn_mems[stage_id], 

372 stage_id, 

373 ) 

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

375 imgs += [(stage_id, img)] 

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

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

378 plt.clf() 

379 plt.close() 

380 # Concatenate every stage plots 

381 if imgs: 

382 imgs = sorted(imgs) 

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

384 stage_canva = Image.new( 

385 "RGB", 

386 ( 

387 canvas[0].size[0] 

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

389 canvas[0].size[1] 

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

391 ), 

392 ) 

393 offset_x, offset_y = 0, 0 

394 for i in canvas: 

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

396 offset_x += i.size[0] 

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

398 offset_x = 0 

399 offset_y += i.size[1] 

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

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

402 

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

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

405 if not stage_logs[stage_id].node_compute_log: 

406 stage_logs[stage_id].node_compute_log = {} 

407 stage_logs[stage_id].node_compute_log.update( 

408 self._ctx.node_compute_log 

409 ) 

410 if not stage_logs[stage_id].accu_mem_type: 

411 stage_logs[stage_id].accu_mem_type = { 

412 mt: 0 for mt in list(MemType) 

413 } 

414 for mem_type in list(MemType): 

415 val = self._ctx.accu_mem_type[mem_type] 

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

417 

418 def __order_layer_switches(self, stages: list) -> Optional[list]: 

419 """Each regular layer's own recompute switches, in the order the kinds take; None where rec_op holds them.""" 

420 per_layer = layer_switches(self._ccfg) 

421 if per_layer is None: 

422 return None 

423 order = list(per_layer) 

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

425 first = sum(len(s[0]) for s in stages) - 1 

426 order = order[:first] + order[first:][::-1] 

427 return order 

428 

429 def __order_layer_kinds(self, stages: list, kinds: list) -> list: 

430 """Check the layers' kinds against the partition, in the order the stages visit them.""" 

431 flatten = list(kinds) 

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

433 total_n_lay_stages = self._ccfg.count_layers(stages) 

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

435 raise(AttributeError( 

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

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

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

439 )) 

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

441 raise AttributeError( 

442 f"layer kinds ({len(flatten)})" 

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

444 ) 

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

446 n_layer_first_chunk = ( 

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

448 ) # Except embedding layer 

449 flatten = ( 

450 flatten[:n_layer_first_chunk] 

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

452 ) 

453 return flatten 

454 

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

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

457 stages = args[0] 

458 spec_stage_id = args[3] 

459 

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

461 spec_stage_id = -1 

462 # Process partition generation 

463 if not stages: 

464 stages = self._ccfg.generate_partitions_vpp() 

465 # multimodal=self._ccfg.multimodal 

466 # ) 

467 self.__name_main_module(self._ccfg) 

468 if not self._ccfg.multimodal: 

469 return self.__estimate_stages_backbone( 

470 stages, args[1], args[2], spec_stage_id, args[4], layer_kinds(self._ccfg) 

471 ) 

472 res = [] 

473 original_ccfg = self._ccfg 

474 kinds = [] 

475 self.evaluator_instances = [] 

476 # Build the layers' kinds, submodule after submodule + Build temporary evaluators 

477 for m in self._ccfg.mm_order: 

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

479 if not self._child_cls: 

480 raise AttributeError("expected non null _child_cls") 

481 tmp_evaluator = type(self._child_cls)( 

482 None, 

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

484 trace_fun=getattr(self, "toggle_func_trace", False), 

485 ) 

486 tmp_evaluator.import_eval_yaml() 

487 num_layer = tmp_evaluator.get_num_layers() 

488 # assert not isinstance(num_layer, tuple) 

489 strategy = tmp_evaluator.get_strategy() 

490 full_rec = strategy["full_rec"] 

491 offset = strategy["offset"] 

492 # A user hook class names one hook per submodule; built-in 

493 # submodules are dispatched by their arch. 

494 if self._ccfg.hooks_dict: 

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

496 else: 

497 check_and_apply_custom_hook(tmp_evaluator) 

498 strategy = tmp_evaluator.get_strategy() 

499 if ( 

500 tmp_evaluator.get_num_layers() != num_layer 

501 or strategy["full_rec"] != full_rec 

502 or strategy["offset"] != offset 

503 ): 

504 # Reverify num layers, recomp, offset after hook 

505 stages[m] = ( 

506 tmp_evaluator.ccfg.generate_partitions_vpp_unimodal() 

507 ) 

508 # Its layers are priced plain, as the hook left them. 

509 tmp_evaluator.ccfg.layer_stack = None 

510 kinds += layer_kinds(self._ccfg.mm_ccfgs[m]) 

511 self.evaluator_instances += [tmp_evaluator] 

512 if args[1]: 

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

514 tmp_evaluator.print_ctx() 

515 self.print_stages(stages[m]) 

516 if args[1]: 

517 logger.info( 

518 "Combined layer kinds for %s\n%s", 

519 self._ccfg.model_name, 

520 pprint.pformat([kind.name if kind else None for kind in kinds], compact=True), 

521 ) 

522 logger.info( 

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

524 self._ccfg.model_name, 

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

526 ) 

527 

528 res = self.__estimate_stages_backbone( 

529 self._ccfg.combine_partition_multimodal(stages), 

530 args[1], 

531 args[2], 

532 spec_stage_id, 

533 args[4], 

534 kinds, 

535 ) 

536 self._ccfg = original_ccfg 

537 return res 

538 

539 @staticmethod 

540 def _stage_buffers(stages: list, pipeline: int) -> Dict: 

541 """What the layer loop fills, per stage, chunk and layer, and per stage.""" 

542 def per_layer(value: Any) -> List[List[List[Any]]]: 

543 """One *value* per layer of every chunk of every stage.""" 

544 return [[[value for _ in c] for c in s] for s in stages] 

545 

546 return { 

547 "stat": per_layer(0), 

548 "dyn": per_layer(0), 

549 "comm": per_layer(0), 

550 "micro": per_layer(1), 

551 "grad_out": [0 for _ in stages], 

552 "root_grad": [0 for _ in stages], 

553 "logs": [Config({}) for _ in range(pipeline)], 

554 } 

555 

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

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

558 stages = args[0] 

559 verbose = args[1] 

560 compute_ppb = args[2] 

561 spec_stage_id = args[3] 

562 

563 if verbose: 

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

565 self._ccfg.print_stages(stages, spec_stage_id) 

566 insights = [] 

567 # Compute peak memory 

568 flatten = self.__order_layer_kinds(stages, args[5]) 

569 self._switch_order = self.__order_layer_switches(stages) 

570 self._ctx.switches = None 

571 if verbose: 

572 logger.info( 

573 "Layer kinds in stage order\n%s", 

574 pprint.pformat( 

575 [kind.name if kind else None for kind in flatten], compact=True 

576 ), 

577 ) 

578 

579 stage_misc = self._stage_buffers(stages, self._ccfg.p) 

580 # tmp_ppb_lay_desc = [] # PPB purpose 

581 ppb_lay_desc = [] 

582 record_lay_types = {} 

583 self.__chunk_stage_lay_loops( 

584 flatten, 

585 stages, 

586 record_lay_types, 

587 stage_misc, 

588 verbose, 

589 compute_ppb, 

590 ppb_lay_desc, # tmp_ppb_lay_desc, 

591 ) 

592 self.__postprocess_stages( 

593 stages, 

594 record_lay_types, 

595 stage_misc, 

596 verbose, 

597 spec_stage_id, 

598 insights, 

599 ) 

600 # PPB Input 

601 ppb_input = None 

602 if compute_ppb: 

603 self._ppb_obj.ppb_withdraw_dominated(ppb_lay_desc) 

604 self._ppb_obj.ppb_scale_times(ppb_lay_desc) 

605 self._ppb_obj.ppb_combine_bodies(ppb_lay_desc) 

606 ppb_input = {"layers_description": ppb_lay_desc} 

607 if args[4]: # Plot 

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

609 self._switch_order = None 

610 self._ctx.switches = None 

611 return insights, ppb_input 

612 

613 def __update_evaluator(self, node, verbose): 

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

615 tmp_eval = self.evaluator_instances.pop(0) 

616 self._ccfg = tmp_eval.ccfg 

617 self._ctx.copy_tmp_buff(tmp_eval.ctx) 

618 self._ctx = tmp_eval.ctx 

619 if verbose: 

620 logger.info( 

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

622 self._ccfg.model_name, 

623 ) 

624 

625 def __apply_next_layer_kind(self, *args): 

626 """Give the next regular layer its kind before evaluating it.""" 

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

628 record_lay_types = args[2] 

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

630 node = args[6] 

631 kind = None 

632 if self.is_regular_layer(node) and flatten: 

633 kind = flatten.pop(0) 

634 # The layer's own recompute switches, which its context carries. 

635 switches = None 

636 if self.is_regular_layer(node) and self._switch_order: 

637 switches = self._switch_order.pop(0) 

638 self._ctx.switches = switches 

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

640 self._ccfg, 

641 self._ctx, 

642 kind, 

643 switches, 

644 ) 

645 if kind is not None: 

646 if verbose: 

647 logger.info("Apply layer kind %s", kind.name) 

648 apply_layer_kind(self, kind) 

649 if verbose: 

650 logger.info( 

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

652 stage_id, 

653 chunk_id, 

654 lay_id, 

655 node, 

656 ) 

657 self._ccfg.print_parallelism() 

658 

659 def __chunk_stage_lay_loops(self, *args): 

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

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

662 sm = args[3] 

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

664 self._ctx.real_lay_ids = [] 

665 count = 0 

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

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

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

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

670 self._ctx.init_tmp_buff() 

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

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

673 if self.is_regular_layer(node): 

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

675 count += 1 

676 else: 

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

678 # Update evaluator (multimodal) 

679 self.__update_evaluator(node, verbose) 

680 # Give the next layer its kind 

681 self.__apply_next_layer_kind( 

682 flatten, 

683 verbose, 

684 record_lay_types, 

685 stage_id, 

686 chunk_id, 

687 lay_id, 

688 node, 

689 ) 

690 self._ctx.current_stage_id = stage_id 

691 self._ctx.current_chunk_id = chunk_id 

692 self._ctx.current_lay_id = lay_id 

693 self._ctx.current_node = node 

694 static_mem = self._inner_static_mem() 

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

696 activation, comm = self._inner_dynamic_mem() 

697 sm["dyn"][stage_id][chunk_id][lay_id] = activation + comm 

698 sm["comm"][stage_id][chunk_id][lay_id] = comm 

699 sm["micro"][stage_id][chunk_id][lay_id] = self._ctx.micro_factor or 1 

700 if self.is_regular_layer(node): 

701 sm["grad_out"][stage_id] += self._ctx.node_grad 

702 elif node == LayerType.OUTPUT_LAYER and self._ccfg.overlaps_grad_reduce: 

703 sm["root_grad"][stage_id] += EvalTailSingle.reduced_grad_out_single(self._ccfg, self._ctx)[0] 

704 if verbose: 

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

706 # PPB Purpose 

707 if compute_ppb: 

708 desc = self._ppb_obj.lay_ppb( 

709 self._ccfg, 

710 self._ctx, 

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

712 record_lay_types[(stage_id, chunk_id, lay_id)][2], 

713 ) 

714 self._ppb_obj.add_to_ppb_list(ppb_lay_desc, desc) 

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

716 

717 def __backward_end(self, stages, stage_id, sm, record_lay_types) -> float: 

718 """The stage's dynamic memory as a micro-batch's backward ends, or 0. 

719 

720 FSDP that holds each layer's reduce-scatter output until the backward 

721 ends, and accumulates gradients over several micro-batches, then 

722 holds the outputs of the stage's layers beside the gradients it 

723 accumulated; FSDP that reduces a layer's gradients while the next 

724 layer's backward runs holds whole gradients meanwhile, and the 

725 root's, the output table's, until the backward ends. Either's peak 

726 can come as the backward ends, when that micro-batch keeps nothing 

727 but the other micro-batches in flight still do. The root's gathered 

728 tables stay, and the stage's first layer runs its backward last, 

729 before its own reduce-scatter starts. 

730 """ 

731 held = sm["grad_out"][stage_id] if self._ccfg.defers_grads and self._ccfg.m > 1 else 0