Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / sapp / sapp_pipeline.py: 92%

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"""High-level orchestrator around :class:`SappSolver`: build, solve, simulate, export YAML.""" 

16import os 

17import sys 

18from typing import Any, Dict, List, Optional, Union 

19 

20import matplotlib.pyplot as plt 

21import yaml 

22 

23import hyper_parallel.auto_parallel.sapp_ppb.simulator.pp_simulator as sim 

24import hyper_parallel.auto_parallel.sapp_ppb.utils.recompute as Recompute 

25from hyper_parallel.auto_parallel.sapp_ppb.sapp.sapp_solver import SappSolver 

26from hyper_parallel.auto_parallel.sapp_ppb.utils.check_rules import check_yaml_depth_before_loading 

27from hyper_parallel.auto_parallel.sapp_ppb.utils.layer import Layer, filter_layer_type 

28from hyper_parallel.auto_parallel.sapp_ppb.utils.logger import logger 

29 

30 

31class SappPipeline: 

32 """pipeline balancer""" 

33 

34 def __init__( 

35 self, 

36 model_name: str, 

37 num_of_stage: int, 

38 num_of_micro_batch: int, 

39 max_memory: int, 

40 layers: List[Layer], 

41 vpp_less_memory: bool = False, 

42 # Add arg dual 

43 dual: bool = False, 

44 num_of_interleave: int = 1, 

45 constant_memory: int = 0, 

46 optimization_level: int = 1, 

47 extracted_training_params: Optional[Dict[str, int]] = None, 

48 seq_split_num: int = 1, 

49 ) -> None: 

50 """Cache pipeline parameters and index the input ``layers`` by HEAD / BODY / TAIL. 

51 

52 Args: 

53 model_name (str): Model identifier, used for dump filenames and log prefixes. 

54 num_of_stage (int): Number of physical pipeline stages. 

55 num_of_micro_batch (int): Number of micro-batches scheduled per iteration. 

56 max_memory (int): Per-device memory budget in MB. 

57 layers (List[Layer]): Ordered list of layer descriptors covering HEAD/BODY/TAIL. 

58 vpp_less_memory (bool, optional): If ``True``, use the less-memory VPP scheduler variant. 

59 Default: ``False``. 

60 dual (bool, optional): Enable dualpipe-V scheduling support. Default: ``False``. 

61 num_of_interleave (int, optional): Virtual-pipeline (VPP) chunk count. Default: ``1``. 

62 constant_memory (int, optional): Constant per-stage memory overhead (MB). Default: ``0``. 

63 optimization_level (int, optional): Solver optimization level (``0-2``). Default: ``1``. 

64 extracted_training_params (Optional[Dict[str, int]], optional): Optional training-config parameters for 

65 seqpp. Default: ``None``. 

66 seq_split_num (int, optional): Number of sequence splits; ``>1`` enables sequence pipeline. 

67 Default: ``1``. 

68 """ 

69 self.model_name_ = model_name 

70 self.num_of_stage_ = num_of_stage 

71 self.num_of_micro_batch_ = num_of_micro_batch 

72 self.num_of_interleave_ = num_of_interleave 

73 self.max_memory_ = max_memory 

74 self.vpp_less_memory_ = vpp_less_memory 

75 # Add arg dual_ 

76 self.dual_ = dual 

77 self.constant_memory_ = constant_memory 

78 self.optimization_level = optimization_level 

79 self.extracted_training_params_ = extracted_training_params 

80 self.seq_split_num_ = seq_split_num 

81 self.seqpipe_ = self.seq_split_num_ > 1 

82 # logger.output("seq chunk: %s",self.seq_split_num_) 

83 

84 self.problem_ = None 

85 self.layers_ = layers 

86 self.layers_sorted_ = { 

87 Layer.type_enum.HEAD: filter_layer_type(layers, 

88 Layer.type_enum.HEAD), 

89 Layer.type_enum.BODY: filter_layer_type(layers, 

90 Layer.type_enum.BODY), 

91 Layer.type_enum.TAIL: filter_layer_type(layers, 

92 Layer.type_enum.TAIL), 

93 } 

