Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/auto_parallel/sapp_ppb/simulator/pp

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

7# http://www.apache.org/licenses/LICENSE-2.0

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"""Pipeline scheduler simulator: builds block dependencies, computes bubbles and peak memory."""

16from __future__ import annotations

18import copy

19import sys

21import numpy as np

23from hyper_parallel.auto_parallel.sapp_ppb.simulator.causal_error import CausalCommError, CausalError

24from hyper_parallel.auto_parallel.sapp_ppb.simulator.pipeline_builder import PipelineBuilder

25from hyper_parallel.auto_parallel.sapp_ppb.simulator.plot_manager import PlotMgr

26from hyper_parallel.auto_parallel.sapp_ppb.simulator.sim_block import BlockSim, RecBlockSim, SendBlockSim

27from hyper_parallel.auto_parallel.sapp_ppb.simulator.utils import apply_color, apply_format, format_2d_inputs

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

30sys.setrecursionlimit(8192)

33class PipelineSimulator:

34 r"""

35 Pipeline Simulator which provide pipeline flow process, bubbles and relative memories for stages.

37 Args:

38 block_time (Union[List[int|float], List[List[int|float]]]): Relative forward computing time for each block.

39 If it is List of List, the outer List indicates number of virtual-pp

40 while the inner List indicates pp_stage.

41 micro_num (int): Micro batch number.

42 comm_time (float, optional): Communication block (send/receive) time. Default: ``0.1``.

43 layer_recompute (Union[bool, List[int|float], List[List[int|float]]], optional): The block recompute

44 information.

45 If it is bool type, the backward block will be extended by block_time depending on whether it is True.

46 Otherwise it represents relative computing time of recompute for each block. Default: ``False``.

47 block_mem (Union[bool, List[int|float], List[List[int|float]]], optional): The block memory information.

48 If it is a number, the memory will be `block_mem` * `block_time`. Otherwise it represents relative memory

49 for each block. Default: ``1``.

50 backward_ratio (Union[List[int|float], List[List[int|float]]], optional): The ratios of backward computing

51 time and forward computing time for each block. Default: ``2``.

53 Example:

54 A PipelineSimulator with pp=4, micro=16, each stage has 8 layers and last stage has extra head and

55 loss computation equivalent to 0.8 layer:

56 >>> sim = PipelineSimulator([8,8,8,8+0.8], 16, comm_time=0.1) # create an instance of PipelineSimulator

57 >>> sim.run() # run simulation to scheduler the pipeline (information will be automatically printed)

58 ————————————— pp: 4, vp: 1, micro: 16 ————————————

59 -------------------- bubble --------------------

60 real = ideal + imba + comm

61 0.2658 = 0.1875 + 0.0615 + 0.0168

62 -------------------- memory --------------------

63 peak memory: 32.00, 24.00, 16.00, 8.80

64 >>> sim.show() # draw the pipeline and memory timeline picture

66 Show imbalance timeline of vp=2, pp=4, micro=8, total 16 layers with extra equivalent 1.2 layer:

67 >>> PipelineSimulator([[2,2,2,2],[1,2,3,2+1.2]], 8, comm_time=0.1).run().show()

68 ————————————— pp: 4, vp: 2, micro: 8 ————————————

69 -------------------- bubble --------------------

70 real = ideal + imba + comm

71 0.4971 = 0.1875 + 0.2447 + 0.0649

72 -------------------- memory --------------------

73 peak memory: 18.00, 18.00, 18.00, 14.80

75 Show timeline of vp=3, pp=8, micro=16, total 48 layers with extra equivalent 0.6 layer.

76 some of layers are recomputed and set memory correspondingly:

77 >>> PipelineSimulator([[2,2,2,2,2,2,2,2],

78 >>> [2,2,2,2,2,2,2,2],

79 >>> [2,2,2,2,2,2,2,2+0.6]], 16, 0.1,

80 >>> [[0,0,0,0,0,0,0,0],

81 >>> [1,0,0,0,0,0,0,0],

82 >>> [2,2,1,0,0,0,0,0]],

83 >>> [[2,2,2,2,2,2,2,2],

84 >>> [1.1,2,2,2,2,2,2,2],

85 >>> [0.2,0.2,1.1,2,2,2,2,2]]).run().show()

86 ————————————— pp: 8, vp: 3, micro: 16 ————————————

87 -------------------- bubble --------------------

88 real = ideal + imba + comm + recompute

89 0.4444 = 0.1458 + 0.1851 + 0.0724 + 0.0412

90 -------------------- memory --------------------

91 peak memory: 40.40, 43.60, 46.80, 50.00, 46.00, 42.00, 38.00, 34.00

93 Show timeline without comm for vp=2, pp=15, micro=16, total 96 layers with extra equivalent 1.2 layer:

94 >>> PipelineSimulator([[3,3,3,3,3,3,3,3,3,3,3,3,3,3,3],

95 >>> [3,3,3,3,3,3,3,3,4,4,4,4,4,4,3+1.2]], 16).run(False).show()

96 ————————————— pp:15, vp: 2, micro: 16 ————————————

97 -------------------- bubble --------------------

98 real = ideal + imba

99 0.5741 = 0.4375 + 0.1366

100 -------------------- memory --------------------

101 peak memory: 96.00, 96.00, 96.00, 96.00, 96.00, 96.00, 96.00, 93.00, 103.00, 97.00, 91.00,

102 85.00, 79.00, 73.00, 70.20

103 """

