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

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"""Build pipeline scheduler chains (1F1B, VPP, VPP-less-memory).""" 

16from __future__ import annotations 

17 

18from typing import Callable, List 

19 

20from hyper_parallel.auto_parallel.sapp_ppb.simulator.sim_block import BlockSim, HeadBlockSim, MicroBlockSim 

21 

22BuilderFn = Callable[..., List[BlockSim]] 

23 

24 

25class PipelineBuilder: 

26 r"""Build pipeline scheduler""" 

27 @staticmethod 

28 def _inter_merge(a: list[MicroBlockSim], b: list[MicroBlockSim], delta: int = 0) -> list[MicroBlockSim]: 

29 r"""merge forward and backward chain for 1f1b""" 

30 res = [] 

31 if delta >= 0: 

32 res.extend(a[:delta]) 

33 a = a[delta:] 

34 else: 

35 res.extend(b[:-delta]) 

36 b = b[-delta:] 

37 stable_count = 0 

38 while a: 

39 block = a.pop(0) 

40 block.phase = 'stable' 

41 res.append(block) 

42 stable_count += 1 

43 if b: 

44 block = b.pop(0) 

45 block.phase = 'stable' 

46 res.append(block) 

47 stable_count += 1 

48 else: 

49 break 

50 if stable_count: 

51 res[-1].phase = 'cooldown' 

52 if a: 

53 res.extend(a) 

54 elif b: 

55 res.extend(b) 

56 return res 

57 

58 @staticmethod 

59 def _build_chain(line: list[MicroBlockSim], p: int) -> list[BlockSim]: 

60 r"""build pipeline chain""" 

61 # pylint: disable=E1120 

62 head = HeadBlockSim(p) 

63 left = head 

64 for item in line: 

65 left.right = item 

66 item.left = left 

67 left = item 

68 if p == 0: 

69 head.right.pre = head 

70 return line 

71 

72 @staticmethod 

73 # pylint: disable=W0613 

74 def build_1f1b(pp: int, micro_num: int, vp: int, p: int, 

75 forward_time: List[float], backward_time: List[float], 

76 block_mem: List[float], block_mem_par: List[float]) -> List[BlockSim]: 

77 """Build a 1F1B schedule chain for one pipeline rank. 

78 

79 Args: 

80 pp: Total number of pipeline stages. 

81 micro_num: Number of micro-batches. 

82 vp: Virtual-pipeline degree (unused here, kept for a common builder signature). 

83 p: Pipeline-stage index this chain belongs to. 

84 forward_time: Per-chunk forward times; only ``forward_time[0]`` is used. 

85 backward_time: Per-chunk backward times; only ``backward_time[0]`` is used. 

86 block_mem: Per-chunk activation memory; only ``block_mem[0]`` is used. 

87 block_mem_par: Per-chunk parameter memory; only ``block_mem_par[0]`` is used. 

88 

89 Returns: 

90 The ordered chain of :class:`BlockSim` nodes. 

91 """ 

92 forward_time = forward_time[0] 

93 backward_time = backward_time[0] 

94 block_mem = block_mem[0] 

95 block_mem_par = block_mem_par[0] 

96 for_line = [MicroBlockSim(p, 'f', i, 0, forward_time, mem=block_mem, mem_par=block_mem_par, phase='warmup') 

97 for i in range(micro_num)] 

98 back_line = [MicroBlockSim(p, 'b', i, 0, backward_time, mem=block_mem, mem_par=block_mem_par, phase='cooldown') 

99 for i in range(micro_num)] 

100 line = PipelineBuilder._inter_merge(for_line, back_line, pp - p - 1) 

101 return PipelineBuilder._build_chain(line, p) 

102 

103 @staticmethod 

104 def build_virtualpipeline(pp: int, micro_num: int, vp: int, p: int, 

105 forward_time: List[float], backward_time: List[float], 

106 block_mem: List[float], 

107 block_mem_par: List[float]) -> List[BlockSim]: 

