Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / recompute / knapsack.py: 99%
139 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
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"""Choosing each layer's option under a memory budget.
17Every layer runs one option of its kind's front. At a stage that keeps k
18micro-batches of a layer in flight, an option keeps k times its memory per
19micro-batch plus its memory held once, less what those charge beyond the
20layer at k (:meth:`~.front.LayerOption.memory`), and costs its forward and
21backward time. Choosing one option per layer so that a stage's layers are as
22fast as possible within the memory left to them is a multiple-choice
23knapsack. :func:`choose` solves it exactly over memory counted in buckets,
24each option's memory rounded up to a whole bucket, so a choice never keeps
25more than the budget.
27Without pipeline parallelism the model is one stage, and :func:`choose` is
28the whole answer. With it, :func:`pp_lite` solves every stage of a split and
29times the pipeline by its slowest stage: a cheap score to rank candidates by
30before the pipeline balancer places their layers unevenly.
31"""
32import math
33from collections import Counter
34from dataclasses import dataclass
35from typing import Callable, Hashable, List, Mapping, Optional, Sequence, Tuple
37import numpy as np
39from hyper_parallel.auto_parallel.sapp_nd.recompute.front import LayerOption
41MEGABYTE = 2 ** 20
44@dataclass(frozen=True)
45class Layers:
46 """Layers of one kind that a stage keeps as many micro-batches of in flight.
48 Attributes:
49 kind: The key of the kind's front.
50 count: How many layers.
51 in_flight: How many micro-batches of each the stage keeps at once.
52 """
54 kind: Hashable
55 count: int
56 in_flight: int = 1
59@dataclass(frozen=True)
60class Assignment:
61 """How many of a group's layers run one option."""
63 layers: Layers
64 option: LayerOption
65 count: int
68@dataclass(frozen=True)
69class StageChoice:
70 """The options a stage's layers run, and what they cost together.
72 Attributes:
73 assignments: Per group of layers, how many run each option, the
74 slower options first.
75 memory: The bytes the layers keep, each option's memory counted
76 exactly rather than in buckets.
77 time: The layers' forward and backward time.
78 """
80 assignments: Tuple[Assignment, ...]
81 memory: float
82 time: float
85@dataclass(frozen=True)
86class Stage:
87 """A pipeline stage: its layers, the memory left to them, and its time outside them.
89 Attributes:
90 groups: The stage's layers.
91 budget: The bytes the layers may keep: the device's memory less
92 everything on the stage that the choice does not change.
93 fixed_time: The stage's time outside its layers, the embedding's or
94 the output layer's.
95 """
97 groups: Tuple[Layers, ...]
98 budget: float
99 fixed_time: float = 0.0
102@dataclass(frozen=True)
103class PipelineChoice:
104 """The choice of every stage, and each stage's time."""
106 stages: Tuple[StageChoice, ...]
107 times: Tuple[float, ...]
109 @property
110 def slowest(self) -> float:
111 """The slowest stage's time, which the pipeline runs at."""
112 return max(self.times)
115def _kept(option: LayerOption, in_flight: int) -> float:
116 """The bytes a layer running *option* keeps with *in_flight* micro-batches in flight."""
117 return option.memory(in_flight)
120def _time(option: LayerOption) -> float:
121 """A layer's forward and backward time under *option*."""
122 return option.forward_time + option.backward_time
125def _useful(options: Sequence[LayerOption], in_flight: int) -> List[LayerOption]:
126 """The options no other beats on memory and time at *in_flight* micro-batches in flight."""
127 kept, fastest = [], math.inf
128 for option in sorted(options, key=lambda option: (_kept(option, in_flight), _time(option))):
129 if _time(option) < fastest:
130 kept.append(option)
131 fastest = _time(option)
132 return kept
135def choose(
136 groups: Sequence[Layers],
137 fronts: Mapping[Hashable, Sequence[LayerOption]],
138 budget: float,
139 bucket: float = MEGABYTE,
140) -> Optional[StageChoice]:
141 """The fastest options for a stage's layers that keep at most *budget* bytes.
143 Of the fastest choices, the one that keeps the least memory.
145 Args:
146 groups: The stage's layers, by kind and micro-batches in flight.
147 fronts: Each kind's options.
148 budget: The bytes the layers may keep.
149 bucket: The memory granularity, in bytes. Each option's memory is
150 rounded up to a whole bucket, so a coarser bucket solves faster
151 and may give up a little of the budget.
153 Returns:
154 The choice, or ``None`` when not even the lightest options fit.
155 """
156 if budget < 0:
157 return None
158 useful = [_useful(fronts[group.kind], group.in_flight) for group in groups]
159 if sum(group.count * _kept(options[-1], group.in_flight) for group, options in zip(groups, useful)) <= budget:
160 # Every layer's fastest option fits: nothing to choose.
161 return _stage_choice(
162 [Assignment(group, options[-1], group.count) for group, options in zip(groups, useful) if group.count]
163 )
164 size = int(budget // bucket)
165 # fastest[w]: the least time of the layers seen so far keeping exactly w buckets.
166 fastest = np.full(size + 1, np.inf)
167 fastest[0] = 0.0
168 steps = []
169 for position, (group, options) in enumerate(zip(groups, useful)):
170 weights = [math.ceil(_kept(option, group.in_flight) / bucket) for option in options]
171 for _ in range(group.count):
172 after = np.full(size + 1, np.inf)
173 picked = np.full(size + 1, -1, dtype=np.int16)
174 for index, (option, weight) in enumerate(zip(options, weights)):
175 if weight > size:
176 continue
177 candidate = fastest[:size + 1 - weight] + _time(option)
178 better = candidate < after[weight:]
179 after[weight:][better] = candidate[better]
180 picked[weight:][better] = index
181 fastest = after
182 steps.append((position, weights, picked))
183 if not np.isfinite(fastest).any():
184 return None
185 used = int(np.argmin(fastest))
186 counts = Counter()
187 for position, weights, picked in reversed(steps):
188 index = int(picked[used])
189 counts[position, index] += 1
190 used -= weights[index]
191 assignments = []
192 for position, (group, options) in enumerate(zip(groups, useful)):
193 chosen = [(options[index], count) for (found, index), count in counts.items() if found == position]
194 for option, count in sorted(chosen, key=lambda pair: -_time(pair[0])):
195 assignments.append(Assignment(group, option, count))
196 return _stage_choice(assignments)
199def least_times(
200 groups: Sequence[Layers],
201 fronts: Mapping[Hashable, Sequence[LayerOption]],
202 budget: float,
203 bucket: float = MEGABYTE,
204) -> Callable[[float], float]:
205 """The least time of a stage's layers within every budget up to *budget*, as :func:`choose` finds it.
207 One pass over the layers, as :func:`choose` makes for one budget, answers
208 all of them.
210 Args:
211 groups: The stage's layers, by kind and micro-batches in flight.
212 fronts: Each kind's options.
213 budget: The largest budget asked, in bytes.
214 bucket: The memory granularity, as for :func:`choose`.
216 Returns:
217 For a budget in bytes, the time of the choice :func:`choose` makes
218 within it; ``inf`` where none fits.
219 """
220 useful = [_useful(fronts[group.kind], group.in_flight) for group in groups]
221 heaviest = sum(group.count * _kept(options[-1], group.in_flight) for group, options in zip(groups, useful))
222 fastest_time = sum(group.count * _time(options[-1]) for group, options in zip(groups, useful))
223 # Past the fastest options' memory every budget has them.
224 size = int(max(0.0, min(budget, heaviest)) // bucket)
225 fastest = np.full(size + 1, np.inf)
226 fastest[0] = 0.0
227 for group, options in zip(groups, useful):
228 weights = [math.ceil(_kept(option, group.in_flight) / bucket) for option in options]
229 for _ in range(group.count):
230 after = np.full(size + 1, np.inf)
231 for option, weight in zip(options, weights):
232 if weight <= size:
233 np.minimum(after[weight:], fastest[:size + 1 - weight] + _time(option), out=after[weight:])
234 fastest = after
235 within = np.minimum.accumulate(fastest)
237 def _within(room: float) -> float:
238 """The least time within *room* bytes."""
239 if room < 0:
240 return math.inf
241 if heaviest <= room:
242 return fastest_time
243 return float(within[min(int(room // bucket), size)])
245 return _within
248def suffix_times(
249 layers: Sequence[Layers],
250 fronts: Mapping[Hashable, Sequence[LayerOption]],
251 budgets: Sequence[float],
252 bucket: float = MEGABYTE,
253) -> List[float]:
254 """For every start, the least time of the layers from it on that keep at most its budget.
256 One pass over the layers from the last to the first, as :func:`choose`
257 solves one stage, gives every suffix's fastest choice at once.
259 Args:
260 layers: The layers in order, one :class:`Layers` of count 1 each.
261 fronts: Each kind's options.
262 budgets: One budget in bytes per start, ``len(layers) + 1`` of them;
263 the last is the empty suffix's.
264 bucket: The memory granularity, as for :func:`choose`.
266 Returns:
267 ``len(layers) + 1`` times, ``inf`` where a suffix cannot fit its
268 budget.
269 """
270 if len(budgets) != len(layers) + 1:
271 raise ValueError(f"{len(budgets)} budgets for {len(layers)} layers; expected one more")
272 size = int(max(0.0, *budgets) // bucket)
273 fastest = np.full(size + 1, np.inf)
274 fastest[0] = 0.0
276 def _least(start: int) -> float:
277 """The least time of the suffix now in *fastest* within its budget."""
278 if budgets[start] < 0:
279 return math.inf
280 return float(np.min(fastest[:int(budgets[start] // bucket) + 1]))
282 times = [math.inf] * (len(layers) + 1)
283 times[-1] = _least(len(layers))
284 for start in range(len(layers) - 1, -1, -1):
285 layer = layers[start]
286 after = np.full(size + 1, np.inf)
287 for option in _useful(fronts[layer.kind], layer.in_flight):
288 weight = math.ceil(_kept(option, layer.in_flight) / bucket)
289 if weight > size:
290 continue
291 np.minimum(after[weight:], fastest[:size + 1 - weight] + _time(option), out=after[weight:])
292 fastest = after
293 times[start] = _least(start)
294 return times
297def _stage_choice(assignments: Sequence[Assignment]) -> StageChoice:
298 """A stage's assignments, with the memory they keep and the time they take."""
299 return StageChoice(
300 assignments=tuple(assignments),
301 memory=sum(item.count * _kept(item.option, item.layers.in_flight) for item in assignments),
302 time=sum(item.count * _time(item.option) for item in assignments),
303 )
306def pp_lite(
307 stages: Sequence[Stage], fronts: Mapping[Hashable, Sequence[LayerOption]], bucket: float = MEGABYTE
308) -> Optional[PipelineChoice]:
309 """The fastest choice for every stage, each within its own budget.
311 Stages share no memory, so choosing each stage's fastest options also
312 makes the slowest stage as fast as it can be.
314 Args:
315 stages: The pipeline's stages.
316 fronts: Each kind's options.
317 bucket: The memory granularity, as for :func:`choose`.
319 Returns:
320 The choice, or ``None`` when some stage cannot fit.
321 """
322 choices = []
323 for stage in stages:
324 choice = choose(stage.groups, fronts, stage.budget, bucket)
325 if choice is None:
326 return None
327 choices.append(choice)
328 return PipelineChoice(
329 stages=tuple(choices),
330 times=tuple(choice.time + stage.fixed_time for choice, stage in zip(choices, stages)),
331 )