104 def __init__(self, block_time: list, micro_num: int, *args: object,

105 comm_time: float = 0.1,

106 layer_recompute: object = False, block_mem: object = 1,

107 block_mem_par: object = 0, constant_mem: float = 0,

108 backward_ratio: object = 2.,

109 backward_time: object = 0,

110 sub_fig: object = None, **kwargs: object) -> None:

111 """Delegate initialisation to :meth:`init` (kept as a named method for subclassing)."""

112 self.init(block_time, micro_num, comm_time, layer_recompute, block_mem,

113 block_mem_par, constant_mem, backward_ratio, backward_time,

114 sub_fig, *args, **kwargs)

115

116 # pylint: disable=W0613

117 def init(self, block_time: list, micro_num: int, comm_time: float,

118 layer_recompute: object, block_mem: object, block_mem_par: object,

119 constant_mem: float, backward_ratio: object, backward_time: object,

120 sub_fig: object, *args: object, **kwargs: object) -> None:

121 """Build the block grid, statistics and communication graph for the simulator."""

122 self.micro_num = micro_num

123 self.pp, self.vp = self._base_init(block_time)

124 self.block_num = 2 * self.vp * self.micro_num

125 self.comm_time = comm_time

126 self._input_format(block_time, layer_recompute, block_mem, block_mem_par,

127 backward_ratio, backward_time)

128 self.constant_mem = constant_mem

129 self._statistic_init()

130 self._comm = True

131 self.adjust_func_list = [self.swap_send_rec]

132 self.sub_fig = sub_fig

133 # Construct pipeline blocks

134 if self.vp == 1:

135 method = '1f1b'

136 else:

137 method = kwargs.get('method', 'vpp')

138 if self.micro_num >= self.pp:

139 self.adjust_func_list = [self.vpp_send_delay, self.residue_delay] + self.adjust_func_list

140 pp_builder = PipelineBuilder.get_builder(method)

141 self.blocks = [pp_builder(self.pp, self.micro_num, self.vp, p, self.block_time[:, p],

142 self.backward_time[:, p], self.block_mem[:, p], self.block_mem_par[:, p])

143 for p in range(self.pp)]

144

145 self._build_block() # create connection among compute blocks

146 self._build_comm_block() # create comm blocks for each compute block

147 self.peak_memory = None

148 self.end_time = None

149 self.lines = None

150 self.canvas = None

151

152 def run(self, comm: bool = True, print_info: bool = True) -> "PipelineSimulator":

153 """Run simulation to schedule the pipeline.

154

155 Args:

156 comm: Whether to build the pipeline considering communication dependency and time.

157 Default: ``True``.

158 print_info: Whether to automatically print bubble and memory information.

159 Default: ``True``.

160

161 Returns:

162 The current :class:`PipelineSimulator` instance (for chaining).

163

164 Raises:

165 CausalError: If the block sequences contain a dependency loop.

166 CausalCommError: If the block-with-comm sequences contain a dependency loop.

167 """