94 

95 def has_some_memory_info(self) -> bool: 

96 """Check if there is all information for memory constraint.""" 

97 return self.problem_.has_some_memory_info() 

98 

99 def construct_problem(self, solver: str = "pulp") -> None: 

100 """Construct the underlying ILP problem using the requested solver backend.""" 

101 if solver == "pulp": 

102 self.problem_ = self._construct_problem_pulp_() 

103 elif solver == "other": 

104 logger.warning( 

105 "No other solver available..., automatically switch to pulp!!!" 

106 ) 

107 self.problem_ = self._construct_problem_pulp_() 

108 else: 

109 logger.warning( 

110 "No other solver available..., automatically switch to pulp!!!" 

111 ) 

112 self.problem_ = self._construct_problem_pulp_() 

113 

114 def solve_problem(self, time_limit: int = 90, dump_folder: Optional[str] = None) -> None: 

115 """Solve the ILP, optionally dumping the LP model into ``dump_folder``.""" 

116 self.problem_.solve(time_limit, dump_folder) 

117 

118 def get_result(self) -> dict[str, list[list[str]]]: 

119 """Get result distribution of the solution (compact form).""" 

120 return self.problem_.result() 

121 

122 def get_memory_activation(self) -> list[float]: 

123 """Get the activation memory per stage for simulator.""" 

124 return self.problem_.get_simulator_memory_activation() 

125 

126 def get_memory_parameter(self) -> list[float]: 

127 """Get the parameter memory per stage for simulator.""" 

128 return self.problem_.get_simulator_memory_parameter() 

129 

130 def get_fw_time(self) -> list[float]: 

131 """Get the forward time per stage for simulator.""" 

132 time = self.problem_.get_simulator_forward_time() 

133 return time 

134 

135 def get_recompute_time(self) -> list[float]: 

136 """Get the recompute time per stage for simulator.""" 

137 time = self.problem_.get_simulator_recompute_time() 

138 return time 

139 

140 def get_time(self) -> list[float]: 

141 """Get the time per stage for simulator.""" 

142 return self.problem_.get_simulator_time() 

143 

144 def naive_layer_per_stage(self, 

145 layer_num: int, 

146 num_of_interleave: int = 1) -> List[List[int]]: 

147 """Return the naive layer-to-stage assignment (``layer_num`` evenly split).""" 

148 logger.output("layer_num = %s", layer_num) 

149 layer_count = layer_num // (self.num_of_stage_ * num_of_interleave) 

150 return [[layer_count] * self.num_of_stage_ for _ in range(num_of_interleave)] 

151 

152 def print_yaml_results(self) -> None: 

153 """Log the solver output in the MindFormers YAML schema.""" 

154 

155 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

156 nass = self.naive_layer_per_stage(layer.nb_layer_, 

157 self.num_of_interleave_) 

158 yaml_format = Recompute.yaml_from_internal( 

159 self.num_of_interleave_, 

160 self.num_of_stage_, 

161 self.problem_.variables_[layer.name_], 

162 nass, 

163 ) 

164 logger.output("layer-to-stage assignment baseline is \n\t%s", nass) 

165 yaml_results = "\nTo put in yaml configuration:" 

166 for y, v in yaml_format.items(): 

167 yaml_results += f"\n\t{y}: {v}" 

168 logger.output(yaml_results) 

169 

170 def get_manual_memory_activation( 

171 self, 

172 each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]], 

173 interleave_num: int = 1) -> List[List[float]]: 

174 """Return the per-stage activation memory for a user-supplied layer assignment.""" 

175 memory_active = [] 

176 if self.has_some_memory_info(): 

177 for inter in range(interleave_num): 

178 memory_active.append([]) 

179 for stage in range(self.num_of_stage_): 

180 memory_activation = 0 

181 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

182 memory_activation += self._get_layer_memory_activation( 

183 each_layer_per_recompute, layer, inter, stage 

184 ) 