108 """Build a virtual-pipeline (VPP) 1F1B chain for one pipeline rank.""" 

109 for_line = [] 

110 back_line = [] 

111 r = micro_num % pp 

112 for inter in range(micro_num // pp): 

113 for i in range(vp): 

114 bi = vp - 1 - i 

115 if inter == 0: 

116 for_line.extend([MicroBlockSim(p, 'f', m, i, forward_time[i], 

117 mem=block_mem[i], mem_par=block_mem_par[i], 

118 phase='warmup') for m in range(r)]) 

119 back_line.extend([MicroBlockSim(p, 'b', m, bi, backward_time[bi], 

120 mem=block_mem[bi], mem_par=block_mem_par[bi], 

121 phase='cooldown') for m in range(r)]) 

122 for_line.extend([MicroBlockSim(p, 'f', r + m + inter * pp, i, forward_time[i], 

123 mem=block_mem[i], mem_par=block_mem_par[i], 

124 phase='warmup') for m in range(pp)]) 

125 back_line.extend([MicroBlockSim(p, 'b', r + m + inter * pp, bi, backward_time[bi], 

126 mem=block_mem[bi], mem_par=block_mem_par[bi], 

127 phase='cooldown') for m in range(pp)]) 

128 line = PipelineBuilder._inter_merge(for_line, back_line, (vp + 1) * pp - 2 * p - 2 + r * (vp - 1)) 

129 return PipelineBuilder._build_chain(line, p) 

130 

131 @staticmethod 

132 def build_virtualpipeline2(pp: int, micro_num: int, vp: int, p: int, 

133 forward_time: List[float], backward_time: List[float], 

134 block_mem: List[float], 

135 block_mem_par: List[float]) -> List[BlockSim]: 

136 """Build a VPP 1F1B chain using the less-memory scheduler variant.""" 

137 for_line = [] 

138 back_line = [] 

139 r = micro_num % pp 

140 for inter in range(micro_num // pp): 

141 for i in range(vp): 

142 bi = vp - 1 - i 

143 if inter == 0: 

144 for_line.extend([MicroBlockSim(p, 'f', m, i, forward_time[i], 

145 mem=block_mem[i], mem_par=block_mem_par[i], 

146 phase='warmup') for m in range(r)]) 

147 back_line.extend([MicroBlockSim(p, 'b', m, bi, backward_time[bi], 

148 mem=block_mem[bi], mem_par=block_mem_par[bi], 

149 phase='cooldown') for m in range(r)]) 

150 for_line.extend([MicroBlockSim(p, 'f', r + m + inter * pp, i, forward_time[i], 

151 mem=block_mem[i], mem_par=block_mem_par[i], 

152 phase='warmup') for m in range(pp)]) 

153 back_line.extend([MicroBlockSim(p, 'b', r + m + inter * pp, bi, backward_time[bi], 

154 mem=block_mem[bi], mem_par=block_mem_par[bi], 

155 phase='cooldown') for m in range(pp)]) 

156 

157 line = PipelineBuilder._inter_merge(for_line, back_line, vp * pp - p - 1) 

158 return PipelineBuilder._build_chain(line, p) 

159 

160 @staticmethod 

161 def get_builder(method: str = '1f1b') -> BuilderFn: 

162 """Return the schedule-builder callable for a given schedule ``method``. 

163 

164 Args: 

165 method: One of ``'1f1b'``, ``'vpp'``, ``'vpp2'``. 

166 

167 Raises: 

168 ValueError: If ``method`` is not one of the supported values. 

169 """ 

170 if method == '1f1b': 

171 return PipelineBuilder.build_1f1b 

172 if method == 'vpp': 

173 return PipelineBuilder.build_virtualpipeline 

174 if method == 'vpp2': 

175 return PipelineBuilder.build_virtualpipeline2 

176 raise ValueError(f"`method` only support ['1f1b', 'vpp', 'vpp2'], but got {method}")