168 self._comm = comm

169 self._check_loop()

170 if comm:

171 self.lines = self._create_lines(*self.adjust_func_list)

172 self._check_comm_loop()

173 for b in range(self.block_num):

174 for p in range(self.pp):

175 self.blocks[p][b].build_with_comm()

176 self.lines[0][-1].build_with_comm()

177 else:

178 for p in range(self.pp):

179 for block in self.blocks[p]:

180 block.build_without_comm()

181 self._statistic_info()

182 if print_info:

183 self.print_info()

184 return self

185

186 def draw(self, comm: bool = True, connect: bool = None) -> "PipelineSimulator":

187 """Show the pipeline and memory timeline.

188

189 Args:

190 comm: Whether to show the comm blocks. Default: ``True``.

191 connect: Whether to show the connect arrow of the send-receive pair when the comm

192 pipeline is built. Default: ``None`` (auto-selected based on ``comm``).

193

194 Returns:

195 The current :class:`PipelineSimulator` instance (for chaining).

196 """

197 self.canvas = PlotMgr(2, ['block', 'memory'], sub_fig=self.sub_fig)

198 if self._comm:

199 connect = True if connect is None else connect

200 self.canvas.draw(self.lines, 0, comm, connect, False, 'timeline')

201 else:

202 connect = False if connect is None else connect

203 self.canvas.draw(self.blocks, 0, comm, connect, False, 'timeline')

204 self.canvas.draw_mem(self.states.get('block_mem_list', []), 1)

205 self.canvas.draw_info(self.bubbles, self.peak_memory)

206 return self

207

208

209 def show(self, comm: bool = True, connect: bool = None,

210 file_name: str = None) -> "PipelineSimulator":

211 """Draw the pipeline and display/save it via the canvas."""

212 draw_result = self.draw(comm, connect)

213 if draw_result is None:

214 raise RuntimeError("draw() returned None.")

215 self.canvas.show(file_name)

216 return self

217

218 def save(self, file_name: str, comm: bool = True,

219 connect: bool = None) -> "PipelineSimulator":

220 """Draw the pipeline and save it to ``file_name``."""

221 draw_result = self.draw(comm, connect)

222 if draw_result is None:

223 raise RuntimeError("draw() returned None.")

224 self.canvas.save(file_name)

225 return self

226

227 def print_info(self) -> "PipelineSimulator":

228 """Log bubble and peak memory information."""

229 bubble_colors = ['1;33', '1;32', '1;31', '1;35', '1;36']

230 header = '\033[1;37m' + '—' * 13 + \

231 f' pp:{self.pp:>2}, vp:{self.vp:>2}, micro:{self.micro_num:>3} ' + \

232 '—' * 12 + '\033[0m'

233 bubble_header = '-' * 20 + ' bubble ' + '-' * 20

234 bubble_keys = apply_format(apply_color(list(self.bubbles.keys()), bubble_colors))

235 bubble_values = apply_format(apply_color(list(self.bubbles.values()), bubble_colors))

236 memory_header = '-' * 20 + ' memory ' + '-' * 20

237 peak_memory = f"peak memory: {', '.join(f'{v:.2f}' for v in self.peak_memory)}"

238 logger.output(

239 "%s\n%s\n%s\n%s\n%s\n%s",

240 header, bubble_header, bubble_keys, bubble_values, memory_header, peak_memory,

241 )

242 return self

243

244 def _base_init(self, block_time) -> tuple:

245 r"""init base setting"""

246 if isinstance(block_time, (list, tuple)):

247 if all(isinstance(item, (list, tuple)) for item in block_time):

248 vp = len(block_time)

249 pp = len(block_time[0])

250 elif all(isinstance(item, (int, float)) for item in block_time):

251 vp = 1

252 pp = len(block_time)

253 else:

254 raise ValueError(f"Unsupported input format block_time: {block_time}")

255 else:

256 raise ValueError(f"Unsupported input format block_time: {block_time}")

257 if self.micro_num < pp:

258 raise ValueError(f" `micro_num`({self.micro_num}) should equal or larger than `pp`({pp})")

259 return pp, vp

260