185 memory_active[inter].append(memory_activation) 

186 return memory_active 

187 

188 @staticmethod 

189 def _get_layer_memory_activation(each_layer_per_recompute, layer, interleave, stage): 

190 """Calculate activation memory for one layer at one pipeline position.""" 

191 memory_activation = 0 

192 unused_recompute_list = Recompute.get_unused_list(each_layer_per_recompute[layer]) 

193 for rec in Recompute.TYPE: 

194 if rec in unused_recompute_list: 

195 continue 

196 value = each_layer_per_recompute[layer][rec][interleave][stage] 

197 if value > 0: 

198 memory_activation += value * layer.memory_activation_rec_[rec] 

199 return memory_activation 

200 

201 def get_manual_memory_parameter( 

202 self, 

203 each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]], 

204 interleave_num: int = 1) -> List[List[float]]: 

205 """Return the per-stage parameter memory for a user-supplied layer assignment.""" 

206 memory_param_stage = [0] * self.num_of_stage_ 

207 for inter in range(interleave_num): 

208 for stage in range(self.num_of_stage_): 

209 for rec in Recompute.TYPE: 

210 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

211 if layer.memory_parameter_ is None: 

212 continue 

213 

214 if rec in Recompute.get_unused_list(each_layer_per_recompute[layer]): 

215 continue 

216 

217 value = each_layer_per_recompute[layer][rec][inter][stage] 

218 if value <= 0: 

219 continue 

220 

221 memory_param_stage[stage] += value * layer.memory_parameter_ 

222 for head in self.layers_sorted_[Layer.type_enum.HEAD]: 

223 if head.memory_parameter_ is not None: 

224 memory_param_stage[0] += head.memory_parameter_ 

225 for tail in self.layers_sorted_[Layer.type_enum.TAIL]: 

226 if tail.memory_parameter_ is not None: 

227 memory_param_stage[self.num_of_stage_ - 

228 1] += tail.memory_parameter_ 

229 memory_param = [memory_param_stage] * interleave_num 

230 return memory_param 

231 

232 def get_manual_time( 

233 self, 

234 each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]], 

235 interleave_num: int = 1) -> List[List[float]]: 

236 """Return the per-stage execution time for a user-supplied layer assignment.""" 

237 time = [] 

238 for i in range(interleave_num): 

239 time.append([]) 

240 for s in range(self.num_of_stage_): 

241 time[i].append(0) 

242 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

243 for r in Recompute.TYPE: 

244 if each_layer_per_recompute[layer][r][i][s] > 0: 

245 time[i][s] += each_layer_per_recompute[layer][r][i][s] * ( 

246 layer.forward_time_ + 

247 layer.backward_time_rec_[r]) 

248 

249 for head in self.layers_sorted_[Layer.type_enum.HEAD]: 

250 time[0][0] += head.time_ 

251 for tail in self.layers_sorted_[Layer.type_enum.TAIL]: 

252 time[interleave_num - 1][self.num_of_stage_ - 1] += tail.time_ 

253 return time 

254 

255 def get_manual_fw_time( 

256 self, 

257 each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]], 

258 interleave_num: int = 1) -> List[List[float]]: 

259 """Return the per-stage forward time for a user-supplied layer assignment.""" 

260 time = [] 

261 for i in range(interleave_num): 

262 time.append([]) 

263 for s in range(self.num_of_stage_): 

264 time[i].append(0) 

265 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

266 for r in Recompute.TYPE: 

267 if (r not in Recompute.get_unused_list(each_layer_per_recompute[layer]) 

268 and each_layer_per_recompute[layer][r][i][s] > 0): 

269 time[i][s] += each_layer_per_recompute[layer][r][i][s] * ( 

270 layer.forward_time_) 

271 for head in self.layers_sorted_[Layer.type_enum.HEAD]: 

272 time[0][0] += head.time_ 

273 for tail in self.layers_sorted_[Layer.type_enum.TAIL]: 

274 time[interleave_num - 1][self.num_of_stage_ - 1] += tail.time_ 

