Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / memory_estimation / _ppb.py: 98%
252 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"""PPB input module"""
16from __future__ import annotations
18import contextlib
19from typing import Any, Callable, Dict, Iterator, Mapping, Optional, Sequence, Tuple, TYPE_CHECKING
21from hyper_parallel.auto_parallel._model_spec import KindActivations
22from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config
23from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import HYPER_SELECTIVE_REC_OP
24from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
25from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
26from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.body import EvalBody
27from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.utils import EvalUtils
28from hyper_parallel.auto_parallel.sapp_nd.recompute.profile import OPTIONAL, SWITCHES, Cost, SwitchProfile
30if TYPE_CHECKING:
31 from hyper_parallel.auto_parallel._op_profiles import LayerKind
32 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
34# The balancer's recompute options, in the order that settles a tie between two.
35_OPTIONS = ("NONE", "SLCT", "COMM", "BOTH", "FULL")
36# Where the balancer reads each option's memory per micro-batch and its backward time.
37_MEMORY_KEY = {
38 "NONE": "memory_activation",
39 "SLCT": "memory_select_rec",
40 "COMM": "memory_select_comm",
41 "BOTH": "memory_both_comm_select",
42 "FULL": "memory_recompute",
43}
44_TIME_KEY = {
45 "NONE": "backward_time",
46 "SLCT": "select_rec_time",
47 "COMM": "select_comm_time",
48 "BOTH": "both_comm_select_time",
49 "FULL": "recompute_time",
50}
53def _layer_switches(ccfg: CostModelConfig, ctx: Optional[Context]) -> Dict[str, Any]:
54 """The switches of the layer *ctx* evaluates, its own where the context carries them, else the config's."""
55 stated = EvalUtils.switches(ccfg, ctx)
56 return dict(stated) if isinstance(stated, Mapping) else dict(vars(stated))
59class _PPB:
60 """Pipeline balance payload builder."""
62 def __init__(self, eval_cfg: Config, inner_dyn_fun: Callable) -> None:
63 """Initialize _PPB with evaluation config and dynamic memory function.
65 Args:
66 eval_cfg: Evaluation configuration object.
67 inner_dyn_fun: Function to compute inner dynamic memory.
68 """
69 self.eval_cfg = eval_cfg
70 self._inner_dynamic_mem = inner_dyn_fun
71 self.mb = EvalUtils.mb
72 # Prices a layer as ``(forward, backward)`` from its config, kind,
73 # layer type and recompute switches; set while a description with
74 # times is being built.
75 self.layer_times: Optional[Callable] = None
76 # Each model's and layer kind's switch profile; set while profiles
77 # are being measured.
78 self.profiles: Optional[Dict[tuple, SwitchProfile]] = None
79 # The most micro-batches any stage keeps in flight, where profiles
80 # split the buffers when set.
81 self.profile_in_flight: Optional[int] = None
82 # The counts of micro-batches in flight the stages keep, at which
83 # profiles state what the split charges beyond the buffers.
84 self.profile_counts: Tuple[int, ...] = ()
85 # Whether profiles measure each switch alone, or only the plain and
86 # the fully recomputed layer.
87 self.profile_each_switch = True
89 @contextlib.contextmanager
90 def timing(self, layer_times: Callable) -> Iterator[None]:
91 """Give every layer description its times while the context is open.
93 Args:
94 layer_times: Prices a layer as ``(forward, backward)`` from its
95 config, kind, layer type and recompute switches.
96 """
97 self.layer_times = layer_times
98 try:
99 yield
100 finally:
101 self.layer_times = None
103 @contextlib.contextmanager
104 def profiling(
105 self,
106 layer_times: Callable,
107 most_in_flight: Optional[int] = None,
108 each_switch: bool = True,
109 in_flight: Sequence[int] = (),
110 ) -> Iterator[Dict[tuple, SwitchProfile]]:
111 """Measure each model's and layer kind's switch profile while the context is open.
113 Args:
114 layer_times: Prices a layer, as for :meth:`timing`.
115 most_in_flight: The most micro-batches any stage keeps in flight.
116 each_switch: Whether to measure each switch alone.
117 in_flight: The counts of micro-batches in flight the stages keep.
119 Yields:
120 The profiles, filled as the stages are walked.
121 """
122 self.layer_times = layer_times
123 self.profiles = {}
124 self.profile_in_flight = most_in_flight
125 self.profile_each_switch = each_switch
126 self.profile_counts = tuple(sorted(set(in_flight)))
127 try:
128 yield self.profiles
129 finally:
130 self.layer_times = None
131 self.profiles = None
132 self.profile_in_flight = None
133 self.profile_each_switch = True
134 self.profile_counts = ()
136 @staticmethod
137 def add_to_ppb_list(ppb_lay_desc: list, desc: dict) -> None:
138 """layer description list preparation"""
139 if desc:
140 already_comp = False
141 body_idx = 0
142 for d in ppb_lay_desc:
143 if all(v == d[k] for k, v in desc.items()):
144 # already exist desc
145 d["nb_layer"] += 1
146 already_comp = True
147 if d["type"] == "BODY":
148 body_idx += 1
149 if desc and not already_comp:
150 desc["nb_layer"] = 1
151 if desc["type"] == "BODY":
152 desc["name"] = f"BODY_{body_idx}"
153 else:
154 desc["name"] = desc["type"]
155 ppb_lay_desc += [desc]
157 @staticmethod
158 def selective_switches(ccfg: CostModelConfig, ctx: Optional[Context] = None) -> Dict[str, Dict[str, int]]:
159 """The recompute switches of each selective option.
161 SLCT is the selective recompute the layer runs, its own where *ctx*
162 carries them, else the config's; COMM recomputes the tensor-parallel
163 gathers alone, and BOTH does both.
164 """
165 rec_op = _layer_switches(ccfg, ctx)
166 keep = dict.fromkeys(SWITCHES, 1)
167 configured = {name: int(bool(rec_op.get(name, 1))) for name in SWITCHES}
168 return {"SLCT": configured, "COMM": dict(keep, gather=0), "BOTH": dict(configured, gather=0)}
170 def lay_ppb(
171 self, ccfg: CostModelConfig, ctx: Context, res_stat: float, kind: Optional[LayerKind] = None
172 ) -> dict:
173 """layer description preparation
175 With ``layer_times`` set, a body also offers COMM and BOTH, and the
176 description carries the layer's times, priced on *kind*, the
177 layer's kind. While ``profiles`` are measured, a body is profiled
178 instead and no layer is described.
179 """
180 desc = {"model_name": ccfg.model_name}
181 timed = self.layer_times is not None
182 original_enable_node_log = ctx.enable_node_log
183 ctx.enable_node_log = False
184 try:
185 if self.profiles is not None:
186 if ctx.current_node == ctx.head_node:
187 # The pricer copies a config where it first sees it, which
188 # must come before the walk gives any layer its kind.
189 self.layer_times(ccfg, None, LayerType.EMBEDDING_LAYER)
190 elif ctx.current_node != ctx.tail_node:
191 self._profile(ccfg, ctx, kind)
192 return {}
193 if ctx.current_node == ctx.head_node:
194 d_emb = self.mb(sum(self._inner_dynamic_mem(ppb=True)))
195 desc["type"] = "HEAD"
196 desc["memory_parameter"] = self.mb(res_stat) + d_emb
197 desc["time"] = 1
198 elif ctx.current_node == ctx.tail_node:
199 d_out = self.mb(sum(self._inner_dynamic_mem(ppb=True)))
200 desc["type"] = "TAIL"
201 desc["memory_parameter"] = self.mb(res_stat) + d_out
202 desc["time"] = 1
203 else:
204 self._body_memory(desc, ccfg, ctx, res_stat, timed)
205 finally:
206 ctx.enable_node_log = original_enable_node_log
207 if timed:
208 self._time_ppb(desc, ccfg, ctx, kind)
209 return desc
211 def _body_memory(self, desc: dict, ccfg: CostModelConfig, ctx: Context, res_stat: float, timed: bool) -> None:
212 """Describe a body layer's memory under each option.
214 The balancer charges an option's memory once per micro-batch in
215 flight, and the layer's parameter memory once. Activations are kept
216 per micro-batch. Of the communication buffers, what grows with the
217 micro-batches in flight, such as the tensor-parallel gathers COMM
218 recomputes, is charged per micro-batch, and the rest, the largest
219 any option keeps, once.
220 """
221 many = self._many(ccfg)
222 ctx.current_node = LayerType.NOT_REC_LAYER
223 dyn = {"NONE": self._dynamic_mem(many)}
224 ctx.current_node = LayerType.SEL_REC_LAYER
225 dyn["SLCT"] = self._dynamic_mem(many)
226 if timed:
227 switches = self.selective_switches(ccfg, ctx)
228 for name in ("COMM", "BOTH"):
229 dyn[name] = self._selective(ccfg, ctx, switches[name], lambda _: self._dynamic_mem(many))
230 ctx.current_node = LayerType.FULL_REC_LAYER
231 dyn["FULL"] = self._dynamic_mem(many)
232 desc["type"] = "BODY"
233 desc["memory_parameter"] = self.mb(res_stat) + self.mb(max(once for _, _, once, _ in dyn.values()))
234 for name in _OPTIONS:
235 if name in dyn:
236 activation, per_micro_batch, _, _ = dyn[name]
237 desc[_MEMORY_KEY[name]] = self.mb(activation) + self.mb(per_micro_batch)
238 desc["time"] = 1
240 @staticmethod
241 def _many(ccfg: CostModelConfig) -> int:
242 """The most micro-batches a stage keeps in flight under 1F1B, and at least 2."""
243 return max(2, min(getattr(ccfg, "p", 1), getattr(ccfg, "m", 1)))
245 def _profile(self, ccfg: CostModelConfig, ctx: Context, kind: Optional[LayerKind]) -> None:
246 """Measure the layer plain, with each op alone recomputed and fully recomputed, once per model and kind.
248 Where a census prices the kind, it prices the plain layer and
249 HyperParallel's selective policy, and its records per op, or else
250 the formulas, every other setting: the layer selective with every
251 op kept is measured too, as the base the settings add up from, and
252 the policy's setting whole.
253 """
254 key = (ccfg.model_name, kind)
255 if key in self.profiles:
256 return
257 many = max(2, self.profile_in_flight) if self.profile_in_flight else self._many(ccfg)
258 # The split is exact at one micro-batch in flight and at *many*.
259 counts = tuple(count for count in self.profile_counts if 1 < count < many)
260 keep = dict.fromkeys(SWITCHES, 1)
261 census = isinstance(getattr(ccfg, "kind_activations", None), KindActivations)
263 def _measure(at: Context) -> Tuple[Any, ...]:
264 """The current layer's memory, the working sets of its backward at both points, and their census terms."""
265 return self._dynamic_mem(many, counts) + self._workings(at) + (EvalBody.census_working_terms(ccfg, at),)
267 ctx.current_node = LayerType.NOT_REC_LAYER
268 plain = _measure(ctx)
269 # A fully recomputed layer's backward runs the plain layer, on
270 # activations the stage does not keep; without a census, the working
271 # set is the same either way.
272 full_working = self._working_extra(ctx, 2, on_saved=False) if census else plain[4]
273 base = self._selective(ccfg, ctx, keep, _measure) if census else plain
274 alone = self._alone(ccfg, ctx, base, _measure, many) if self.profile_each_switch else {}
275 policy = dict(HYPER_SELECTIVE_REC_OP)
276 whole = {frozenset(name for name, state in policy.items() if not state): self._selective(
277 ccfg, ctx, policy, _measure)} if census else {}
278 ctx.current_node = LayerType.FULL_REC_LAYER
279 full = self._dynamic_mem(many, counts) + (full_working, plain[5], None)
281 def _cost(memory: Tuple[Any, ...], backward: float) -> Cost:
282 """A measurement as a cost: activations and growing buffers per micro-batch, the rest once.
284 Where a census's records per op price the layer, its working set
285 as warm-up ends before the clamp, and what it keeps of them.
286 """
287 activation, per_micro_batch, once, excess, working, first_working, terms = memory
288 kept = 0.0
289 if terms is not None:
290 kept, held = terms
291 working -= max(0.0, kept - held)
292 return Cost(activation + per_micro_batch, once, backward, excess, working, first_working, kept)
294 forward, backward = self.layer_times(ccfg, kind, LayerType.NOT_REC_LAYER)
295 self.profiles[key] = SwitchProfile(
296 forward_time=forward,
297 plain=_cost(plain, backward),
298 alone=self._alone_costs(ccfg, kind, alone, _cost, _cost(base, backward)),
299 full=_cost(full, self.layer_times(ccfg, kind, LayerType.FULL_REC_LAYER)[1]),
300 counts=counts,
301 # A selective layer that keeps every op recomputes nothing.
302 selective_base=_cost(base, backward) if census else None,
303 census_held=self._census_held(base),
304 whole={
305 recompute: _cost(memory, self.layer_times(ccfg, kind, LayerType.SEL_REC_LAYER, policy)[1])
306 for recompute, memory in whole.items()
307 },
308 )
310 def _alone_costs(
311 self, ccfg: CostModelConfig, kind: Optional[LayerKind], alone: Dict[str, Tuple[Any, ...]],
312 cost: Callable, base: Cost,
313 ) -> Dict[str, Cost]:
314 """Each switch's cost, its op alone recomputed; an optional one the layer keeps nothing under costs *base*."""
315 keep = dict.fromkeys(SWITCHES, 1)
316 costs = {
317 name: cost(memory, self.layer_times(ccfg, kind, LayerType.SEL_REC_LAYER, dict(keep, **{name: 0}))[1])
318 for name, memory in alone.items()
319 }
320 if self.profile_each_switch:
321 costs.update({name: base for name in OPTIONAL if name not in costs})
322 return costs
324 @staticmethod
325 def _census_held(measured: Tuple[Any, ...]) -> Optional[float]:
326 """What the backward of a layer :meth:`_profile` measured holds, where its census's records per op price it."""
327 terms = measured[6]
328 return None if terms is None else terms[1]
330 def _alone(
331 self, ccfg: CostModelConfig, ctx: Context, base: Tuple[Any, ...], measure: Callable, many: int
332 ) -> Dict[str, Tuple[Any, ...]]:
333 """The current layer measured with each op alone recomputed, by *measure* as *base* was.
335 Only gather acts on the buffers, so any other op recomputed alone
336 leaves the split, and the working sets, as the base has them; but
337 where a census prices the layer, what the layer keeps sets its
338 working set as warm-up ends, and each op is measured whole.
339 """
340 keep = dict.fromkeys(SWITCHES, 1)
341 census = isinstance(getattr(ccfg, "kind_activations", None), KindActivations)
342 alone = {}
343 for name in SWITCHES:
344 if name in OPTIONAL and not self._drops_any(ccfg, name):
345 continue
346 switches = dict(keep, **{name: 0})
347 if name == "gather" or census:
348 alone[name] = self._selective(ccfg, ctx, switches, measure)
349 else:
350 measured = self._selective(ccfg, ctx, switches, lambda _: self._dynamic_mem(many))
351 alone[name] = measured[:3] + base[3:]
352 return alone
354 @staticmethod
355 def _drops_any(ccfg: CostModelConfig, switch: str) -> bool:
356 """Whether the current layer keeps anything an optional *switch* drops.
358 attUp, the one optional switch, drops the heads an MLA layer's
359 up-projections build, as the records' part of it states: a layer
360 that compresses its keys and values has them.
361 """
362 return switch != "attUp" or bool(getattr(ccfg, "dc_kv", 0))
364 def _workings(self, ctx: Context) -> Tuple[float, float]:
365 """What the current layer's backward holds beyond what it keeps at one micro-batch, at both points.
367 As warm-up ends, where the stage keeps what the layer keeps for this
368 backward already, and on a stage's first layer as a micro-batch's
369 backward ends, where it keeps none of it.
370 """
371 return self._working_extra(ctx, 2, on_saved=True), self._working_extra(ctx, 1, on_saved=False)
373 def _working_extra(self, ctx: Context, gathered: int, on_saved: bool) -> float:
374 """What the working set of the current layer's backward holds beyond what it keeps at one micro-batch.
376 Under FSDP that reshards, the backward holds *gathered* layers'
377 parameters: two, its own and the next's, in the working set that
378 ends warm-up, and one in the first layer's as the backward ends.
379 *on_saved* says the stage counts what the layer keeps for this
380 backward already, which a census's working set then leaves out.
381 """
382 kept = sum(self._inner_dynamic_mem(ppb=True))
383 ctx.working_set = gathered
384 ctx.working_on_saved = on_saved
385 try:
386 return sum(self._inner_dynamic_mem(default_micro_factor=1)) - kept
387 finally:
388 ctx.working_set = 0
389 ctx.working_on_saved = False
391 def _dynamic_mem(
392 self, many: int, counts: Tuple[int, ...] = ()
393 ) -> Tuple[float, float, float, Tuple[float, ...]]:
394 """``(activation, buffers per micro-batch, buffers once, excess)`` of the current layer.
396 The buffers are split by how they grow from one micro-batch in
397 flight to *many*. At each of *counts*, the excess is what the split
398 charges beyond the buffers kept with that many micro-batches in
399 flight.
400 """
401 at = [self._inner_dynamic_mem(default_micro_factor=count)[1] for count in counts]
402 _, more = self._inner_dynamic_mem(default_micro_factor=many)
403 activation, one = self._inner_dynamic_mem(ppb=True)
404 per_micro_batch = (more - one) / (many - 1)
405 excess = tuple(one + (count - 1) * per_micro_batch - kept for count, kept in zip(counts, at))
406 return activation, per_micro_batch, one - per_micro_batch, excess
408 @staticmethod
409 def _selective(ccfg: CostModelConfig, ctx: Context, switches: Dict[str, int], measure: Callable) -> Any:
410 """*measure* of *ctx*, the current layer selective with *switches*.
412 The context carries them to the evaluators, over the layer's own,
413 and the config's are left as they are.
414 """
415 before = ctx.switches
416 stated = _layer_switches(ccfg, ctx)
417 ctx.current_node = LayerType.SEL_REC_LAYER
418 ctx.switches = {**stated, **switches}
419 try:
420 return measure(ctx)
421 finally:
422 ctx.switches = before
424 def _time_ppb(self, desc: dict, ccfg: CostModelConfig, ctx: Context, kind: Optional[LayerKind]) -> None:
425 """Add the layer's forward time and the backward time of each option, where the balancer reads them."""
426 end = {"HEAD": LayerType.EMBEDDING_LAYER, "TAIL": LayerType.OUTPUT_LAYER}.get(desc["type"])
427 if end is not None:
428 desc["forward_time"], desc["backward_time"] = self.layer_times(ccfg, None, end)
429 else:
430 desc["forward_time"], desc["backward_time"] = self.layer_times(
431 ccfg, kind, LayerType.NOT_REC_LAYER
432 )
433 for name, switches in self.selective_switches(ccfg, ctx).items():
434 desc[_TIME_KEY[name]] = self.layer_times(ccfg, kind, LayerType.SEL_REC_LAYER, switches)[1]
435 desc["recompute_time"] = self.layer_times(ccfg, kind, LayerType.FULL_REC_LAYER)[1]
436 desc["time"] = desc["forward_time"]
438 @staticmethod
439 def _beaten(desc: dict, name: str) -> bool:
440 """Whether another option of *desc* needs no more memory and no more backward time than *name*.
442 Of two options that tie, the one first in :data:`_OPTIONS` wins.
443 """
444 def _cost(option: str) -> tuple:
445 """The option's memory and backward time."""
446 return desc[_MEMORY_KEY[option]], desc[_TIME_KEY[option]]
448 mine = _cost(name)
449 return any(
450 other != name
451 and all(theirs <= own for theirs, own in zip(_cost(other), mine))
452 and (_cost(other) != mine or _OPTIONS.index(other) < _OPTIONS.index(name))
453 for other in _OPTIONS
454 if _MEMORY_KEY[other] in desc
455 )
457 def ppb_withdraw_dominated(self, ppb_lay_desc: list) -> None:
458 """Withdraw the selective options no body is better off with.
460 An option is withdrawn from every body when, in each of them, another
461 option needs no more memory and no more backward time. The balancer
462 takes its options from the first body, so every body keeps the same
463 ones. Descriptions without times are left as they are.
464 """
465 bodies = [d for d in ppb_lay_desc if d["type"] == "BODY"]
466 if not bodies or any("backward_time" not in d for d in bodies):
467 return
468 withdrawn = [name for name in ("SLCT", "COMM", "BOTH") if all(self._beaten(d, name) for d in bodies)]
469 for d in bodies:
470 for name in withdrawn:
471 del d[_MEMORY_KEY[name]], d[_TIME_KEY[name]]
473 @staticmethod
474 def ppb_scale_times(ppb_lay_desc: list) -> None:
475 """Express the times in units of the first body's forward time.
477 The balancer only compares times with one another, and its solver
478 finds feasible problems infeasible at the estimate's own magnitudes,
479 around 1e13. Descriptions without times are left as they are.
480 """
481 unit = next((d["forward_time"] for d in ppb_lay_desc if d["type"] == "BODY" and d.get("forward_time")), None)
482 if not unit:
483 return
484 for d in ppb_lay_desc:
485 for key in d:
486 if key == "time" or key.endswith("_time"):
487 d[key] /= unit
489 def ppb_combine_bodies(self, ppb_lay_desc: list) -> None:
490 """combine descriptions into a new body
492 Memories add up over the combined bodies, and so do the times of a
493 timed description.
494 """
495 if not self.eval_cfg.ppb_combined:
496 return
497 for new_body in self.eval_cfg.ppb_combined:
498 desc = {
499 "model_name": "combined",
500 "type": "BODY",
501 "memory_parameter": 0,
502 "memory_activation": 0,
503 "memory_recompute": 0,
504 "time": 1,
505 "nb_layer": 1,
506 "name": "COMBINED",
507 }
508 idx = -1
509 for mod, t in new_body:
510 target = next(
511 (
512 d
513 for d in ppb_lay_desc
514 if d["model_name"] == mod and d["type"] == t.upper()
515 ),
516 None,
517 )
518 if target:
519 desc["model_name"] += "_" + mod
520 desc["name"] += "_" + target["name"]
521 for key, value in target.items():
522 if key.startswith("memory") or key.endswith("_time"):
523 desc[key] = desc.get(key, 0) + value
524 target_idx = ppb_lay_desc.index(target)
525 idx = target_idx if idx < 0 else min(idx, target_idx)
526 del ppb_lay_desc[target_idx]
527 if "forward_time" in desc:
528 desc["time"] = desc["forward_time"]
529 idx = max(idx, 0)
530 ppb_lay_desc.insert(idx, desc)