261 def _input_format(self, block_time, layer_recompute, block_mem, block_mem_par,

262 backward_ratio, backward_time) -> None:

263 r"""format inputs as 2d array"""

264 self.block_time = format_2d_inputs(block_time, self.vp, self.pp)

265 if isinstance(layer_recompute, bool):

266 self.layer_recompute = self.block_time if layer_recompute else format_2d_inputs(0, self.vp, self.pp)

267 else:

268 self.layer_recompute = format_2d_inputs(layer_recompute, self.vp, self.pp)

269 if isinstance(block_mem, (int, float)):

270 self.block_mem = self.block_time * block_mem

271 else:

272 self.block_mem = format_2d_inputs(block_mem, self.vp, self.pp)

273

274 if isinstance(block_mem_par, (int, float)):

275 self.block_mem_par = self.block_time * block_mem_par

276 else:

277 self.block_mem_par = format_2d_inputs(block_mem_par, self.vp, self.pp)

278

279 self.backward_ratio = format_2d_inputs(backward_ratio, self.vp, self.pp)

280

281 if isinstance(backward_time, (int, float)) and backward_time == 0:

282 self._provided_backward_time = None

283 else:

284 self._provided_backward_time = format_2d_inputs(backward_time, self.vp, self.pp)

285

286 def _statistic_init(self) -> None:

287 r"""init statistic info"""

288 self.forward_time = self.block_time

289 if self._provided_backward_time is not None:

290 self.backward_time = self._provided_backward_time

291 else:

292 self.backward_time = self.block_time * self.backward_ratio + self.layer_recompute

293 self.states = {'last_time': np.zeros(self.pp),

294 'warmup_time': np.zeros(self.pp),

295 'cooldown_time': np.zeros(self.pp),

296 'stable_free_time': (np.zeros((self.vp, self.pp)), np.zeros((self.vp, self.pp))),

297 'block_mem_list': [np.array([[0, 0]]) for _ in range(self.pp)]}

298 self.model_compute_time = (np.sum(self.forward_time) + \

299 np.sum(self.backward_time - self.layer_recompute)) * self.micro_num

300 self.hardware_compute_time = (np.sum(self.forward_time) + np.sum(self.backward_time)) * self.micro_num

301 self.bubbles = {'real': 0,

302 'ideal': (self.pp - 1) / self.vp / self.micro_num,

303 'imba': 0,

304 'comm': 0}

305 if np.sum(self.layer_recompute) > 1e-5:

306 self.bubbles['recompute'] = self.hardware_compute_time / self.model_compute_time - 1

307 p, v, m = self.pp, self.vp, self.micro_num

308 if self.vp == 1:

309 if self.pp == 2:

310 self.bubbles['comm'] = 4 * m

311 elif self.pp % 2 == 0:

312 self.bubbles['comm'] = 4 * p * m + 4 * p ** 2 - 14 * p

313 else:

314 self.bubbles['comm'] = 4 * p * m + 4 * p ** 2 - 12 * p

315 elif self.pp <= 5:

316 comm_coef_list = [[4, -2, 0], [6, -2, -6], [4, 0, 12], [6, -2, 40]]

317 self.bubbles['comm'] = np.dot(np.array([p * v * m, m * p, 1]), comm_coef_list[self.pp - 2])

318 elif self.pp % 2 == 0:

319 self.bubbles['comm'] = 4 * p * v * m + 4 * p ** 2 - 13 * p

320 else:

321 self.bubbles['comm'] = 6 * p * v * m - 2 * v * p ** 2 + 4 * v * p - 2 * p * m + 6 * p ** 2 - 16 * p

322

323 self.bubbles['comm'] *= self.comm_time / self.model_compute_time

324

325 def _update_block_mem(self, block, current_mem, p):

326 r"""Update memory for one block and return (updated_block, current_mem) or None if skipped.

327

328 Args:

329 block: A simulation block (compute, recompute, or communication).

330 current_mem: Current memory usage in MB for stage *p*.

331 p: Stage index.

332

333 Returns:

334 ``(block, current_mem)`` if the block affected memory, or

335 ``(None, current_mem)`` if the block was skipped.

336 """

337 if block.type == 'c' and block.state == 'f':