275 return time 

276 

277 def get_manual_recompute_time( 

278 self, 

279 each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]], 

280 interleave_num: int = 1) -> List[List[float]]: 

281 """Return the per-stage recompute-only time for a user-supplied layer assignment.""" 

282 logger.output("each_layer_per_recompute = %s", each_layer_per_recompute) 

283 time_all_rec = [] 

284 time_no_rec = [] 

285 for i in range(interleave_num): 

286 time_all_rec.append([]) 

287 time_no_rec.append([]) 

288 for s in range(self.num_of_stage_): 

289 time_all_rec[i].append(0) 

290 time_no_rec[i].append(0) 

291 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

292 self._add_manual_recompute_time( 

293 each_layer_per_recompute, layer, i, s, time_all_rec, time_no_rec) 

294 

295 return [[r - n for r, n in zip(ar, nr)] 

296 for ar, nr in zip(time_all_rec, time_no_rec)] 

297 

298 def _add_manual_recompute_time(self, each_layer_per_recompute, layer, interleave, stage, 

299 time_all_rec, time_no_rec): 

300 """Accumulate recompute time for a single layer and stage.""" 

301 logger.output("backward_time_rec_(%s) = %s", layer, layer.backward_time_rec_) 

302 unused_rec = Recompute.get_unused_list(each_layer_per_recompute[layer]) 

303 for rec in Recompute.TYPE: 

304 layer_num = each_layer_per_recompute[layer][rec][interleave][stage] 

305 if rec in unused_rec or layer_num <= 0: 

306 continue 

307 if layer.backward_time_rec_[rec] is None: 

308 raise ValueError("No backward tme is specified for this " 

309 "recomputation. Recomputation " 

310 f"'{Recompute.YAML_NAME[rec]}' is likely not considered") 

311 logger.output("r = %s; i = %s; s = %s", rec, interleave, stage) 

312 time_all_rec[interleave][stage] += layer_num * layer.backward_time_rec_[rec] 

313 time_no_rec[interleave][stage] += layer_num * layer.backward_time_rec_[Recompute.TYPE.NONE] 

314 

315 def simulate(self, show: bool = True, file_name: Optional[str] = None, 

316 sub_fig: Optional[plt.Figure] = None) -> float: 

317 """Run the simulator on the solved schedule and return its estimated total time.""" 

318 forward_time = self.get_fw_time() 

319 recompute_overhead = self.get_recompute_time() 

320 stage_mem_par = 0 

321 stage_mem_act = 0 

322 if self.has_some_memory_info(): 

323 stage_mem_par = self.get_memory_parameter() 

324 stage_mem_act = self.get_memory_activation() 

325 

326 return self.simulation( 

327 forward_time, 

328 recompute_overhead, 

329 stage_mem_par, 

330 stage_mem_act, 

331 self.constant_memory_, 

332 show, 

333 file_name, 

334 sub_fig 

335 ) 

336 

337 def simulate_naive(self, layers: List[Layer], output_folder: str) -> None: 

338 """Simulate the naive (even) layer-to-stage assignments for sanity comparison.""" 

339 num_layers = 0 

340 rec_considered = {} 

341 for layer in layers: 

342 if layer.type_ == Layer.type_enum.BODY: 

343 num_layers = layer.nb_layer_ 

344 rec_considered = layer.recompute_considered_ 

345 

346 all_recomp = {"offset": 0} 

347 no_recomp = {"offset": 0} 

348 for rec in [Recompute.TYPE.FULL, Recompute.TYPE.SLCT, Recompute.TYPE.COMM]: 

349 if rec_considered.get(rec, False): 

350 all_recomp[Recompute.YAML_NAME[rec]] = True 

351 no_recomp[Recompute.YAML_NAME[rec]] = False 

352 

353 self.simulate_yaml( 

354 yaml_format=all_recomp, 

355 show=True, 

356 interleave_num=self.num_of_interleave_, 

357 file_name=os.path.join(output_folder, 

358 "result_naive_all_recomp.svg"), 

359 ) 

