Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / balancing_adapter.py: 82%
233 statements
« 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"""Utilities for adapting balancing for a given pipeline configuration"""
16from dataclasses import dataclass
17import copy
18from typing import List, Union
20from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger
23@dataclass
24class Pipeline:
25 """Class storing a pipeline configuration"""
27 pp: int
28 vpp: int
30 def __init__(self, pp, vpp):
31 self.pp = pp
32 self.vpp = vpp
34 def __str__(self):
35 return "PP: " + str(self.pp) + ", VPP: " + str(self.vpp)
37 def chunk_stage(self):
38 """Product of chunks and stages"""
39 return self.pp * self.vpp
42def front_loaded_offset(layers: int, pp: int, vpp: int) -> Union[List[int], List[List[int]]]:
43 """An offset placing every one of *layers* on the first pipeline stage.
45 A vision tower runs before the language model and on the stage that
46 holds the model's first layers, so its layers stay there whatever
47 pipeline depth is priced. The shape is the one
48 ``is_consistent_pp_config`` wants: a list per stage, or one per chunk
49 and stage under interleaving, of each stage's layers beyond an even
50 share.
51 """
52 per_stage = max(0, layers // max(1, pp) // max(1, vpp))
53 head = layers - per_stage
54 if vpp > 1:
55 chunks = [[-per_stage] * pp for _ in range(vpp)]
56 chunks[0][0] = head
57 return chunks
58 stages = [-per_stage] * pp
59 stages[0] = head
60 return stages
63def infer_pp_and_vpp(offset):
64 """Return a pipeline configuration inferred from an offset"""
65 if is_zero_d(offset):
66 return Pipeline(1, 1)
67 if is_one_d(offset):
68 return Pipeline(len(offset), 1)
69 if is_two_d(offset):
70 return Pipeline(len(offset[0]), len(offset))
71 raise TypeError(f"Offset {offset} has a wrong type!")
74class BalancingAdapter:
75 """Adapt pipeline balancing to a given 'new' pipeline configuration"""
77 layers: int
78 prev_pip: Pipeline
80 def __init__(self, layers, offset, recompute, manual_ppb):
81 self.layers = layers
82 self.prev_offset = offset
83 self.prev_recompute = recompute
84 self.from_config = manual_ppb
85 # The full recompute a recompute dimension states for every
86 # candidate, while the search prices one of its modes; None leaves
87 # each candidate the recompute derived below.
88 self.stated_recompute = None
89 logger.debug(
90 "init: layers = %d, from_config = %s, offset = %s, recompute = %s",
91 layers,
92 str(self.from_config),
93 str(offset),
94 str(recompute),
95 )
96 self.prev_pip = infer_pp_and_vpp(offset)
98 def treat_pp_list(self, new_pip, stages):
99 """Treat 1D offset or recompute config"""
100 current_pp_len = len(stages)
101 logger.debug(
102 "new pp (%d) = prev pp (%d)? l = %s, len(l) = %d",
103 new_pip.pp,
104 self.prev_pip.pp,
105 str(stages),
106 len(stages),
107 )
108 if new_pip.pp == current_pp_len:
109 return stages
110 new_l = []
111 if new_pip.pp > current_pp_len:
112 while new_pip.pp % len(stages) != 0:
113 stages.insert(len(stages) // 2, 0)
114 logger.debug("stages: %s", str(stages))
115 factor = new_pip.pp // current_pp_len
116 for s in stages:
117 rest = s % factor
118 for _ in range(factor):
119 if rest > 0:
120 new_l.append(s // factor + 1)
121 rest -= 1
122 else:
123 new_l.append(s // factor)
124 if new_pip.pp < current_pp_len:
125 for _ in range(current_pp_len % new_pip.pp):
126 stages.insert(len(stages) // 2, 0)
127 factor = current_pp_len // new_pip.pp
128 for i in range(new_pip.pp):
129 total_rec_layers = 0
130 for j in range(factor):
131 total_rec_layers += stages[i * factor + j]
132 new_l.append(total_rec_layers)
133 return new_l
135 def treat_vpp_list(self, new_pip, ll):
136 """Treat 2D offset or recompute config"""
137 logger.debug("treat_vpp_list start: ll = %s", str(ll))
138 prev_vpp = len(ll)
139 if prev_vpp < new_pip.vpp:
140 if prev_vpp == 1:
141 ll.append([0] * self.prev_pip.pp)
142 ll[1][-1] = ll[0][-1]
143 ll[0][-1] = 0
144 prev_vpp = 2
145 for _ in range(new_pip.vpp - prev_vpp):
146 ll.insert(prev_vpp // 2, [0] * self.prev_pip.pp)
147 logger.debug("B: ll = %s", str(ll))
149 prev_layer_per_stage = [
150 prev_vpp * (self.layers // self.prev_pip.chunk_stage())
151 + sum(stages[p] for stages in make_two_d(self.prev_offset))
152 for p in range(self.prev_pip.pp)
153 ]
154 new_layer_per_stage = [
155 new_pip.vpp * (self.layers // self.prev_pip.chunk_stage())
156 + sum(stages[p] for stages in ll)
157 for p in range(self.prev_pip.pp)
158 ]
159 logger.debug(
160 "prev_layer_per_stage = %s", str(prev_layer_per_stage)
161 )
162 logger.debug("new_layer_per_stage = %s", str(new_layer_per_stage))
163 for s in range(self.prev_pip.pp):
164 for v in range(new_pip.vpp):
165 if (
166 prev_layer_per_stage[s] - new_layer_per_stage[s] > 0
167 and ll[v][s] < 1
168 ):
169 ll[v][s] += 1
170 new_layer_per_stage[s] += 1
172 elif prev_vpp > new_pip.vpp:
173 for _ in range(prev_vpp - new_pip.vpp):
174 ll[-2] = [sum(x) for x in zip(ll[-1], ll[-2])]
175 del ll[-1]
176 logger.debug("C: ll = %s", str(ll))
177 new_vpp_list = []
178 for stages in ll:
179 new_vpp_list.append(self.treat_pp_list(new_pip, stages))
180 logger.debug("D: new_vpp_list = %s", str(new_vpp_list))
181 return new_vpp_list
183 def treat_recompute_list(self, new_pip, recompute):
184 """Treat recompute config recursively"""
185 if all(isinstance(x, int) for x in recompute):
186 if all(isinstance(x, int) for x in recompute):
187 if new_pip.vpp == 1:
188 return self.treat_pp_list(new_pip, recompute)
189 return self.treat_vpp_list(new_pip, recompute)
190 elif all(isinstance(x, int) for stages in recompute for x in stages):
191 if new_pip.vpp == 1:
192 return self.treat_vpp_list(new_pip, recompute)[0]
193 return self.treat_vpp_list(new_pip, recompute)
194 return recompute
196 def treat_recompute(self, new_pp, new_vpp):
197 """Treat recompute config for the new given Pipeline config"""
198 new_pip = Pipeline(new_pp, new_vpp)
199 if self.stated_recompute is not None:
200 return self.stated_recompute
201 if not self.from_config:
202 return self.default_recompute(
203 new_pip, copy.deepcopy(self.prev_recompute)
204 )
205 if new_pip == self.prev_pip or isinstance(self.prev_recompute, bool):
206 return copy.deepcopy(self.prev_recompute)
207 return self.treat_recompute_list(
208 new_pip, copy.deepcopy(self.prev_recompute)
209 )
211 def treat_offset(self, new_pp, new_vpp):
212 """Treat offset for the new given Pipeline config"""
213 new_pip = Pipeline(new_pp, new_vpp)
214 if not self.from_config:
215 return self.make_valid(new_pip, self.default_offset(new_pip))
216 offset = copy_offset(self.prev_offset)
217 if new_pip == self.prev_pip:
218 logger.debug("Same pipeline, no offset change")
219 return offset
220 logger.debug(
221 "change offset %s from PP = %d, VPP = %d, to PP = %d, VPP = %d",
222 str(offset),
223 self.prev_pip.pp,
224 self.prev_pip.vpp,
225 new_pip.pp,
226 new_pip.vpp,
227 )
229 if is_zero_d(offset):
230 offset = []
231 for _ in range(new_pip.vpp):
232 offset.append([])
233 for _ in range(new_pip.pp):
234 offset[-1].append(0)
235 return self.make_valid(new_pip, offset)
236 if is_one_d(offset):
237 if new_pip.vpp == 1:
238 return self.make_valid(
239 new_pip, self.treat_pp_list(new_pip, offset)
240 )
241 return self.make_valid(
242 new_pip, self.treat_vpp_list(new_pip, [offset])
243 )
244 if is_two_d(offset):
245 return self.make_valid(
246 new_pip, self.treat_vpp_list(new_pip, offset)
247 )
248 raise TypeError(f"Offset {offset} has a wrong type!")
250 def check_offset(self, new_pip, offset):
251 """Check offset validity"""
252 if is_zero_d(offset):
253 return (new_pip.pp == 1) and (new_pip.vpp == 1)
254 if is_one_d(offset):
255 return (
256 len(offset) == new_pip.pp
257 and sum(offset) == self.layers % new_pip.pp
258 )
259 if is_two_d(offset):
260 return (
261 len(offset) == new_pip.vpp
262 and all(len(stages) == new_pip.pp for stages in offset)
263 and sum(sum(stages) for stages in offset)
264 == self.layers % (new_pip.chunk_stage())
265 )
266 return False
268 def offset_checker(self, new_pp, new_vpp, offset):
269 """Log an error message if offset is invalid"""
270 new_pip = Pipeline(new_pp, new_vpp)
271 if not self.check_offset(new_pip, offset):
272 logger.error(
273 "offset %s is wrong!! pp = %d, vpp = %d, L = %d",
274 str(offset),
275 new_pip.pp,
276 new_pip.vpp,
277 self.layers,
278 )
279 return False
280 return True
282 def make_valid(self, new_pip, offset):
283 """Transform an invalid offset into a valid one"""
284 logger.debug("offset to make valid = %s", str(offset))
285 flat = make_one_d(offset)
287 delta = self.layers % (new_pip.chunk_stage()) - sum(flat)
288 logger.debug("delta = %s", str(delta))
289 if delta < 0:
290 for _ in range(-delta):
291 top = max(flat)
292 for i, _ in enumerate(flat):
293 if flat[i] == top:
294 flat[i] = flat[i] - 1
295 break
296 elif delta > 0:
297 for _ in range(delta):
298 bot = min(flat)
299 for i, _ in enumerate(flat):
300 # for i, _ in reversed(list(enumerate(flat))):
301 if flat[i] == bot:
302 flat[i] = flat[i] + 1
303 break
304 logger.debug("valid flat offset = %s", str(flat))
306 # cases where sum > (new_pip.vpp * new_pip.pp)
307 delta = sum(flat) // (new_pip.chunk_stage())
308 for i, _ in enumerate(flat):
309 flat[i] -= delta
311 logger.debug("new flat offset = %s", str(flat))
312 if new_pip.vpp == 1:
313 return flat
315 for v in range(new_pip.vpp):
316 for s in range(new_pip.pp):
317 offset[v][s] = flat[v * new_pip.pp + s]
318 return offset
320 def default_offset(self, new_pip):
321 """Construct default offset"""
322 rest = self.layers % (new_pip.chunk_stage())
323 offset = []
324 for v in range(new_pip.vpp):
325 offset.insert(0, [])
326 for s in range(new_pip.pp):
327 if rest > 0 and v > 0 and s > 0:
328 offset[0].insert(0, 1)
329 rest -= 1
330 else:
331 offset[0].insert(0, 0)
332 if offset[0][0] == 1:
333 offset[0][0] = 0
334 offset[-1][-1] = 1
335 return offset
337 def default_recompute(self, _, recompute_cfg):
338 """Construct Default recompute config when not taken from config"""
339 logger.debug(
340 "recompute config has type %s and value %s",
341 type(recompute_cfg),
342 str(recompute_cfg),
343 )
344 return True
347def copy_offset(offset):
348 """Copy an offset"""
349 if is_zero_d(offset):
350 return offset
351 return copy.deepcopy(offset)
354def is_zero_d(offset):
355 """Check if offset is an int"""
356 return isinstance(offset, int)
359def is_one_d(offset):
360 """Check if offset is an int list"""
361 return all(isinstance(x, int) for x in offset)
364def is_two_d(offset):
365 """Check if offset is an int list list"""
366 return all(isinstance(x, int) for stages in offset for x in stages)
369def make_one_d(offset):
370 """Transform offset is an int list"""
371 if is_zero_d(offset):
372 return [0]
373 if is_one_d(offset):
374 return offset
375 if is_two_d(offset):
376 return [x for stages in offset for x in stages]
377 raise TypeError(f"Offset {offset} has a wrong type!")
380def make_two_d(offset):
381 """Transform offset is an int list list"""
382 if is_zero_d(offset):
383 return [[0]]
384 if is_one_d(offset):
385 return [offset]
386 if is_two_d(offset):
387 return offset
388 raise TypeError(f"Offset {offset} has a wrong type!")