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