360 

361 if num_layers % self.num_of_stage_ == 0: 

362 self.simulate_yaml( 

363 yaml_format=no_recomp, 

364 show=True, 

365 interleave_num=self.num_of_interleave_, 

366 file_name=os.path.join(output_folder, 

367 "result_naive_no_recomp.svg"), 

368 ) 

369 else: 

370 logger.warning("num layer cannot be divided by num stage") 

371 

372 def simulate_comparison(self, manual_config_file: str, output_folder: str) -> None: 

373 """Render side-by-side automatic vs manual simulations for every entry in the YAML.""" 

374 with open(manual_config_file, encoding="utf-8") as fp: 

375 check_yaml_depth_before_loading(fp) 

376 fp.seek(0) 

377 data = yaml.safe_load(fp) 

378 yaml_data = {} 

379 for manual in data.values(): 

380 yaml_data[Recompute.OFFSET] = manual.get(Recompute.OFFSET) 

381 if isinstance(yaml_data[Recompute.OFFSET], list) and all( 

382 isinstance(item, int) for item in yaml_data[Recompute.OFFSET]): 

383 yaml_data[Recompute.OFFSET] = [yaml_data[Recompute.OFFSET]] 

384 

385 for rec in Recompute.YAML_NAME.values(): 

386 yaml_data[rec] = manual.get(rec) 

387 if isinstance(yaml_data[rec], list) and all( 

388 isinstance(item, int) for item in yaml_data[rec]): 

389 yaml_data[rec] = [yaml_data[rec]] 

390 interleave_num = manual.get("interleave_num", 

391 self.num_of_interleave_) 

392 show = manual.get("show", False) 

393 file_name = manual.get("file_name") 

394 full_file_name = os.path.join(output_folder, 

395 file_name) if (file_name) else None 

396 

397 fig = plt.figure(figsize=(24, 8)) 

398 sub_figs = fig.subfigures(1, 2, wspace=0.07) 

399 sub_figs[0].suptitle('Automatic', fontsize='x-large') 

400 try: 

401 simulate_result = self.simulate( 

402 show=False, 

403 file_name=os.path.join(output_folder, "Auto_" + file_name), 

404 sub_fig=sub_figs[0], 

405 ) 

406 except Exception: 

407 logger.exception("Failed to simulate auto pipeline.") 

408 raise 

409 

410 if simulate_result is None: 

411 raise RuntimeError("simulate() returned None.") 

412 

413 sub_figs[1].suptitle('Manual', fontsize='x-large') 

414 self.simulate_yaml(yaml_data, False, interleave_num, full_file_name, sub_figs[1]) 

415 plt.savefig(os.path.join(output_folder, "Comparison_" + file_name)) 

416 if show: 

417 plt.show() 

418 

419 def simulate_only_manual(self, manual_config_file: str, output_folder: str) -> None: 

420 """Render only the manual simulation for every entry in ``manual_config_file``.""" 

421 with open(manual_config_file, encoding="utf-8") as fp: 

422 check_yaml_depth_before_loading(fp) 

423 fp.seek(0) 

424 data = yaml.safe_load(fp) 

425 yaml_data = {} 

426 for manual in data.values(): 

427 yaml_data[Recompute.OFFSET] = manual.get(Recompute.OFFSET) 

428 if isinstance(yaml_data[Recompute.OFFSET], list) and all( 

429 isinstance(item, int) for item in yaml_data[Recompute.OFFSET]): 

430 yaml_data[Recompute.OFFSET] = [yaml_data[Recompute.OFFSET]] 

431 

432 for rec in Recompute.YAML_NAME.values(): 

433 yaml_data[rec] = manual.get(rec) 

434 if isinstance(yaml_data[rec], list) and all( 

435 isinstance(item, int) for item in yaml_data[rec]): 

436 yaml_data[rec] = [yaml_data[rec]] 

437 interleave_num = manual.get("interleave_num", 

438 self.num_of_interleave_) 