338 current_mem += block.mem

339 elif block.type == 'c' and block.state == 'b':

340 if not self._comm or not block.rec_block:

341 current_mem -= block.mem

342 else:

343 return None, current_mem

344 elif block.type == 'r' and block.host.state == 'b':

345 current_mem -= block.host.mem

346 block = block.host

347 else:

348 return None, current_mem

349 self.states['block_mem_list'][p] = np.append(self.states['block_mem_list'][p],

350 np.array([[block.end, current_mem]]), axis=0)

351 return block, current_mem

352

353 def _statistic_info(self) -> None:

354 r"""Compute per-stage peak memory and pipeline step time.

355

356 Memory accounting starts from ``constant_mem + first_compute_block.mem_par``

357 (the parameter memory of the first compute block in the stage timeline).

358 When communication blocks are enabled (``self._comm``), the first block

359 in ``self.lines[p]`` may be a receive block rather than a compute block;

360 in that case using ``blocks[0].mem_par`` would be incorrect. This

361 implementation explicitly finds the first compute block to ensure the

362 initial parameter memory is accounted correctly.

363

364 .. note::

365 For pipelines with ``comm=True`` where the first block in a stage's

366 timeline is a communication block, the reported ``peak_memory`` may

367 differ from earlier versions that used ``blocks[0].mem_par``.

368 This is a correctness fix — the previous behavior was accidentally

369 using a communication block's parameter memory instead of a compute

370 block's.

371 """

372 for p in range(self.pp):

373 blocks = self.lines[p] if self._comm else self.blocks[p]

374 first_compute = next((b for b in blocks if b.type == 'c'), None)

375 current_mem = self.constant_mem + (first_compute.mem_par if first_compute else 0)

376

377 for block in blocks:

378 _, current_mem = self._update_block_mem(block, current_mem, p)

379 self.states['block_mem_list'][p] = np.append(self.states['block_mem_list'][p],

380 np.array([[blocks[-1].end, current_mem]]), axis=0)

381 self.peak_memory = [np.max((self.states['block_mem_list'][p].T)[1]) for p in range(self.pp)]

382 self.end_time = max(np.max((self.states['block_mem_list'][p].T)[0]) for p in range(self.pp))

383 self.bubbles['real'] = (self.pp * self.end_time - self.model_compute_time) / self.model_compute_time

384 self.bubbles['imba'] = self.bubbles['real'] - self.bubbles['ideal'] + 1e-10

385 if not self._comm:

386 self.bubbles.pop('comm')

387 else:

388 self.bubbles['imba'] -= self.bubbles['comm']

389 if self.bubbles.get('recompute'):

390 self.bubbles['imba'] -= self.bubbles['recompute']

391

392 def _get_pre_label(self, label: tuple) -> tuple:

393 r"""get pre block label"""

394 t, s, m, v, p = label

395 if (s, v, p) == ('f', 0, 0):

396 return ('h', p)

397 if (s, p) == ('f', 0):

398 res = (t, s, m, v - 1, self.pp - 1)

399 return res

400 if (s, p) == ('b', self.pp - 1):

401 if v == self.vp - 1:

402 res = (t, 'f', m, self.vp - 1, p)

403 return res

404 res = (t, s, m, v + 1, 0)

405 return res

406 if s == 'f':

407 res = (t, s, m, v, p - 1)

408 return res

409 if s == 'b':

410 res = (t, s, m, v, p + 1)

411 return res

412 raise ValueError(f"Illegal label: {label}")

413

414 def _build_block(self) -> None:

415 r"""Build `pre` relation for computation blocks."""

416 books = {self.blocks[0][0].pre.label: self.blocks[0][0].pre}

417 for p in range(self.pp):

418 for item in self.blocks[p]:

419 books[item.label] = item

420 for p in range(self.pp):

421 block = self.blocks[p][0]

422 while block is not None:

423 pre_label = self._get_pre_label(block.label)

424 block.pre = books.get(pre_label, None)

425 block = block.right

426

427 def _build_comm_block(self) -> None:

428 r"""Build `send_block` and `rec_block` relation among a computation block and two comm blocks."""

429 for p in range(self.pp):

430 block = self.blocks[p][0]

431 while block is not None:

