Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / recompute / front.py: 100%
82 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"""The recompute front of a layer kind: the options no other option beats.
17An option runs a layer one way: plain, recomputing some of the ops the
18recompute switches name, or fully recomputed. Of the 128 settings of the
19seven switches, 256 on an MLA layer, which may recompute its up-projections,
20and full recompute, the front keeps those that no other option beats
21on memory per micro-batch, memory held once, memory at each count of
22micro-batches in flight a stage keeps, the working set of its backward, and
23backward time together. The plain layer and full recompute are always on it.
25:func:`layer_fronts` is the interim form of the search's entry point for
26layer options (shared decision S3): it measures each layer kind on the memory
27backbone's config at the evaluator's current strategy, and phase 5 of the
28search re-implements it on the model IR.
29"""
30import itertools
31from dataclasses import dataclass
32from typing import Any, Dict, FrozenSet, Iterator, List, Mapping, Optional, Sequence, Tuple
34from hyper_parallel.auto_parallel._op_profiles import LayerKind
35from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.estimate_v2 import EvaluatorV2
36from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.estimate import LayerTimes
37from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.utils_classes import CustomConfig
38from hyper_parallel.auto_parallel.sapp_nd.recompute.profile import SWITCHES, Cost, SwitchProfile
41@dataclass(frozen=True)
42class LayerOption:
43 """One way of running a layer of a kind, and what it costs.
45 Attributes:
46 recompute: The switches set to recompute their op; ``None`` for full
47 recompute.
48 memory_per_micro_batch: Bytes kept for each micro-batch in flight.
49 memory_once: Bytes kept once, however many micro-batches are in flight.
50 forward_time: The forward time.
51 backward_time: The backward time, recompute included.
52 link_bandwidth: For an option that offloads, the bytes each
53 micro-batch moves to the host after its forward and back before
54 its backward; 0 for one that keeps them on the device.
55 names: The names the pipeline balancer and the parsers know the
56 option by: NONE, SLCT, COMM, BOTH or FULL.
57 excess: ``(count, bytes)`` for each count of micro-batches in flight
58 at which the two memories charge more than the layer keeps, by
59 that many bytes; see :attr:`Cost.excess`.
60 working_extra: What the working set of the option's backward holds
61 beyond what it keeps at one micro-batch, the plain layer for full
62 recompute, as warm-up ends; see :attr:`Cost.working`.
63 first_working_extra: The same for a stage's first layer's backward,
64 the last it runs; see :attr:`Cost.first_working`.
65 """
67 recompute: Optional[FrozenSet[str]]
68 memory_per_micro_batch: float
69 memory_once: float
70 forward_time: float
71 backward_time: float
72 link_bandwidth: float = 0.0
73 names: Tuple[str, ...] = ()
74 excess: Tuple[Tuple[int, float], ...] = ()
75 working_extra: float = 0.0
76 first_working_extra: float = 0.0
78 def memory(self, in_flight: int) -> float:
79 """The bytes a layer running the option keeps with *in_flight* micro-batches in flight."""
80 kept = in_flight * self.memory_per_micro_batch + self.memory_once
81 return kept - next((excess for count, excess in self.excess if count == in_flight), 0.0)
83 @property
84 def switches(self) -> Optional[Dict[str, int]]:
85 """The switch setting the option stands for, 1 to keep an op and 0 to recompute it; ``None`` if full."""
86 if self.recompute is None:
87 return None
88 return {name: int(name not in self.recompute) for name in SWITCHES}
91@dataclass(frozen=True)
92class KindFront:
93 """The front of one layer kind of one model.
95 Attributes:
96 model_name: The model the kind belongs to, a sub-model of a
97 multimodal one.
98 kind: The layer kind; ``None`` for a model whose layers are all priced
99 on its config as it stands.
100 options: The kind's options, fastest first.
101 """
103 model_name: str
104 kind: Optional[LayerKind]
105 options: Tuple[LayerOption, ...]
108def _candidates(profile: SwitchProfile) -> Iterator[Tuple[Optional[FrozenSet[str]], Cost]]:
109 """Every option and its cost, those that recompute fewer ops first and full recompute last.
111 The settings are of the switches that act on the kind
112 (:meth:`SwitchProfile.acting`).
113 """
114 switches = profile.acting()
115 for size in range(len(switches) + 1):
116 for names in itertools.combinations(switches, size):
117 yield frozenset(names), profile.selective(names)
118 yield None, profile.full
121def _names(recompute: Optional[FrozenSet[str]], configured: Mapping[str, Any]) -> Tuple[str, ...]:
122 """The names an option is known by; the configured switches name SLCT, COMM and BOTH."""
123 if recompute is None:
124 return ("FULL",)
125 if not recompute:
126 return ("NONE",)
127 selective = frozenset(name for name in SWITCHES if not int(bool(configured.get(name, 1))))
128 names = []
129 if selective and recompute == selective:
130 names.append("SLCT")
131 if recompute == {"gather"}:
132 names.append("COMM")
133 if selective and recompute == selective | {"gather"}:
134 names.append("BOTH")
135 return tuple(names)
138def price_option(
139 profile: SwitchProfile, recompute: Optional[FrozenSet[str]], configured: Optional[Mapping[str, Any]] = None
140) -> LayerOption:
141 """The option recomputing *recompute*, priced from the kind's *profile*.
143 Args:
144 profile: The kind's costs, each switch measured alone.
145 recompute: The switches to recompute; ``None`` for full recompute.
146 configured: The switches the config sets, which name the option.
148 Returns:
149 The option, whether or not it is on the kind's front.
150 """
151 cost = profile.full if recompute is None else profile.selective(recompute)
152 return LayerOption(
153 recompute=None if recompute is None else frozenset(recompute),
154 memory_per_micro_batch=cost.memory_per_micro_batch,
155 memory_once=cost.memory_once,
156 forward_time=profile.forward_time,
157 backward_time=cost.backward_time,
158 names=_names(recompute, configured or {}),
159 excess=tuple((count, excess) for count, excess in zip(profile.counts, cost.excess) if excess),
160 working_extra=cost.working,
161 first_working_extra=cost.first_working,
162 )
165def _workings(recompute: Optional[FrozenSet[str]], cost: Cost, profile: SwitchProfile) -> Tuple[float, float]:
166 """The working sets of the backward of a layer recomputing *recompute*: as warm-up ends, and on a first layer.
168 What the layer keeps at one micro-batch, the plain layer for full
169 recompute, and what the backward holds beyond.
170 """
171 run = profile.plain if recompute is None else cost
172 kept = run.memory_per_micro_batch + run.memory_once
173 return kept + cost.working, kept + cost.first_working
176def _compared(
177 recompute: Optional[FrozenSet[str]], cost: Cost, profile: SwitchProfile
178) -> Tuple[float, ...]:
179 """What an option is compared on: its two memories, backward time, memory at each count and working sets."""
180 at = (
181 count * cost.memory_per_micro_batch + cost.memory_once - excess
182 for count, excess in zip(profile.counts, cost.excess)
183 )
184 return cost.values() + tuple(at) + _workings(recompute, cost, profile)
187def build_front(
188 profile: SwitchProfile, configured: Optional[Mapping[str, Any]] = None
189) -> Tuple[LayerOption, ...]:
190 """The options of a layer kind that no other option beats.
192 An option beats another when it needs no more memory per micro-batch, no
193 more memory once and no more backward time, no more memory at any of the
194 profile's counts of micro-batches in flight, and no larger a working set
195 in its backward. Of options that cost the same, the one that recomputes
196 fewer ops stays. The plain layer and full recompute are always kept.
198 Args:
199 profile: The kind's costs, each switch measured alone.
200 configured: The switches the config sets, which name the options the
201 pipeline balancer knows as SLCT, COMM and BOTH.
203 Returns:
204 The options, fastest first.
205 """
206 candidates: List[Tuple[Optional[FrozenSet[str]], Cost]] = list(_candidates(profile))
207 compared = [_compared(recompute, cost, profile) for recompute, cost in candidates]
208 kept = []
209 for index, (recompute, _) in enumerate(candidates):
210 mine = compared[index]
211 beaten = any(
212 all(theirs <= own for theirs, own in zip(other, mine))
213 and (other != mine or other_index < index)
214 for other_index, other in enumerate(compared)
215 if other_index != index
216 )
217 if not beaten or recompute is None or not recompute:
218 kept.append(price_option(profile, recompute, configured))
219 return tuple(sorted(kept, key=lambda option: (option.backward_time, -option.memory_per_micro_batch)))
222def configured_switches(evaluator: EvaluatorV2) -> Mapping[str, Any]:
223 """The recompute switches the evaluator's config sets."""
224 rec_op = getattr(evaluator.ccfg, "rec_op", None)
225 return vars(rec_op) if rec_op is not None else {}
228def layer_profiles(
229 evaluator: EvaluatorV2,
230 device_type: Any,
231 ccfg: Optional[CustomConfig] = None,
232 most_in_flight: Optional[int] = None,
233 each_switch: bool = True,
234 in_flight: Sequence[int] = (),
235) -> Dict[Tuple[str, Optional[LayerKind]], SwitchProfile]:
236 """What each switch of every layer kind saves and costs, at the evaluator's current strategy.
238 Args:
239 evaluator: The memory evaluator, set to the strategy to price.
240 device_type: The device the times are priced on.
241 ccfg: Estimator options; the search's defaults when omitted.
242 most_in_flight: The most micro-batches any stage keeps in flight, up
243 to which an option's memory is exact; see
244 :meth:`EvaluatorV2.estimate_switch_profiles`.
245 each_switch: Whether to measure each switch alone; without, a profile
246 prices only the plain and the fully recomputed layer.
247 in_flight: The counts of micro-batches in flight the stages keep, at
248 each of which an option's memory is exact too.
250 Returns:
251 ``{(model name, layer kind): SwitchProfile}``, in model order.
252 """
253 return evaluator.estimate_switch_profiles(
254 LayerTimes(device_type, ccfg), most_in_flight=most_in_flight, each_switch=each_switch, in_flight=in_flight
255 )
258def layer_fronts(
259 evaluator: EvaluatorV2,
260 device_type: Any,
261 ccfg: Optional[CustomConfig] = None,
262 most_in_flight: Optional[int] = None,
263) -> Tuple[KindFront, ...]:
264 """The recompute front of every layer kind of the evaluator's model, at its current strategy.
266 Measures each kind's switches with the memory backbone and the
267 performance estimate the search scores with, and leaves the evaluator's
268 config as it found it.
270 Args:
271 evaluator: The memory evaluator, set to the strategy to price.
272 device_type: The device the times are priced on.
273 ccfg: Estimator options; the search's defaults when omitted.
274 most_in_flight: The most micro-batches any stage keeps in flight, up
275 to which an option's memory is exact; see
276 :meth:`EvaluatorV2.estimate_switch_profiles`.
278 Returns:
279 One front per model and layer kind, in model order.
280 """
281 profiles = layer_profiles(evaluator, device_type, ccfg, most_in_flight)
282 configured = configured_switches(evaluator)
283 return tuple(
284 KindFront(model_name, kind, build_front(profile, configured))
285 for (model_name, kind), profile in profiles.items()
286 )