439 show = manual.get("show", False) 

440 file_name = manual.get("file_name") 

441 full_file_name = os.path.join(output_folder, 

442 file_name) if (file_name) else None 

443 

444 fig = plt.figure(figsize=(12, 8)) 

445 self.simulate_yaml(yaml_data, False, interleave_num, full_file_name, fig) 

446 plt.savefig(os.path.join(output_folder, "manual_file_" + file_name)) 

447 if show: 

448 plt.show() 

449 

450 def simulate_yaml(self, yaml_format: Dict[str, Any], show: bool = True, 

451 interleave_num: int = 1, 

452 file_name: Optional[str] = None, 

453 sub_fig: Optional[plt.Figure] = None) -> float: 

454 """Simulate a manual pipeline configuration encoded as a YAML-compatible dict.""" 

455 layer_num = 0 

456 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

457 layer_num += layer.nb_layer_ 

458 nass = self.naive_layer_per_stage(layer_num, 

459 num_of_interleave=interleave_num) 

460 layer_per_recompute = Recompute.internal_from_yaml( 

461 interleave_num, self.num_of_stage_, yaml_format, nass) 

462 each_layer_per_recompute = self.split_layer_per_recompute(layer_per_recompute) 

463 return self.simulate_manual( 

464 each_layer_per_recompute, 

465 show, 

466 interleave_num=interleave_num, 

467 file_name=file_name, 

468 sub_fig=sub_fig 

469 ) 

470 

471 ####################################################################### 

472 ## ## 

473 ## Print Solver Model ## 

474 ## ## 

475 ####################################################################### 

476 def _calculate_activation_memory(self, each_layer_per_recompute, v, s): 

477 """Calculate activation memory for next and current stage""" 

478 act_mem_next = 0 

479 act_mem_curr = 0 

480 

481 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

482 for rec in Recompute.TYPE: 

483 if self.problem_.recompute_considered_[rec]: 

484 if each_layer_per_recompute[layer][rec][v + 1][s] > 0: # next 

485 act_mem_next += (each_layer_per_recompute[layer][rec][v + 1][s] * 

486 layer.memory_activation_rec_[rec]) 

487 if each_layer_per_recompute[layer][rec][v][s] > 0: # current 

488 act_mem_curr += (each_layer_per_recompute[layer][rec][v][s] * 

489 layer.memory_activation_rec_[rec]) 

490 

491 return act_mem_next, act_mem_curr 

492 

493 def _compute_parameter_memory_manually_solver(self, each_layer_per_recompute, s, interleave_num=1): 

494 """Solver memory model: parameter memory""" 

495 param_mem = 0 

496 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

497 if layer.memory_parameter_ is not None: 

498 param_mem += self._calculate_layer_parameter_memory( 

499 layer, each_layer_per_recompute[layer], s, interleave_num) 

500 return param_mem 

501 

502 def _calculate_layer_parameter_memory(self, layer, layer_per_recompute, s, interleave_num): 

503 """Calculate parameter memory for a single layer""" 

504 layer_mem = 0 

505 for inter in range(interleave_num): 

506 for rec in Recompute.TYPE: 

507 if self.problem_.recompute_considered_[rec]: 

508 if layer_per_recompute[rec][inter][s] > 0: 

509 layer_mem += layer_per_recompute[rec][inter][s] * layer.memory_parameter_ 

510 return layer_mem 

511 

512 def _calculate_activation_memory_solver(self, each_layer_per_recompute, s, interleave_num, activation_nums): 

513 """Calculate activation memory for a given stage""" 

514 act_mem = 0 

515 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

516 for inter in range(interleave_num): 

517 for rec in Recompute.TYPE: 

518 if self.problem_.recompute_considered_[rec]: 

519 if each_layer_per_recompute[layer][rec][inter][s] > 0: 

520 act_mem += (each_layer_per_recompute[layer][rec][inter][s] * 

521 layer.memory_activation_rec_[rec] * 

522 activation_nums[inter][s]) 

523 return act_mem 