432 pre = block.pre

433 if pre.stage != block.stage:

434 block.rec_block = RecBlockSim(p, block.state, block.id, block.chunk, self.comm_time)

435 pre.send_block = SendBlockSim(pre.stage, pre.state, pre.id, pre.chunk, self.comm_time)

436 block.rec_block.host = block

437 block.rec_block.dual = pre.send_block

438 pre.send_block.host = pre

439 pre.send_block.dual = block.rec_block

440 block.depend_pre = block.rec_block

441 block.rec_block.depend_pre = pre.send_block

442 pre.send_block.depend_pre = pre

443 else:

444 block.depend_pre = pre

445 block = block.right

446

447 def _check_loop(self) -> None:

448 r"""check the existence of dependency"""

449 loop = self.blocks[0][-1].loop()

450 if loop:

451 raise CausalError('Block dependency exist loops!', self.blocks, loop)

452 for p in range(self.pp):

453 for block in self.blocks[p]:

454 block.flag = False

455

456 def _check_comm_loop(self) -> None:

457 r"""check the existence of comm dependency"""

458 loop = self.lines[0][-1].comm_loop()

459 if loop:

460 raise CausalCommError('Block comm dependency exist loops!', self.lines, loop)

461 for p in range(self.pp):

462 for block in self.lines[p]:

463 block.flag = False

464

465 def _create_lines(self, *adjust_func) -> list[list[BlockSim]]:

466 r"""create block line for each stage with comm"""

467 lines = [copy.copy(self.blocks[p]) for p in range(self.pp)]

468 for p in range(self.pp):

469 for b in range(self.block_num):

470 block = self.blocks[p][b]

471 pre = block.pre

472 if block.rec_block:

473 lines[p].insert(lines[p].index(block), block.rec_block)

474 if pre.type == 'h':

475 lines[pre.stage].insert(0, pre.send_block)

476 else:

477 lines[pre.stage].insert(lines[pre.stage].index(pre) + 1, pre.send_block)

478 for func in adjust_func:

479 lines = func(lines)

480 for p in range(self.pp):

481 for b, block in enumerate(lines[p]):

482 if b == 0:

483 block.depend_left = block.left if block.left else block.host.left

484 else:

485 block.depend_left = lines[p][b - 1]

486 return lines

487

488 def _get_block_phase(self, p: int, b: int) -> str:

489 r"""get block phase"""

490 r = self.micro_num % self.pp

491 if b < (self.vp + 1) * self.pp - 2 * p - 2 + r:

492 return 'warmup'

493 if b > self.block_num - (self.vp + 1) * self.pp + 2 * p:

494 return 'cooldown'

495 return 'stable'

496

497 def _send_block_delay(self, lines, p: int, b: int, distance: int) -> None:

498 r"""adjust send block: delay send block"""

499 i_send = lines[p].index(self.blocks[p][b].send_block)

500 send_block = lines[p].pop(i_send)

501 i_new = lines[p].index(self.blocks[p][b + distance]) + 1

502 lines[p].insert(i_new, send_block)

503

504 def _process_swap_gap3(self, block, lines, p, b, i_b):

505 r"""process swap when gap == 3."""

506 if p % 2 == 0 and lines[p][i_b + 1].type == 'r' and lines[p][i_b + 2].type == 's':

507 lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]

508 if p % 2 == 1 and lines[p][i_b + 1].type == 's' and lines[p][i_b + 2].type == 'r':

509 if block.phase == 'warmup' and self.blocks[p][b + 1].phase == 'cooldown':

510 return False

511 lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]

512 if lines[p][i_b + 1].dual.stage == lines[p][i_b + 2].dual.stage:

513 pd = lines[p][i_b + 1].dual.stage

514 j_b1 = lines[pd].index(lines[p][i_b + 1].dual)

515 j_b2 = lines[pd].index(lines[p][i_b + 2].dual)

516 if j_b1 > j_b2:

517 lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]

518 return True

519

520 def _process_swap_gap4(self, lines, p, i_b):

521 r"""process swap when gap == 4."""

522 if lines[p][i_b + 1].dual.stage == lines[p][i_b + 2].dual.stage and \

