Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / simulator / pp_simulator.py: 80%
313 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
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"""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 sub_fig: object = None, **kwargs: object) -> None:
110 """Delegate initialisation to :meth:`init` (kept as a named method for subclassing)."""
111 self.init(block_time, micro_num, comm_time, layer_recompute, block_mem,
112 block_mem_par, constant_mem, backward_ratio, sub_fig, *args, **kwargs)
114 # pylint: disable=W0613
115 def init(self, block_time: list, micro_num: int, comm_time: float,
116 layer_recompute: object, block_mem: object, block_mem_par: object,
117 constant_mem: float, backward_ratio: object,
118 sub_fig: object, *args: object, **kwargs: object) -> None:
119 """Build the block grid, statistics and communication graph for the simulator."""
120 self.micro_num = micro_num
121 self.pp, self.vp = self._base_init(block_time)
122 self.block_num = 2 * self.vp * self.micro_num
123 self.comm_time = comm_time
124 self._input_format(block_time, layer_recompute, block_mem, block_mem_par, backward_ratio)
125 self.constant_mem = constant_mem
126 self._statistic_init()
127 self._comm = True
128 self.adjust_func_list = [self.swap_send_rec]
129 self.sub_fig = sub_fig
130 # Construct pipeline blocks
131 if self.vp == 1:
132 method = '1f1b'
133 else:
134 method = kwargs.get('method', 'vpp')
135 if self.micro_num >= self.pp:
136 self.adjust_func_list = [self.vpp_send_delay, self.residue_delay] + self.adjust_func_list
137 pp_builder = PipelineBuilder.get_builder(method)
138 self.blocks = [pp_builder(self.pp, self.micro_num, self.vp, p, self.block_time[:, p],
139 self.backward_time[:, p], self.block_mem[:, p], self.block_mem_par[:, p])
140 for p in range(self.pp)]
142 self._build_block() # create connection among compute blocks
143 self._build_comm_block() # create comm blocks for each compute block
144 self.peak_memory = None
145 self.end_time = None
146 self.lines = None
147 self.canvas = None
149 def run(self, comm: bool = True, print_info: bool = True) -> "PipelineSimulator":
150 """Run simulation to schedule the pipeline.
152 Args:
153 comm: Whether to build the pipeline considering communication dependency and time.
154 Default: ``True``.
155 print_info: Whether to automatically print bubble and memory information.
156 Default: ``True``.
158 Returns:
159 The current :class:`PipelineSimulator` instance (for chaining).
161 Raises:
162 CausalError: If the block sequences contain a dependency loop.
163 CausalCommError: If the block-with-comm sequences contain a dependency loop.
164 """
165 self._comm = comm
166 self._check_loop()
167 if comm:
168 self.lines = self._create_lines(*self.adjust_func_list)
169 self._check_comm_loop()
170 for b in range(self.block_num):
171 for p in range(self.pp):
172 self.blocks[p][b].build_with_comm()
173 self.lines[0][-1].build_with_comm()
174 else:
175 for p in range(self.pp):
176 for block in self.blocks[p]:
177 block.build_without_comm()
178 self._statistic_info()
179 if print_info:
180 self.print_info()
181 return self
183 def draw(self, comm: bool = True, connect: bool = None) -> "PipelineSimulator":
184 """Show the pipeline and memory timeline.
186 Args:
187 comm: Whether to show the comm blocks. Default: ``True``.
188 connect: Whether to show the connect arrow of the send-receive pair when the comm
189 pipeline is built. Default: ``None`` (auto-selected based on ``comm``).
191 Returns:
192 The current :class:`PipelineSimulator` instance (for chaining).
193 """
194 self.canvas = PlotMgr(2, ['block', 'memory'], sub_fig=self.sub_fig)
195 if self._comm:
196 connect = True if connect is None else connect
197 self.canvas.draw(self.lines, 0, comm, connect, False, 'timeline')
198 else:
199 connect = False if connect is None else connect
200 self.canvas.draw(self.blocks, 0, comm, connect, False, 'timeline')
201 self.canvas.draw_mem(self.states.get('block_mem_list', []), 1)
202 self.canvas.draw_info(self.bubbles, self.peak_memory)
203 return self
206 def show(self, comm: bool = True, connect: bool = None,
207 file_name: str = None) -> "PipelineSimulator":
208 """Draw the pipeline and display/save it via the canvas."""
209 draw_result = self.draw(comm, connect)
210 if draw_result is None:
211 raise RuntimeError("draw() returned None.")
212 self.canvas.show(file_name)
213 return self
215 def save(self, file_name: str, comm: bool = True,
216 connect: bool = None) -> "PipelineSimulator":
217 """Draw the pipeline and save it to ``file_name``."""
218 draw_result = self.draw(comm, connect)
219 if draw_result is None:
220 raise RuntimeError("draw() returned None.")
221 self.canvas.save(file_name)
222 return self
224 def print_info(self) -> "PipelineSimulator":
225 """Log bubble and peak memory information."""
226 bubble_colors = ['1;33', '1;32', '1;31', '1;35', '1;36']
227 header = '\033[1;37m' + '—' * 13 + \
228 f' pp:{self.pp:>2}, vp:{self.vp:>2}, micro:{self.micro_num:>3} ' + \
229 '—' * 12 + '\033[0m'
230 bubble_header = '-' * 20 + ' bubble ' + '-' * 20
231 bubble_keys = apply_format(apply_color(list(self.bubbles.keys()), bubble_colors))
232 bubble_values = apply_format(apply_color(list(self.bubbles.values()), bubble_colors))
233 memory_header = '-' * 20 + ' memory ' + '-' * 20
234 peak_memory = f"peak memory: {', '.join(f'{v:.2f}' for v in self.peak_memory)}"
235 logger.output(
236 "%s\n%s\n%s\n%s\n%s\n%s",
237 header, bubble_header, bubble_keys, bubble_values, memory_header, peak_memory,
238 )
239 return self
241 def _base_init(self, block_time) -> tuple:
242 r"""init base setting"""
243 if isinstance(block_time, (list, tuple)):
244 if all(isinstance(item, (list, tuple)) for item in block_time):
245 vp = len(block_time)
246 pp = len(block_time[0])
247 elif all(isinstance(item, (int, float)) for item in block_time):
248 vp = 1
249 pp = len(block_time)
250 else:
251 raise ValueError(f"Unsupported input format block_time: {block_time}")
252 else:
253 raise ValueError(f"Unsupported input format block_time: {block_time}")
254 if self.micro_num < pp:
255 raise ValueError(f" `micro_num`({self.micro_num}) should equal or larger than `pp`({pp})")
256 return pp, vp
258 def _input_format(self, block_time, layer_recompute, block_mem, block_mem_par, backward_ratio) -> None:
259 r"""format inputs as 2d array"""
260 self.block_time = format_2d_inputs(block_time, self.vp, self.pp)
261 if isinstance(layer_recompute, bool):
262 self.layer_recompute = self.block_time if layer_recompute else format_2d_inputs(0, self.vp, self.pp)
263 else:
264 self.layer_recompute = format_2d_inputs(layer_recompute, self.vp, self.pp)
265 if isinstance(block_mem, (int, float)):
266 self.block_mem = self.block_time * block_mem
267 else:
268 self.block_mem = format_2d_inputs(block_mem, self.vp, self.pp)
270 if isinstance(block_mem_par, (int, float)):
271 self.block_mem_par = self.block_time * block_mem_par
272 else:
273 self.block_mem_par = format_2d_inputs(block_mem_par, self.vp, self.pp)
275 self.backward_ratio = format_2d_inputs(backward_ratio, self.vp, self.pp)
277 def _statistic_init(self) -> None:
278 r"""init statistic info"""
279 self.forward_time = self.block_time
280 self.backward_time = self.block_time * self.backward_ratio + self.layer_recompute
281 self.states = {'last_time': np.zeros(self.pp),
282 'warmup_time': np.zeros(self.pp),
283 'cooldown_time': np.zeros(self.pp),
284 'stable_free_time': (np.zeros((self.vp, self.pp)), np.zeros((self.vp, self.pp))),
285 'block_mem_list': [np.array([[0, 0]]) for _ in range(self.pp)]}
286 self.model_compute_time = (np.sum(self.forward_time) + \
287 np.sum(self.forward_time * self.backward_ratio)) * self.micro_num
288 self.hardware_compute_time = (np.sum(self.forward_time) + np.sum(self.backward_time)) * self.micro_num
289 self.bubbles = {'real': 0,
290 'ideal': (self.pp - 1) / self.vp / self.micro_num,
291 'imba': 0,
292 'comm': 0}
293 if np.sum(self.layer_recompute) > 1e-5:
294 self.bubbles['recompute'] = self.hardware_compute_time / self.model_compute_time - 1
295 p, v, m = self.pp, self.vp, self.micro_num
296 if self.vp == 1:
297 if self.pp == 2:
298 self.bubbles['comm'] = 4 * m
299 elif self.pp % 2 == 0:
300 self.bubbles['comm'] = 4 * p * m + 4 * p ** 2 - 14 * p
301 else:
302 self.bubbles['comm'] = 4 * p * m + 4 * p ** 2 - 12 * p
303 elif self.pp <= 5:
304 comm_coef_list = [[4, -2, 0], [6, -2, -6], [4, 0, 12], [6, -2, 40]]
305 self.bubbles['comm'] = np.dot(np.array([p * v * m, m * p, 1]), comm_coef_list[self.pp - 2])
306 elif self.pp % 2 == 0:
307 self.bubbles['comm'] = 4 * p * v * m + 4 * p ** 2 - 13 * p
308 else:
309 self.bubbles['comm'] = 6 * p * v * m - 2 * v * p ** 2 + 4 * v * p - 2 * p * m + 6 * p ** 2 - 16 * p
311 self.bubbles['comm'] *= self.comm_time / self.model_compute_time
313 def _statistic_info(self) -> None:
314 r"""compute statistic info"""
315 for p in range(self.pp):
316 blocks = self.lines[p] if self._comm else self.blocks[p]
317 current_mem = self.constant_mem + blocks[0].mem_par
319 for block in blocks:
320 if block.type == 'c' and block.state == 'f':
321 current_mem += block.mem
322 elif block.type == 'c' and block.state == 'b':
323 if not self._comm or not block.rec_block:
324 current_mem -= block.mem
325 elif block.type == 'r' and block.host.state == 'b':
326 current_mem -= block.host.mem
327 block = block.host
328 else:
329 continue
330 self.states['block_mem_list'][p] = np.append(self.states['block_mem_list'][p],
331 np.array([[block.end, current_mem]]), axis=0)
332 self.states['block_mem_list'][p] = np.append(self.states['block_mem_list'][p],
333 np.array([[blocks[-1].end, current_mem]]), axis=0)
334 self.peak_memory = [np.max((self.states['block_mem_list'][p].T)[1]) for p in range(self.pp)]
335 self.end_time = max(np.max((self.states['block_mem_list'][p].T)[0]) for p in range(self.pp))
336 self.bubbles['real'] = (self.pp * self.end_time - self.model_compute_time) / self.model_compute_time
337 self.bubbles['imba'] = self.bubbles['real'] - self.bubbles['ideal'] + 1e-10
338 if not self._comm:
339 self.bubbles.pop('comm')
340 else:
341 self.bubbles['imba'] -= self.bubbles['comm']
342 if self.bubbles.get('recompute'):
343 self.bubbles['imba'] -= self.bubbles['recompute']
345 def _get_pre_label(self, label: tuple) -> tuple:
346 r"""get pre block label"""
347 t, s, m, v, p = label
348 if (s, v, p) == ('f', 0, 0):
349 return ('h', p)
350 if (s, p) == ('f', 0):
351 res = (t, s, m, v - 1, self.pp - 1)
352 return res
353 if (s, p) == ('b', self.pp - 1):
354 if v == self.vp - 1:
355 res = (t, 'f', m, self.vp - 1, p)
356 return res
357 res = (t, s, m, v + 1, 0)
358 return res
359 if s == 'f':
360 res = (t, s, m, v, p - 1)
361 return res
362 if s == 'b':
363 res = (t, s, m, v, p + 1)
364 return res
365 raise ValueError(f"Illegal label: {label}")
367 def _build_block(self) -> None:
368 r"""Build `pre` relation for computation blocks."""
369 books = {self.blocks[0][0].pre.label: self.blocks[0][0].pre}
370 for p in range(self.pp):
371 for item in self.blocks[p]:
372 books[item.label] = item
373 for p in range(self.pp):
374 block = self.blocks[p][0]
375 while block is not None:
376 pre_label = self._get_pre_label(block.label)
377 block.pre = books.get(pre_label, None)
378 block = block.right
380 def _build_comm_block(self) -> None:
381 r"""Build `send_block` and `rec_block` relation among a computation block and two comm blocks."""
382 for p in range(self.pp):
383 block = self.blocks[p][0]
384 while block is not None:
385 pre = block.pre
386 if pre.stage != block.stage:
387 block.rec_block = RecBlockSim(p, block.state, block.id, block.chunk, self.comm_time)
388 pre.send_block = SendBlockSim(pre.stage, pre.state, pre.id, pre.chunk, self.comm_time)
389 block.rec_block.host = block
390 block.rec_block.dual = pre.send_block
391 pre.send_block.host = pre
392 pre.send_block.dual = block.rec_block
393 block.depend_pre = block.rec_block
394 block.rec_block.depend_pre = pre.send_block
395 pre.send_block.depend_pre = pre
396 else:
397 block.depend_pre = pre
398 block = block.right
400 def _check_loop(self) -> None:
401 r"""check the existence of dependency"""
402 loop = self.blocks[0][-1].loop()
403 if loop:
404 raise CausalError('Block dependency exist loops!', self.blocks, loop)
405 for p in range(self.pp):
406 for block in self.blocks[p]:
407 block.flag = False
409 def _check_comm_loop(self) -> None:
410 r"""check the existence of comm dependency"""
411 loop = self.lines[0][-1].comm_loop()
412 if loop:
413 raise CausalCommError('Block comm dependency exist loops!', self.lines, loop)
414 for p in range(self.pp):
415 for block in self.lines[p]:
416 block.flag = False
418 def _create_lines(self, *adjust_func) -> list[list[BlockSim]]:
419 r"""create block line for each stage with comm"""
420 lines = [copy.copy(self.blocks[p]) for p in range(self.pp)]
421 for p in range(self.pp):
422 for b in range(self.block_num):
423 block = self.blocks[p][b]
424 pre = block.pre
425 if block.rec_block:
426 lines[p].insert(lines[p].index(block), block.rec_block)
427 if pre.type == 'h':
428 lines[pre.stage].insert(0, pre.send_block)
429 else:
430 lines[pre.stage].insert(lines[pre.stage].index(pre) + 1, pre.send_block)
431 for func in adjust_func:
432 lines = func(lines)
433 for p in range(self.pp):
434 for b, block in enumerate(lines[p]):
435 if b == 0:
436 block.depend_left = block.left if block.left else block.host.left
437 else:
438 block.depend_left = lines[p][b - 1]
439 return lines
441 def _get_block_phase(self, p: int, b: int) -> str:
442 r"""get block phase"""
443 r = self.micro_num % self.pp
444 if b < (self.vp + 1) * self.pp - 2 * p - 2 + r:
445 return 'warmup'
446 if b > self.block_num - (self.vp + 1) * self.pp + 2 * p:
447 return 'cooldown'
448 return 'stable'
450 def _send_block_delay(self, lines, p: int, b: int, distance: int) -> None:
451 r"""adjust send block: delay send block"""
452 i_send = lines[p].index(self.blocks[p][b].send_block)
453 send_block = lines[p].pop(i_send)
454 i_new = lines[p].index(self.blocks[p][b + distance]) + 1
455 lines[p].insert(i_new, send_block)
457 def _process_swap(self, block, lines, p, b, i_b, i_bn) -> bool:
458 r"""process swap in condition"""
459 if i_bn - i_b == 3:
460 if p % 2 == 0 and lines[p][i_b + 1].type == 'r' and lines[p][i_b + 2].type == 's':
461 lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]
462 if p % 2 == 1 and lines[p][i_b + 1].type == 's' and lines[p][i_b + 2].type == 'r':
463 if block.phase == 'warmup' and self.blocks[p][b + 1].phase == 'cooldown':
464 return False
465 lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]
466 if lines[p][i_b + 1].dual.stage == lines[p][i_b + 2].dual.stage:
467 pd = lines[p][i_b + 1].dual.stage
468 j_b1 = lines[pd].index(lines[p][i_b + 1].dual)
469 j_b2 = lines[pd].index(lines[p][i_b + 2].dual)
470 if j_b1 > j_b2:
471 lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]
472 if i_bn - i_b == 4:
473 if lines[p][i_b + 1].dual.stage == lines[p][i_b + 2].dual.stage and \
474 lines[p][i_b + 2].dual.stage == lines[p][i_b + 3].dual.stage:
475 if lines[p][i_b + 1].type == 's' and lines[p][i_b + 2].type == 's' \
476 and lines[p][i_b + 3].type == 'r':
477 lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]
478 return True
480 def swap_send_rec(self, lines: list[list[BlockSim]]) -> list[list[BlockSim]]:
481 """Adjust send blocks: swap adjacent send/receive pairs where ordering is ambiguous."""
482 for p in range(self.pp):
483 for b, block in enumerate(self.blocks[p]):
484 if b >= len(self.blocks[p]) - 1:
485 continue
486 i_b = lines[p].index(block)
487 i_bn = lines[p].index(self.blocks[p][b + 1])
488 try:
489 swap_processed = self._process_swap(block, lines, p, b, i_b, i_bn)
490 except (ValueError, IndexError) as error:
491 raise RuntimeError(
492 "Failed to process swap in pipeline simulator."
493 ) from error
494 if not swap_processed:
495 continue
496 return lines
498 def vpp_send_delay(self, lines: list[list[BlockSim]]) -> list[list[BlockSim]]:
499 """Adjust VPP send blocks by delaying them one slot during the stable phase."""
500 if self.micro_num % self.pp != 0:
501 return lines
502 for p in range(self.pp):
503 for b, block in enumerate(self.blocks[p]):
504 if block.send_block is not None and block.phase == 'stable':
505 self._send_block_delay(lines, p, b, 1)
506 return lines
508 def residue_delay(self, lines: list[list[BlockSim]]) -> list[list[BlockSim]]:
509 """Adjust send blocks when ``micro_num % pp`` leaves a residue micro-batch."""
510 r = self.micro_num % self.pp
511 if r == 0:
512 return lines
513 for p in range(self.pp):
514 for b, block in enumerate(self.blocks[p]):
515 if block.send_block is None:
516 continue
517 if p == self.pp - 1 and block.id < self.pp + r and block.state == 'f':
518 self._send_block_delay(lines, -1, b, r + max(0, block.id - self.pp + 1))
519 elif p == 0 and block.id < self.pp + r and block.state == 'b':
520 if self.micro_num // self.pp == 1:
521 self._send_block_delay(lines, 0, b, r)
522 else:
523 self._send_block_delay(lines, 0, b, r + self.pp)
524 elif block.phase == 'stable':
525 self._send_block_delay(lines, p, b, 1)
526 return lines
529if __name__ == '__main__':
531 # PipelineSimulator([[4, 4, 4, 4], [4, 4, 4, 4], [4, 4, 4, 4 + 0.8]], 8, 0.1,
532 # [[1, 0, 0, 0], [1, 0, 0, 0], [1, 1, 0, 0]],
533 # [[1.1, 2, 2, 2], [1.1, 2, 2, 2], [1.1, 1.1, 2, 2]], method='vpp').run().show()
534 PipelineSimulator(
535 [[186.0, 171.0, 132.0, 132.0, 132.0, 132.0, 132.0, 132.0,
536 132.0, 132.0, 132.0, 132.0, 132.0, 132.0, 132.0, 133.0]], 32,
537 block_mem_act=[[1146, 908, 736, 736, 736, 736, 736, 736,
538 736, 736, 2623, 2623, 4510, 4510, 8284, 8284]],
539 block_mem_par=[[14130, 21126, 36252, 36252, 36252, 36252, 36252, 36252,
540 36252, 36252, 36252, 36252, 36252, 36252, 36252, 38297]],
541 layer_recompute=[[135.0, 171.0, 132.0, 132.0, 132.0, 132.0, 132.0, 132.0,
542 132.0, 132.0, 99.0, 99.0, 66.0, 66.0, 0, 0]],
543 less_memory=False).run().show()