524 

525 

526 def debug_print_manual_theoretical_memory( 

527 self, 

528 each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]], 

529 interleave_num: int = 1) -> None: 

530 """Log the per-stage theoretical memory implied by the solver model (debug aid).""" 

531 logger.info("%s Manual Theoretical Memory Analysis %s", "=" * 20, "=" * 20) 

532 

533 if self.vpp_less_memory_: 

534 if self.seqpipe_: 

535 activation_nums = self.problem_.compute_activation_seq_nums( 

536 self.num_of_stage_, interleave_num, self.seq_split_num_, self.num_of_micro_batch_, True) 

537 else: 

538 activation_nums = self.problem_.compute_less_activation_nums( 

539 self.num_of_stage_, interleave_num) 

540 # Add if dual to decide whether dualpipe_v is used 

541 elif self.dual_: 

542 activation_nums = self.problem_.compute_activation_nums_dual( 

543 self.num_of_stage_, interleave_num, self.num_of_micro_batch_) 

544 else: 

545 if self.seqpipe_: 

546 activation_nums = self.problem_.compute_activation_seq_nums( 

547 self.num_of_stage_, interleave_num, self.seq_split_num_, self.num_of_micro_batch_, False) 

548 else: 

549 activation_nums = self.problem_.compute_activation_nums( 

550 self.num_of_stage_, interleave_num, self.num_of_micro_batch_) 

551 

552 logger.info("Activation nums = %s", activation_nums) 

553 

554 # compute for each stage 

555 for s in range(self.num_of_stage_): 

556 

557 # parameter memory 

558 param_mem = self._compute_parameter_memory_manually_solver(each_layer_per_recompute, s, interleave_num) 

559 

560 # head memory 

561 if s == 0: 

562 for head in self.layers_sorted_[Layer.type_enum.HEAD]: 

563 if head.memory_parameter_ is not None: 

564 param_mem += head.memory_parameter_ 

565 

566 # tail memory 

567 if s == self.num_of_stage_ - 1: 

568 for tail in self.layers_sorted_[Layer.type_enum.TAIL]: 

569 if tail.memory_parameter_ is not None: 

570 param_mem += tail.memory_parameter_ 

571 

572 # act memory 

573 act_mem = self._calculate_activation_memory_solver(each_layer_per_recompute, s, 

574 interleave_num, activation_nums) 

575 

576 # overhead 

577 overhead = 0 

578 

579 total = param_mem + act_mem + overhead + self.constant_memory_ 

580 

581 logger.info("Stage %d Manual Memory Analysis:", s) 

582 logger.info("Parameter Memory: %.2f", param_mem) 

583 logger.info("Activation Memory: %.2f", act_mem) 

584 logger.info("Memory Overhead: %.2f", overhead) 

585 logger.info("Constant Memory: %.2f", self.constant_memory_) 

586 logger.info("Total Theoretical Memory: %.2f", total) 

587 

588 def split_layer_per_recompute( 

589 self, 

590 layer_per_recompute: Dict[Recompute.TYPE, List[List[int]]] 

591 ) -> Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]]: 

592 """Split aggregate per-recompute layer counts into counts per BODY layer.""" 

593 each_layer_per_recompute = {} 

594 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

595 rest = layer.nb_layer_ 

596 each_layer_per_recompute[layer] = {r: [] for r in Recompute.TYPE} 

597 for rec in Recompute.TYPE: 

598 for i in range(self.num_of_interleave_): 

599 each_layer_per_recompute[layer][rec].append([0]*self.num_of_stage_) 

600 for s in range(self.num_of_stage_): 

601 subtract = min(layer_per_recompute[rec][i][s], rest) 

602 layer_per_recompute[rec][i][s] -= subtract 

603 rest -= subtract 

604 each_layer_per_recompute[layer][rec][i][s] += subtract 

605 return each_layer_per_recompute 

606 

607 def fuse_layer_per_recompute( 

608 self, 

609 each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]] 

610 ) -> Dict[Recompute.TYPE, List[List[int]]]: 

