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

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 use_backward_time: bool = False, 

50 ) -> None: 

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

52 

53 Args: 

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

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

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

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

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

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

60 Default: ``False``. 

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

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

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

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

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

66 seqpp. Default: ``None``. 

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

68 Default: ``1``. 

69 """ 

70 self.model_name_ = model_name 

71 self.num_of_stage_ = num_of_stage 

72 self.num_of_micro_batch_ = num_of_micro_batch 

73 self.num_of_interleave_ = num_of_interleave 

74 self.max_memory_ = max_memory 

75 self.vpp_less_memory_ = vpp_less_memory 

76 # Add arg dual_ 

77 self.dual_ = dual 

78 self.constant_memory_ = constant_memory 

79 self.optimization_level = optimization_level 

80 self.extracted_training_params_ = extracted_training_params 

81 self.seq_split_num_ = seq_split_num 

82 self.use_backward_time_ = use_backward_time 

83 self.seqpipe_ = self.seq_split_num_ > 1 

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

85 

86 self.problem_ = None 

87 self._simulator: Optional[sim.PipelineSimulator] = None 

88 self.layers_ = layers 

89 self.layers_sorted_ = { 

90 Layer.type_enum.HEAD: filter_layer_type(layers, 

91 Layer.type_enum.HEAD), 

92 Layer.type_enum.BODY: filter_layer_type(layers, 

93 Layer.type_enum.BODY), 

94 Layer.type_enum.TAIL: filter_layer_type(layers, 

95 Layer.type_enum.TAIL), 

96 } 

97 

98 @property 

99 def simulator(self): 

100 """Pipeline simulator instance (available after :meth:`simulate`).""" 

101 return self._simulator 

102 

103 def has_some_memory_info(self) -> bool: 

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

105 return self.problem_.has_some_memory_info() 

106 

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

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

109 if solver == "pulp": 

110 self.problem_ = self._construct_problem_pulp_() 

111 elif solver == "other": 

112 logger.warning( 

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

114 ) 

115 self.problem_ = self._construct_problem_pulp_() 

116 else: 

117 logger.warning( 

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

119 ) 

120 self.problem_ = self._construct_problem_pulp_() 

121 

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

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

124 self.problem_.solve(time_limit, dump_folder) 

125 

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

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

128 return self.problem_.result() 

129 

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

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

132 return self.problem_.get_simulator_memory_activation() 

133 

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

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

136 return self.problem_.get_simulator_memory_parameter() 

137 

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

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

140 time = self.problem_.get_simulator_forward_time() 

141 return time 

142 

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

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

145 time = self.problem_.get_simulator_recompute_time() 

146 return time 

147 

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

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

150 return self.problem_.get_simulator_time() 

151 

152 def naive_layer_per_stage(self, 

153 layer_num: int, 

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

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

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

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

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

159 

160 def print_yaml_results(self) -> None: 

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

162 

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

164 nass = self.naive_layer_per_stage(layer.nb_layer_, 

165 self.num_of_interleave_) 

166 yaml_format = Recompute.yaml_from_internal( 

167 self.num_of_interleave_, 

168 self.num_of_stage_, 

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

170 nass, 

171 ) 

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

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

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

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

176 logger.output(yaml_results) 

177 

178 def get_manual_memory_activation( 

179 self, 

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

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

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

183 memory_active = [] 

184 if self.has_some_memory_info(): 

185 for inter in range(interleave_num): 

186 memory_active.append([]) 

187 for stage in range(self.num_of_stage_): 

188 memory_activation = 0 

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

190 memory_activation += self._get_layer_memory_activation( 

191 each_layer_per_recompute, layer, inter, stage 

192 ) 

193 memory_active[inter].append(memory_activation) 

194 return memory_active 

195 

196 @staticmethod 

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

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

199 memory_activation = 0 

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

201 for rec in Recompute.TYPE: 

202 if rec in unused_recompute_list: 

203 continue 

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

205 if value > 0: 

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

207 return memory_activation 

208 

209 def get_manual_memory_parameter( 

210 self, 

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

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

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

214 memory_param_stage = [0] * self.num_of_stage_ 

215 for inter in range(interleave_num): 

216 for stage in range(self.num_of_stage_): 

217 for rec in Recompute.TYPE: 

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

219 if layer.memory_parameter_ is None: 

220 continue 

221 

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

223 continue 

224 

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

226 if value <= 0: 

227 continue 

228 

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

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

231 if head.memory_parameter_ is not None: 

232 memory_param_stage[0] += head.memory_parameter_ 

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

234 if tail.memory_parameter_ is not None: 

235 memory_param_stage[self.num_of_stage_ - 

236 1] += tail.memory_parameter_ 

237 memory_param = [memory_param_stage] * interleave_num 

238 return memory_param 

239 

240 def get_manual_time( 

241 self, 

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

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

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

245 time = [] 

246 for i in range(interleave_num): 

247 time.append([]) 

248 for s in range(self.num_of_stage_): 

249 time[i].append(0) 

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

251 for r in Recompute.TYPE: 

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

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

254 layer.forward_time_ + 

255 layer.backward_time_rec_[r]) 

256 

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

258 time[0][0] += head.forward_time_ + head.backward_time_rec_[Recompute.TYPE.NONE] 

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

260 time[interleave_num - 1][self.num_of_stage_ - 1] += ( 

261 tail.forward_time_ 

262 + tail.backward_time_rec_[Recompute.TYPE.NONE] 

263 ) 

264 return time 

265 

266 def get_manual_fw_time( 

267 self, 

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

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

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

271 time = [] 

272 for i in range(interleave_num): 

273 time.append([]) 

274 for s in range(self.num_of_stage_): 

275 time[i].append(0) 

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

277 for r in Recompute.TYPE: 

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

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

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

281 layer.forward_time_) 

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

283 time[0][0] += head.forward_time_ 

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

285 time[interleave_num - 1][self.num_of_stage_ - 1] += tail.forward_time_ 

286 return time 

287 

288 def get_manual_backward_time( 

289 self, 

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

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

292 """Return the per-stage backward time for a user-supplied layer assignment.""" 

293 time = [] 

294 for i in range(interleave_num): 

295 time.append([]) 

296 for s in range(self.num_of_stage_): 

297 time[i].append(0) 

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

299 for r in Recompute.TYPE: 

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

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

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

303 layer.backward_time_rec_[r]) 

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

305 time[0][0] += head.backward_time_rec_[Recompute.TYPE.NONE] 

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

307 time[interleave_num - 1][self.num_of_stage_ - 1] += ( 

308 tail.backward_time_rec_[Recompute.TYPE.NONE] 

309 ) 

310 return time 

311 

312 def get_manual_recompute_time( 

313 self, 

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

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

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

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

318 time_all_rec = [] 

319 time_no_rec = [] 

320 for i in range(interleave_num): 

321 time_all_rec.append([]) 

322 time_no_rec.append([]) 

323 for s in range(self.num_of_stage_): 

324 time_all_rec[i].append(0) 

325 time_no_rec[i].append(0) 

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

327 self._add_manual_recompute_time( 

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

329 

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

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

332 

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

334 time_all_rec, time_no_rec): 

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

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

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

338 for rec in Recompute.TYPE: 

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

340 if rec in unused_rec or layer_num <= 0: 

341 continue 

342 if layer.backward_time_rec_[rec] is None: 

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

344 "recomputation. Recomputation " 

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

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

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

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

349 

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

351 sub_fig: Optional[plt.Figure] = None, comm_time: float = 0.0) -> float: 

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

353 forward_time = self.get_fw_time() 

354 recompute_overhead = self.get_recompute_time() 

355 backward_time = self.problem_.get_simulator_backward_time() if self.use_backward_time_ else 0 

356 stage_mem_par = 0 

357 stage_mem_act = 0 

358 if self.has_some_memory_info(): 

359 stage_mem_par = self.get_memory_parameter() 

360 stage_mem_act = self.get_memory_activation() 

361 

362 return self.simulation( 

363 forward_time, 

364 recompute_overhead, 

365 stage_mem_par, 

366 stage_mem_act, 

367 self.constant_memory_, 

368 backward_time=backward_time, 

369 show=show, 

370 file_name=file_name, 

371 sub_fig=sub_fig, 

372 comm_time=comm_time, 

373 ) 

374 

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

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

377 num_layers = 0 

378 rec_considered = {} 

379 for layer in layers: 

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

381 num_layers = layer.nb_layer_ 

382 rec_considered = layer.recompute_considered_ 

383 

384 all_recomp = {"offset": 0} 

385 no_recomp = {"offset": 0} 

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

387 if rec_considered.get(rec, False): 

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

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

390 

391 self.simulate_yaml( 

392 yaml_format=all_recomp, 

393 show=True, 

394 interleave_num=self.num_of_interleave_, 

395 file_name=os.path.join(output_folder, 

396 "result_naive_all_recomp.svg"), 

397 ) 

398 

399 if num_layers % self.num_of_stage_ == 0: 

400 self.simulate_yaml( 

401 yaml_format=no_recomp, 

402 show=True, 

403 interleave_num=self.num_of_interleave_, 

404 file_name=os.path.join(output_folder, 

405 "result_naive_no_recomp.svg"), 

406 ) 

407 else: 

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

409 

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

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

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

413 check_yaml_depth_before_loading(fp) 

414 fp.seek(0) 

415 data = yaml.safe_load(fp) 

416 yaml_data = {} 

417 for manual in data.values(): 

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

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

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

421 yaml_data[Recompute.OFFSET] = [yaml_data[Recompute.OFFSET]] 

422 

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

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

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

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

427 yaml_data[rec] = [yaml_data[rec]] 

428 interleave_num = manual.get("interleave_num", 

429 self.num_of_interleave_) 

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

431 file_name = manual.get("file_name") 

432 full_file_name = os.path.join(output_folder, 

433 file_name) if (file_name) else None 

434 

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

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

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

438 try: 

439 simulate_result = self.simulate( 

440 show=False, 

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

442 sub_fig=sub_figs[0], 

443 ) 

444 except Exception: 

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

446 raise 

447 

448 if simulate_result is None: 

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

450 

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

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

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

454 if show: 

455 plt.show() 

456 

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

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

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

460 check_yaml_depth_before_loading(fp) 

461 fp.seek(0) 

462 data = yaml.safe_load(fp) 

463 yaml_data = {} 

464 for manual in data.values(): 

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

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

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

468 yaml_data[Recompute.OFFSET] = [yaml_data[Recompute.OFFSET]] 

469 

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

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

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

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

474 yaml_data[rec] = [yaml_data[rec]] 

475 interleave_num = manual.get("interleave_num", 

476 self.num_of_interleave_) 

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

478 file_name = manual.get("file_name") 

479 full_file_name = os.path.join(output_folder, 

480 file_name) if (file_name) else None 

481 

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

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

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

485 if show: 

486 plt.show() 

487 

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

489 interleave_num: int = 1, 

490 file_name: Optional[str] = None, 

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

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

493 layer_num = 0 

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

495 layer_num += layer.nb_layer_ 

496 nass = self.naive_layer_per_stage(layer_num, 

497 num_of_interleave=interleave_num) 

498 layer_per_recompute = Recompute.internal_from_yaml( 

499 interleave_num, self.num_of_stage_, yaml_format, nass) 

500 each_layer_per_recompute = self.split_layer_per_recompute(layer_per_recompute) 

501 return self.simulate_manual( 

502 each_layer_per_recompute, 

503 show, 

504 interleave_num=interleave_num, 

505 file_name=file_name, 

506 sub_fig=sub_fig 

507 ) 

508 

509 ####################################################################### 

510 ## ## 

511 ## Print Solver Model ## 

512 ## ## 

513 ####################################################################### 

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

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

516 act_mem_next = 0 

517 act_mem_curr = 0 

518 

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

520 for rec in Recompute.TYPE: 

521 if self.problem_.recompute_considered_[rec]: 

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

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

524 layer.memory_activation_rec_[rec]) 

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

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

527 layer.memory_activation_rec_[rec]) 

528 

529 return act_mem_next, act_mem_curr 

530 

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

532 """Solver memory model: parameter memory""" 

533 param_mem = 0 

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

535 if layer.memory_parameter_ is not None: 

536 param_mem += self._calculate_layer_parameter_memory( 

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

538 return param_mem 

539 

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

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

542 layer_mem = 0 

543 for inter in range(interleave_num): 

544 for rec in Recompute.TYPE: 

545 if self.problem_.recompute_considered_[rec]: 

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

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

548 return layer_mem 

549 

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

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

552 act_mem = 0 

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

554 for inter in range(interleave_num): 

555 for rec in Recompute.TYPE: 

556 if self.problem_.recompute_considered_[rec]: 

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

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

559 layer.memory_activation_rec_[rec] * 

560 activation_nums[inter][s]) 

561 return act_mem 

562 

563 

564 def debug_print_manual_theoretical_memory( 

565 self, 

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

567 interleave_num: int = 1) -> None: 

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

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

570 

571 if self.vpp_less_memory_: 

572 if self.seqpipe_: 

573 activation_nums = self.problem_.compute_activation_seq_nums( 

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

575 else: 

576 activation_nums = self.problem_.compute_less_activation_nums( 

577 self.num_of_stage_, interleave_num) 

578 # Add if dual to decide whether dualpipe_v is used 

579 elif self.dual_: 

580 activation_nums = self.problem_.compute_activation_nums_dual( 

581 self.num_of_stage_, interleave_num, self.num_of_micro_batch_) 

582 else: 

583 if self.seqpipe_: 

584 activation_nums = self.problem_.compute_activation_seq_nums( 

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

586 else: 

587 activation_nums = self.problem_.compute_activation_nums( 

588 self.num_of_stage_, interleave_num, self.num_of_micro_batch_) 

589 

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

591 

592 # compute for each stage 

593 for s in range(self.num_of_stage_): 

594 

595 # parameter memory 

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

597 

598 # head memory 

599 if s == 0: 

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

601 if head.memory_parameter_ is not None: 

602 param_mem += head.memory_parameter_ 

603 

604 # tail memory 

605 if s == self.num_of_stage_ - 1: 

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

607 if tail.memory_parameter_ is not None: 

608 param_mem += tail.memory_parameter_ 

609 

610 # act memory 

611 act_mem = self._calculate_activation_memory_solver(each_layer_per_recompute, s, 

612 interleave_num, activation_nums) 

613 

614 # overhead 

615 overhead = 0 

616 

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

618 

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

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

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

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

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

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

625 

626 def split_layer_per_recompute( 

627 self, 

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

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

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

631 each_layer_per_recompute = {} 

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

633 rest = layer.nb_layer_ 

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

635 for rec in Recompute.TYPE: 

636 for i in range(self.num_of_interleave_): 

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

638 for s in range(self.num_of_stage_): 

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

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

641 rest -= subtract 

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

643 return each_layer_per_recompute 

644 

645 def fuse_layer_per_recompute( 

646 self, 

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

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

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

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

651 for rec in Recompute.TYPE: 

652 for i in range(self.num_of_interleave_): 

653 all_layers_per_recompute[rec].append([]) 

654 for s in range(self.num_of_stage_): 

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

656 each_layer_per_recompute[layer][rec][i][s] 

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

658 )) 

659 return all_layers_per_recompute 

660 

661 

662 def simulate_manual( 

663 self, 

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

665 show: bool = True, 

666 interleave_num: int = 1, 

667 file_name: Optional[str] = None, 

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

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

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

671 

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

673 for rec in Recompute.TYPE: 

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

675 logger.error( 

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

677 layer, 

678 rec, 

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

680 interleave_num, 

681 ) 

682 return sys.maxsize 

683 

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

685 for rec in Recompute.TYPE: 

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