« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025 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"""generate pipeline partitions"""
16from hyper_parallel.auto_parallel.sapp_nd.nd.common._cost_model_variables import _CostModVar
17from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
20class PartitionGenerator(_CostModVar):
21 """partition generator class"""
23 def combine_partition_multimodal(self, mm_stages: dict):
24 """Call after generate_partitions
25 Assuming same PP and VPP degree for each mod
26 Assuming offsets are well set and compatible
27 (otherwise need to preprocess offsets for module compatibility)
28 """
29 combined = []
30 for idx in range(self.p):
31 combined_stages = []
32 for v_idx in range(self.vp):
33 combined_chunks = []
34 for m in self.mm_order:
35 combined_chunks += mm_stages[m][idx][v_idx]
36 combined_stages += [combined_chunks]
37 combined += [combined_stages]
38 return combined
40 def generate_partitions_vpp(self):
41 """infer layers arrangement"""
42 if not self.multimodal:
43 return self.generate_partitions_vpp_unimodal()
44 lpartitions = {}
45 for k, v in self.mm_ccfgs.items():
46 lpartitions[k] = v.generate_partitions_vpp_unimodal()
47 return lpartitions
49 def __process_select_recompute(self, *args):
50 """Build layer types for selective recompute partitions."""
51 args = dict(
52 zip(
53 [
54 "s_idx",
55 "v_idx",
56 "num_assigned",
57 "num_layer_per_stage",
58 ],
59 args,
60 )
61 )
62 chunk = []
63 sel_rec = self.sel_rec
64 if isinstance(sel_rec, list):
65 if isinstance(sel_rec[0], str):
66 sel_rec = [
67 self.n_lay // self.p
68 for _ in range(self.p)
69 ] # Operator names list treated as True
70 if self.vp > 1:
71 if isinstance(sel_rec[0], int):
72 # Preprocess input: Even recomputed layers
73 # distribution throughout chunks
74 sel_rec = [
75 [
76 (
77 s // self.vp
78 + (1 if v_idx < s % self.vp else 0)
79 )
80 for s in sel_rec
81 ]
82 for v_idx in range(self.vp)
83 ]
84 if isinstance(sel_rec[0][0], str):
85 sel_rec = [
86 [
87 args["num_layer_per_stage"]
88 for _ in range(self.p)
89 ]
90 for _ in range(self.vp)
91 ]
92 # Operator names list treated
93 # as True (chunk level)
94 chunk = [
95 (
96 LayerType.SEL_REC_LAYER
97 if j < sel_rec[args["v_idx"]][args["s_idx"]]
98 else LayerType.NOT_REC_LAYER
99 )
100 for j in range(args["num_assigned"])
101 ] # Assuming list of list
102 else:
103 chunk = [
104 (
105 LayerType.SEL_REC_LAYER
106 if j < sel_rec[args["s_idx"]]
107 else LayerType.NOT_REC_LAYER
108 )
109 for j in range(args["num_assigned"])
110 ]
111 else:
112 chunk = [
113 LayerType.SEL_REC_LAYER for _ in range(args["num_assigned"])
114 ]
115 return chunk
117 def __process_full_recompute(self, *args):
118 """Build layer types for full recompute partitions."""
119 args = dict(
120 zip(["s_idx", "v_idx", "num_assigned"], args)
121 )
122 chunk = []
123 full_rec = self.full_rec
124 if isinstance(full_rec, list):
125 if self.vp > 1:
126 if isinstance(full_rec[0], int):
127 # Preprocess input: Even recomputed layers
128 # distribution throughout chunks
129 full_rec = [
130 [
131 (
132 f // self.vp
133 + (1 if v_idx < f % self.vp else 0)
134 )
135 for f in full_rec
136 ]
137 for v_idx in range(self.vp)
138 ]
139 chunk = [
140 (
141 LayerType.FULL_REC_LAYER
142 if j < full_rec[args["v_idx"]][args["s_idx"]]
143 else LayerType.NOT_REC_LAYER
144 )
145 for j in range(args["num_assigned"])
146 ] # Assuming list of list
147 else:
148 chunk = [
149 (
150 LayerType.FULL_REC_LAYER
151 if j < full_rec[args["s_idx"]]
152 else LayerType.NOT_REC_LAYER
153 )
154 for j in range(args["num_assigned"])
155 ]
156 else:
157 chunk = [
158 LayerType.FULL_REC_LAYER for _ in range(args["num_assigned"])
159 ]
160 return chunk
162 def __process_offset(self, *args):
163 """Compute assigned layer count after applying PP offsets."""
164 args = dict(
165 zip(
166 [
167 "s_idx",
168 "v_idx",
169 "lay",
170 "num_layer_per_stage",
171 ],
172 args,
173 )
174 )
175 num_assigned = 0
176 # try:
177 if isinstance(self.offset, list):
178 if isinstance(self.offset[0], list): # VPP
179 num_assigned = min(
180 args["lay"],
181 args["num_layer_per_stage"]
182 + self.offset[args["v_idx"]][args["s_idx"]],
183 ) # Assuming list of list
184 else:
185 # add extra to last chunk
186 if args["v_idx"] < self.vp - 1:
187 num_assigned = min(
188 args["lay"],
189 args["num_layer_per_stage"]
190 + self.offset[args["s_idx"]] // self.vp,
191 )
192 else:
193 num_assigned = min(
194 args["lay"],
195 args["num_layer_per_stage"]
196 + self.offset[args["s_idx"]] // self.vp
197 + (self.offset[args["s_idx"]] % self.vp),
198 )
199 # print(args["num_layer_per_stage"], args["lay"])
200 # print(f"num_assigned for stage {args['s_idx']} chunk {args['v_idx']}: {num_assigned}")
201 return num_assigned
203 def generate_partitions_vpp_unimodal(self):
204 """infer layers arrangement for a module"""
205 lay = self.n_lay
206 if self.is_mtp_in_offset:
207 lay += self.n_mtp
208 if self.emb_out_in_offset:
209 lay += 2
210 num_layer_per_stage = lay // self.p // self.vp
211 # print("lay",lay,"num_layer_per_stage",num_layer_per_stage, self.emb_out_in_offset)
212 num_layer_per_stage = max(
213 int(self.config_format == "json"), num_layer_per_stage
214 )
215 partitions = []
216 for idx in range(self.p):
217 partitions.append([])
218 for v_idx in range(self.vp):
219 # Process offset
220 if self.offset:
221 num_assigned = self.__process_offset(
222 idx, v_idx, lay, num_layer_per_stage
223 )
224 else:
225 num_assigned = min(lay, num_layer_per_stage)
227 # Process full recompute
228 if self.full_rec:
229 partitions[-1] += [
230 self.__process_full_recompute(
231 idx, v_idx, num_assigned
232 )
233 ]
234 # Process select recompute
235 elif self.sel_rec:
236 partitions[-1] += [
237 self.__process_select_recompute(
238 idx,
239 v_idx,
240 num_assigned,
241 num_layer_per_stage,
242 )
243 ]
244 else:
245 partitions[-1] += [
246 [LayerType.NOT_REC_LAYER for _ in range(num_assigned)]
247 ]
248 lay -= num_assigned
250 partitions = self.insert_emb_out_partitions(partitions)
251 return partitions
253 def insert_emb_out_partitions(self, partitions):
254 """end of generation"""
255 sched = self.pp_sched
256 if sched == "zero_bubble_v":
257 return self.insert_emb_out_partitions_zbv(partitions)
258 return self.insert_emb_out_partitions_1f1b(partitions)
260 def insert_emb_out_partitions_zbv(self, partitions):
261 """Emb and Out"""
262 first, last = self.first_and_last_non_empty_stage(partitions)
263 if not self.is_mtp_in_offset:
264 last_lay = partitions[last[0]][last[1]][-1]
265 partitions[last[0]][last[1]] += [last_lay] * self.n_mtp
266 if not self.emb_out_in_offset:
267 partitions[first[0]][first[1]].insert(0, LayerType.EMBEDDING_LAYER)
268 partitions[last[0]][last[1]].append(LayerType.OUTPUT_LAYER)
269 else:
270 partitions[first[0]][first[1]][0] = LayerType.EMBEDDING_LAYER
271 partitions[last[0]][last[1]][-1] = LayerType.OUTPUT_LAYER
272 return partitions
274 def insert_emb_out_partitions_1f1b(self, partitions):
275 """Insert embedding, output, and MTP layers for 1F1B schedules."""
276 first, last = self.first_and_last_non_empty_stage(partitions)
277 if not self.is_mtp_in_offset:
278 last_lay = partitions[last[0]][last[1]][-1]
279 partitions[last[0]][last[1]] += [last_lay] * self.n_mtp
280 if not self.emb_out_in_offset:
281 partitions[first[0]][first[1]].insert(0, LayerType.EMBEDDING_LAYER)
282 partitions[last[0]][last[1]].append(LayerType.OUTPUT_LAYER)
283 else:
284 partitions[first[0]][first[1]][0] = LayerType.EMBEDDING_LAYER
285 partitions[last[0]][last[1]][-1] = LayerType.OUTPUT_LAYER
286 return partitions
288 def first_and_last_non_empty_stage(self, partitions):
289 """Return the first and last non-empty stage/chunk indexes."""
290 # Useful for submodule partitioning in multimodals
291 first, last = None, None
292 for chunk_id in range(self.vp):
293 for stage_id in range(self.p):
294 if partitions[stage_id][chunk_id]:
295 last = (stage_id, chunk_id)
296 if not first:
297 first = (stage_id, chunk_id)
298 elif first and last:
299 return first, last
300 if not first:
301 first = (0, 0)
302 if not last:
303 last = (self.p - 1, self.vp - 1)
304 return first, last
306 # def insert_emb_out_partitions_1f1b(self, partitions):
307 # """Emb and Out"""
308 # # Insert emb in the first non empty chunk
309 # put_emb = None
310 # for chunk_id in range(self.vp):
311 # for stage_id in range(self.p):
312 # if partitions[stage_id][chunk_id]:
313 # partitions[stage_id][chunk_id].insert(
314 # 0, LayerType.EMBEDDING_LAYER
315 # )
316 # put_emb = (stage_id, chunk_id)
317 # break
318 # if put_emb:
319 # break
320 # # Insert out in the last non empty chunk
321 # if self.n_lay > 1:
322 # last_layer = None
323 # for chunk_id in range(self.vp - 1, -1, -1):
324 # for stage_id in range(self.p - 1, -1, -1):
325 # if partitions[stage_id][chunk_id]:
326 # last_layer = (stage_id, chunk_id)
327 # break
328 # if last_layer:
329 # break
330 # if not last_layer:
331 # last_layer = (self.p - 1, self.vp - 1)
332 # if not self.is_mtp_in_offset:
333 # # duplicate last_layer for mtp
334 # partitions[self.p - 1][self.vp - 1] += [
335 # partitions[last_layer[0]][last_layer[1]][-1]
336 # ] * self.n_mtp
337 # last_layer = (self.p - 1, self.vp - 1)
338 # partitions[last_layer[0]][last_layer[1]].append(
339 # LayerType.OUTPUT_LAYER
340 # )
341 # return partitions
342 # if not self.is_mtp_in_offset:
343 # partitions[put_emb[0]][put_emb[1]] += [
344 # partitions[put_emb[0]][put_emb[1]][-1]
345 # ] * self.n_mtp
346 # partitions[put_emb[0]][put_emb[1]].append(LayerType.OUTPUT_LAYER)
347 # return partitions