611 """Fuse per-layer recompute counts back into aggregate per-recompute-type totals.""" 

612 all_layers_per_recompute = {r: [] for r in Recompute.TYPE} 

613 for rec in Recompute.TYPE: 

614 for i in range(self.num_of_interleave_): 

615 all_layers_per_recompute[rec].append([]) 

616 for s in range(self.num_of_stage_): 

617 all_layers_per_recompute[rec][i].append(sum( 

618 each_layer_per_recompute[layer][rec][i][s] 

619 for layer in self.layers_sorted_[Layer.type_enum.BODY] 

620 )) 

621 return all_layers_per_recompute 

622 

623 

624 def simulate_manual( 

625 self, 

626 each_layer_per_recompute: Optional[Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]]] = None, 

627 show: bool = True, 

628 interleave_num: int = 1, 

629 file_name: Optional[str] = None, 

630 sub_fig: Optional[plt.Figure] = None) -> float: 

631 """Run the simulator on a user-supplied per-layer recompute strategy.""" 

632 logger.output("Simulating given strategy: %s", each_layer_per_recompute) 

633 

634 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

635 for rec in Recompute.TYPE: 

636 if len(each_layer_per_recompute[layer][rec]) != interleave_num: 

637 logger.error( 

638 "For layer %s with recompute %s, %s does not match interleave number %s", 

639 layer, 

640 rec, 

641 len(each_layer_per_recompute[layer][rec]), 

642 interleave_num, 

643 ) 

644 return sys.maxsize 

645 

646 for layer in self.layers_sorted_[Layer.type_enum.BODY]: 

647 for rec in Recompute.TYPE: 

648 if any(x < 0 for sublist in each_layer_per_recompute[layer][rec] 

649 for x in sublist): 

650 raise ValueError( 

651 f"for {rec}, there is strategy less than 0 in " 

652 f"{each_layer_per_recompute[layer][rec]}" 

653 ) 

654 

655 forward_time = self.get_manual_fw_time(each_layer_per_recompute, 

656 interleave_num) 

657 recompute_overhead = self.get_manual_recompute_time( 

658 each_layer_per_recompute, interleave_num) 

659 stage_mem_par = 0 

660 stage_mem_act = 0 

661 if self.has_some_memory_info(): 

662 stage_mem_par = self.get_manual_memory_parameter( 

663 each_layer_per_recompute, interleave_num=interleave_num) 

664 stage_mem_act = self.get_manual_memory_activation( 

665 each_layer_per_recompute, interleave_num=interleave_num) 

666 

667 self.debug_print_manual_theoretical_memory(each_layer_per_recompute, interleave_num) 

668 

669 return self.simulation( 

670 forward_time, 

671 recompute_overhead, 

672 stage_mem_par, 

673 stage_mem_act, 

674 self.constant_memory_, 

675 show, 

676 file_name, 

677 sub_fig 

678 ) 

679 

680 def simulation( 

681 self, 

682 forward_time: List[List[float]], 

683 recompute_overhead: Union[int, List[List[float]]] = 0, 

684 stage_mem_par: Union[int, List[List[float]]] = 0, 

685 stage_mem_act: Union[int, List[List[float]]] = 0, 

686 constant_mem: int = 0, 

687 show: bool = True, 

688 file_name: Optional[str] = None, 

689 sub_fig: Optional[plt.Figure] = None, 

690 ) -> float: 

691 """Run the low-level :class:`PipelineSimulator` and return its reported end time.""" 

692 if self.has_some_memory_info(): 

693 logger.output( 

694 "PipelineSimulator(\n\t%s, %s," 

695 "\n\tblock_mem_act=%s," 

696 "\n\tblock_mem_par=%s," 

697 "\n\tlayer_recompute=%s," 

698 "\n\tless_memory=%s )", 

699 forward_time, 

700 self.num_of_micro_batch_, 

701 stage_mem_act, 

702 stage_mem_par, 

703 recompute_overhead, 

704 self.vpp_less_memory_,