Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / balancing_adapter.py: 81%
219 statements
« 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"""Utilities for adapting balancing for a given pipeline configuration"""
16from dataclasses import dataclass
17import copy
19from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger
22@dataclass
23class Pipeline:
24 """Class storing a pipeline configuration"""
26 pp: int
27 vpp: int
29 def __init__(self, pp, vpp):
30 self.pp = pp
31 self.vpp = vpp
33 def __str__(self):
34 return "PP: " + str(self.pp) + ", VPP: " + str(self.vpp)
36 def chunk_stage(self):
37 """Product of chunks and stages"""
38 return self.pp * self.vpp
41def infer_pp_and_vpp(offset):
42 """Return a pipeline configuration inferred from an offset"""
43 if is_zero_d(offset):
44 return Pipeline(1, 1)
45 if is_one_d(offset):
46 return Pipeline(len(offset), 1)
47 if is_two_d(offset):
48 return Pipeline(len(offset[0]), len(offset))
49 raise TypeError(f"Offset {offset} has a wrong type!")
52class BalancingAdapter:
53 """Adapt pipeline balancing to a given 'new' pipeline configuration"""
55 layers: int
56 prev_pip: Pipeline
58 def __init__(self, layers, offset, recompute, manual_ppb):
59 self.layers = layers
60 self.prev_offset = offset
61 self.prev_recompute = recompute
62 self.from_config = manual_ppb
63 logger.debug(
64 "init: layers = %d, from_config = %s, offset = %s, recompute = %s",
65 layers,
66 str(self.from_config),
67 str(offset),
68 str(recompute),
69 )
70 self.prev_pip = infer_pp_and_vpp(offset)
72 def treat_pp_list(self, new_pip, stages):
73 """Treat 1D offset or recompute config"""
74 current_pp_len = len(stages)
75 logger.debug(
76 "new pp (%d) = prev pp (%d)? l = %s, len(l) = %d",
77 new_pip.pp,
78 self.prev_pip.pp,
79 str(stages),
80 len(stages),
81 )
82 if new_pip.pp == current_pp_len:
83 return stages
84 new_l = []
85 if new_pip.pp > current_pp_len:
86 while new_pip.pp % len(stages) != 0:
87 stages.insert(len(stages) // 2, 0)
88 logger.debug("stages: %s", str(stages))
89 factor = new_pip.pp // current_pp_len
90 for s in stages:
91 rest = s % factor
92 for _ in range(factor):
93 if rest > 0:
94 new_l.append(s // factor + 1)
95 rest -= 1
96 else:
97 new_l.append(s // factor)
98 if new_pip.pp < current_pp_len:
99 for _ in range(current_pp_len % new_pip.pp):
100 stages.insert(len(stages) // 2, 0)
101 factor = current_pp_len // new_pip.pp
102 for i in range(new_pip.pp):
103 total_rec_layers = 0
104 for j in range(factor):
105 total_rec_layers += stages[i * factor + j]
106 new_l.append(total_rec_layers)
107 return new_l
109 def treat_vpp_list(self, new_pip, ll):
110 """Treat 2D offset or recompute config"""
111 logger.debug("treat_vpp_list start: ll = %s", str(ll))
112 prev_vpp = len(ll)
113 if prev_vpp < new_pip.vpp:
114 if prev_vpp == 1:
115 ll.append([0] * self.prev_pip.pp)
116 ll[1][-1] = ll[0][-1]
117 ll[0][-1] = 0
118 prev_vpp = 2
119 for _ in range(new_pip.vpp - prev_vpp):
120 ll.insert(prev_vpp // 2, [0] * self.prev_pip.pp)
121 logger.debug("B: ll = %s", str(ll))
123 prev_layer_per_stage = [
124 prev_vpp * (self.layers // self.prev_pip.chunk_stage())
125 + sum(stages[p] for stages in make_two_d(self.prev_offset))
126 for p in range(self.prev_pip.pp)
127 ]
128 new_layer_per_stage = [
129 new_pip.vpp * (self.layers // self.prev_pip.chunk_stage())
130 + sum(stages[p] for stages in ll)
131 for p in range(self.prev_pip.pp)
132 ]
133 logger.debug(
134 "prev_layer_per_stage = %s", str(prev_layer_per_stage)
135 )
136 logger.debug("new_layer_per_stage = %s", str(new_layer_per_stage))
137 for s in range(self.prev_pip.pp):
138 for v in range(new_pip.vpp):
139 if (
140 prev_layer_per_stage[s] - new_layer_per_stage[s] > 0
141 and ll[v][s] < 1
142 ):
143 ll[v][s] += 1
144 new_layer_per_stage[s] += 1
146 elif prev_vpp > new_pip.vpp:
147 for _ in range(prev_vpp - new_pip.vpp):
148 ll[-2] = [sum(x) for x in zip(ll[-1], ll[-2])]
149 del ll[-1]
150 logger.debug("C: ll = %s", str(ll))
151 new_vpp_list = []
152 for stages in ll:
153 new_vpp_list.append(self.treat_pp_list(new_pip, stages))
154 logger.debug("D: new_vpp_list = %s", str(new_vpp_list))
155 return new_vpp_list
157 def treat_recompute_list(self, new_pip, recompute):
158 """Treat recompute config recursively"""
159 if all(isinstance(x, int) for x in recompute):
160 if all(isinstance(x, int) for x in recompute):
161 if new_pip.vpp == 1:
162 return self.treat_pp_list(new_pip, recompute)
163 return self.treat_vpp_list(new_pip, recompute)
164 elif all(isinstance(x, int) for stages in recompute for x in stages):
165 if new_pip.vpp == 1:
166 return self.treat_vpp_list(new_pip, recompute)[0]
167 return self.treat_vpp_list(new_pip, recompute)
168 return recompute
170 def treat_recompute(self, new_pp, new_vpp):
171 """Treat recompute config for the new given Pipeline config"""
172 new_pip = Pipeline(new_pp, new_vpp)
173 if not self.from_config:
174 return self.default_recompute(
175 new_pip, copy.deepcopy(self.prev_recompute)
176 )
177 if new_pip == self.prev_pip or isinstance(self.prev_recompute, bool):
178 return copy.deepcopy(self.prev_recompute)
179 return self.treat_recompute_list(
180 new_pip, copy.deepcopy(self.prev_recompute)
181 )
183 def treat_offset(self, new_pp, new_vpp):
184 """Treat offset for the new given Pipeline config"""
185 new_pip = Pipeline(new_pp, new_vpp)
186 if not self.from_config:
187 return self.make_valid(new_pip, self.default_offset(new_pip))
188 offset = copy_offset(self.prev_offset)
189 if new_pip == self.prev_pip:
190 logger.debug("Same pipeline, no offset change")
191 return offset
192 logger.debug(
193 "change offset %s from PP = %d, VPP = %d, to PP = %d, VPP = %d",
194 str(offset),
195 self.prev_pip.pp,
196 self.prev_pip.vpp,
197 new_pip.pp,
198 new_pip.vpp,
199 )
201 if is_zero_d(offset):
202 offset = []
203 for _ in range(new_pip.vpp):
204 offset.append([])
205 for _ in range(new_pip.pp):
206 offset[-1].append(0)
207 return self.make_valid(new_pip, offset)
208 if is_one_d(offset):
209 if new_pip.vpp == 1:
210 return self.make_valid(
211 new_pip, self.treat_pp_list(new_pip, offset)
212 )
213 return self.make_valid(
214 new_pip, self.treat_vpp_list(new_pip, [offset])
215 )
216 if is_two_d(offset):
217 return self.make_valid(
218 new_pip, self.treat_vpp_list(new_pip, offset)
219 )
220 raise TypeError(f"Offset {offset} has a wrong type!")
222 def check_offset(self, new_pip, offset):
223 """Check offset validity"""
224 if is_zero_d(offset):
225 return (new_pip.pp == 1) and (new_pip.vpp == 1)
226 if is_one_d(offset):
227 return (
228 len(offset) == new_pip.pp
229 and sum(offset) == self.layers % new_pip.pp
230 )
231 if is_two_d(offset):
232 return (
233 len(offset) == new_pip.vpp
234 and all(len(stages) == new_pip.pp for stages in offset)
235 and sum(sum(stages) for stages in offset)
236 == self.layers % (new_pip.chunk_stage())
237 )
238 return False
240 def offset_checker(self, new_pp, new_vpp, offset):
241 """Log an error message if offset is invalid"""
242 new_pip = Pipeline(new_pp, new_vpp)
243 if not self.check_offset(new_pip, offset):
244 logger.error(
245 "offset %s is wrong!! pp = %d, vpp = %d, L = %d",
246 str(offset),
247 new_pip.pp,
248 new_pip.vpp,
249 self.layers,
250 )
251 return False
252 return True
254 def make_valid(self, new_pip, offset):
255 """Transform an invalid offset into a valid one"""
256 logger.debug("offset to make valid = %s", str(offset))
257 flat = make_one_d(offset)
259 delta = self.layers % (new_pip.chunk_stage()) - sum(flat)
260 logger.debug("delta = %s", str(delta))
261 if delta < 0:
262 for _ in range(-delta):
263 top = max(flat)
264 for i, _ in enumerate(flat):
265 if flat[i] == top:
266 flat[i] = flat[i] - 1
267 break
268 elif delta > 0:
269 for _ in range(delta):
270 bot = min(flat)
271 for i, _ in enumerate(flat):
272 # for i, _ in reversed(list(enumerate(flat))):
273 if flat[i] == bot:
274 flat[i] = flat[i] + 1
275 break
276 logger.debug("valid flat offset = %s", str(flat))
278 # cases where sum > (new_pip.vpp * new_pip.pp)
279 delta = sum(flat) // (new_pip.chunk_stage())
280 for i, _ in enumerate(flat):
281 flat[i] -= delta
283 logger.debug("new flat offset = %s", str(flat))
284 if new_pip.vpp == 1:
285 return flat
287 for v in range(new_pip.vpp):
288 for s in range(new_pip.pp):
289 offset[v][s] = flat[v * new_pip.pp + s]
290 return offset
292 def default_offset(self, new_pip):
293 """Construct default offset"""
294 rest = self.layers % (new_pip.chunk_stage())
295 offset = []
296 for v in range(new_pip.vpp):
297 offset.insert(0, [])
298 for s in range(new_pip.pp):
299 if rest > 0 and v > 0 and s > 0:
300 offset[0].insert(0, 1)
301 rest -= 1
302 else:
303 offset[0].insert(0, 0)
304 if offset[0][0] == 1:
305 offset[0][0] = 0
306 offset[-1][-1] = 1
307 return offset
309 def default_recompute(self, _, recompute_cfg):
310 """Construct Default recompute config when not taken from config"""
311 logger.debug(
312 "recompute config has type %s and value %s",
313 type(recompute_cfg),
314 str(recompute_cfg),
315 )
316 return True
319def copy_offset(offset):
320 """Copy an offset"""
321 if is_zero_d(offset):
322 return offset
323 return copy.deepcopy(offset)
326def is_zero_d(offset):
327 """Check if offset is an int"""
328 return isinstance(offset, int)
331def is_one_d(offset):
332 """Check if offset is an int list"""
333 return all(isinstance(x, int) for x in offset)
336def is_two_d(offset):
337 """Check if offset is an int list list"""
338 return all(isinstance(x, int) for stages in offset for x in stages)
341def make_one_d(offset):
342 """Transform offset is an int list"""
343 if is_zero_d(offset):
344 return [0]
345 if is_one_d(offset):
346 return offset
347 if is_two_d(offset):
348 return [x for stages in offset for x in stages]
349 raise TypeError(f"Offset {offset} has a wrong type!")
352def make_two_d(offset):
353 """Transform offset is an int list list"""
354 if is_zero_d(offset):
355 return [[0]]
356 if is_one_d(offset):
357 return [offset]
358 if is_two_d(offset):
359 return offset
360 raise TypeError(f"Offset {offset} has a wrong type!")