523 lines[p][i_b + 2].dual.stage == lines[p][i_b + 3].dual.stage:

524 if lines[p][i_b + 1].type == 's' and lines[p][i_b + 2].type == 's' \

525 and lines[p][i_b + 3].type == 'r':

526 lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]

527

528 def _process_swap(self, block, lines, p, b, i_b, i_bn) -> bool:

529 r"""process swap in condition"""

530 if i_bn - i_b == 3:

531 return self._process_swap_gap3(block, lines, p, b, i_b)

532 if i_bn - i_b == 4:

533 self._process_swap_gap4(lines, p, i_b)

534 return True

535

536 def swap_send_rec(self, lines: list[list[BlockSim]]) -> list[list[BlockSim]]:

537 """Adjust send blocks: swap adjacent send/receive pairs where ordering is ambiguous."""

538 for p in range(self.pp):

539 for b, block in enumerate(self.blocks[p]):

540 if b >= len(self.blocks[p]) - 1:

541 continue

542 i_b = lines[p].index(block)

543 i_bn = lines[p].index(self.blocks[p][b + 1])

544 try:

545 swap_processed = self._process_swap(block, lines, p, b, i_b, i_bn)

546 except (ValueError, IndexError) as error:

547 raise RuntimeError(

548 "Failed to process swap in pipeline simulator."

549 ) from error

550 if not swap_processed:

551 continue

552 return lines

553

554 def vpp_send_delay(self, lines: list[list[BlockSim]]) -> list[list[BlockSim]]:

555 """Adjust VPP send blocks by delaying them one slot during the stable phase."""

556 if self.micro_num % self.pp != 0:

557 return lines

558 for p in range(self.pp):

559 for b, block in enumerate(self.blocks[p]):

560 if block.send_block is not None and block.phase == 'stable':

561 self._send_block_delay(lines, p, b, 1)

562 return lines

563

564 def residue_delay(self, lines: list[list[BlockSim]]) -> list[list[BlockSim]]:

565 """Adjust send blocks when ``micro_num % pp`` leaves a residue micro-batch."""

566 r = self.micro_num % self.pp

567 if r == 0:

568 return lines

569 for p in range(self.pp):

570 for b, block in enumerate(self.blocks[p]):

571 if block.send_block is None:

572 continue

573 if p == self.pp - 1 and block.id < self.pp + r and block.state == 'f':

574 self._send_block_delay(lines, -1, b, r + max(0, block.id - self.pp + 1))

575 elif p == 0 and block.id < self.pp + r and block.state == 'b':

576 if self.micro_num // self.pp == 1:

577 self._send_block_delay(lines, 0, b, r)

578 else:

579 self._send_block_delay(lines, 0, b, r + self.pp)

580 elif block.phase == 'stable':

581 self._send_block_delay(lines, p, b, 1)

582 return lines

583

584

585if __name__ == '__main__':

586

587 # PipelineSimulator([[4, 4, 4, 4], [4, 4, 4, 4], [4, 4, 4, 4 + 0.8]], 8, 0.1,

588 # [[1, 0, 0, 0], [1, 0, 0, 0], [1, 1, 0, 0]],

589 # [[1.1, 2, 2, 2], [1.1, 2, 2, 2], [1.1, 1.1, 2, 2]], method='vpp').run().show()

590 PipelineSimulator(

591 [[186.0, 171.0, 132.0, 132.0, 132.0, 132.0, 132.0, 132.0,

592 132.0, 132.0, 132.0, 132.0, 132.0, 132.0, 132.0, 133.0]], 32,

593 block_mem_act=[[1146, 908, 736, 736, 736, 736, 736, 736,

594 736, 736, 2623, 2623, 4510, 4510, 8284, 8284]],

595 block_mem_par=[[14130, 21126, 36252, 36252, 36252, 36252, 36252, 36252,

596 36252, 36252, 36252, 36252, 36252, 36252, 36252, 38297]],

597 layer_recompute=[[135.0, 171.0, 132.0, 132.0, 132.0, 132.0, 132.0, 132.0,

598 132.0, 132.0, 99.0, 99.0, 66.0, 66.0, 0, 0]],

599 less_memory=False).run().show()

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / simulator / pp_simulator.py: 82%

328 statements