Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / recompute / candidate.py: 97%
616 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 every layer's recompute option for one candidate of the search.
17The search keeps the candidates that fit with every layer fully recomputed,
18the lightest policy. :func:`choose_recompute` gives a candidate's layers the
19fastest options that fit instead. Each pipeline stage's layers share what the
20stage has left under the config's own recompute, plus what those layers keep
21under it: the memory model's peak for the stage is the sum of what the choice
22does not change and of what each layer keeps at the micro-batches the stage
23holds of it in flight. :func:`~.knapsack.pp_lite` then picks the options,
24stage by stage.
26At the end of warm-up, the memory model also charges each stage one layer's
27working set at one micro-batch: the plain layer's if the stage's last layer is
28fully recomputed, the layer's own otherwise, what it keeps then and, under FSDP
29that reshards, the gathered parameters its backward holds besides. So that
30layer's options each keep the working set they end warm-up on too, and the
31knapsack weighs it with them.
33A runtime that runs every layer one way, such as HyperParallel's trainer with
34its ``activation_checkpoint.mode``, gets the fastest of its modes that fits
35instead, from the same budgets: see :data:`MODES`. One that runs a mode per
36layer, as the same trainer does with ``activation_checkpoint.layers``, gets
37each layer the fastest of its modes, chosen as a layer's options are, and
38:func:`trainer_plan` states the choice as the trainer reads it, where the
39installed trainer declares that field (:func:`trainer_runs_layer_plans`). A model
40priced on several submodules, a vision tower beside a language model, gets no
41choice per layer, whose budgets would be built from one submodule's layers;
42its runtime's mode is stated on every submodule, and the whole model priced
43under each (:func:`mode_ranges`, :func:`whole_modes`).
45With a host link, a choice per layer may also offload: a stage's first layers
46run plain and move what they keep per micro-batch to the host, tensor by
47tensor, and back before their backward, so that they keep only the rest. A
48micro-batch's copies to the host keep pace with its own forward: one stream
49copies what each layer keeps once the layer's forward ends, in the order the
50layers run, and is done when the forward is, in the link's share of the
51forward left. So a stage's first layers move all or part of what they keep,
52and its last layers nothing when their window is too short. The copies back
53run in the longer backward, one layer ahead, while this micro-batch's
54backward has released what the layers after the next one keep; where that
55does not cover them, the stage keeps the largest layer's offloaded bytes in
56transit. Each byte moved costs the step the link's copy cost, out in the
57forward and back in the backward, so where the copies cost time a stage
58offloads no more than it needs to fit. Offload is priced at one chunk per
59stage.
61Where a stage's FSDP holds each layer's gradient output until the backward
62ends, and accumulates gradients over several micro-batches, its peak can also
63come as a later micro-batch's backward ends: the memory model then counts
64what the micro-batches still in flight keep, beside every layer's gradient
65output and the working set of the first layer's backward, the last the stage
66runs. The stage's layers then fit at both points. Each pair of options of the
67first layer and of the layer that ends warm-up leaves the others one budget,
68since they keep a share of what they keep as warm-up ends: exact at one chunk
69per stage, and with two, the largest share of the chunks'.
70"""
71import importlib
72import itertools
73import math
74from collections import Counter
75from dataclasses import dataclass, fields, is_dataclass, replace
76from types import SimpleNamespace
77from typing import Any, Dict, FrozenSet, Hashable, Iterator, List, Mapping, Optional, Sequence, Tuple
79from hyper_parallel.auto_parallel._exec_spec import RECOMPUTE_OPS, RecomputeRange
80from hyper_parallel.auto_parallel._model_spec import KindActivations
81from hyper_parallel.auto_parallel._op_profiles import LayerKind
82from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.estimate_v2 import EvaluatorV2
83from hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware import HostLink
84from hyper_parallel.auto_parallel.sapp_nd.nd.common.arch_hooks import layer_kinds
85from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
86from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.utils_classes import CustomConfig
87from hyper_parallel.auto_parallel.sapp_nd.recompute.front import (
88 LayerOption,
89 build_front,
90 configured_switches,
91 layer_profiles,
92 price_option,
93)
94from hyper_parallel.auto_parallel.sapp_nd.recompute.knapsack import (
95 MEGABYTE,
96 Assignment,
97 Layers,
98 Stage,
99 StageChoice,
100 choose,
101 least_times,
102 suffix_times,
103)
104from hyper_parallel.auto_parallel.sapp_nd.recompute.profile import SWITCHES
106# The pipeline schedules whose micro-batches in flight and end-of-warm-up
107# working set the budget follows.
108SCHEDULES = ("1f1b",)
109# The ways a runtime can run every layer: nothing recomputed, the switches the
110# config sets recomputed, or everything recomputed.
111MODES = ("off", "selective", "full")
112_ENDS = (LayerType.EMBEDDING_LAYER, LayerType.OUTPUT_LAYER)
113# Marks the front of a layer that ends warm-up, whose options keep its working set.
114_ENDS_WARM_UP = "ends warm-up"
115# A time shorter than another by this share of it or less is the same time: the
116# same options' times summed in another order can differ in their last bits.
117_ROUNDING = 1e-9
120@dataclass(frozen=True)
121class LayerRange:
122 """Consecutive layers of one kind that run one option.
124 Attributes:
125 first: The first layer's index, counting the model's layers in order
126 without the embedding and output layers.
127 count: How many layers.
128 kind: Their kind; ``None`` for a model priced on its config as it
129 stands.
130 option: The option they run.
131 mode: For a choice of a mode per layer, the mode they run, one of
132 :data:`MODES`; ``None`` otherwise.
133 """
135 first: int
136 count: int
137 kind: Optional[LayerKind]
138 option: LayerOption
139 mode: Optional[str] = None
142@dataclass(frozen=True)
143class RecomputeChoice:
144 """Every layer's option for one candidate, and what the options change.
146 Attributes:
147 ranges: The layers' options, in model order.
148 stage_memory: Each stage's peak memory under the options, in MB.
149 stage_savings: The time each stage's layers save against the
150 recompute the config gives them, in the performance estimate's
151 units.
152 mode: The mode every layer runs, one of :data:`MODES`, when one mode
153 was chosen for all of them; ``None`` for a choice per layer.
154 score: The candidate's score under the choice, where the whole model
155 was priced with it, a model priced on several submodules, or
156 where a caller priced one alternative on its own; ``None`` where
157 the search takes *stage_savings* off the score of the config's
158 own recompute.
159 dimensions: The parallel degrees the choice was priced at, for an
160 alternative weighed against others at one fixed strategy; empty
161 for a choice the search did not weigh that way.
162 score_parts: The parts *score* splits into, in the order the
163 debugger fills them, where a caller kept them for a plot.
164 """
166 ranges: Tuple[LayerRange, ...]
167 stage_memory: Tuple[float, ...]
168 stage_savings: Tuple[float, ...]
169 mode: Optional[str] = None
170 score: Optional[float] = None
171 dimensions: Tuple[Tuple[str, int], ...] = ()
172 score_parts: Tuple[float, ...] = ()
174 @property
175 def memory(self) -> float:
176 """The peak memory in MB, the heaviest stage's."""
177 return max(self.stage_memory)
180@dataclass(frozen=True)
181class _Layer:
182 """A body layer: where it sits, its front's key and how the config runs it."""
184 index: int
185 key: Hashable
186 in_flight: int
187 own: LayerOption
190@dataclass(frozen=True)
191class _Peaks:
192 """What a stage keeps whatever its layers run, at each point its peak can come, in bytes.
194 Attributes:
195 warm_up: As warm-up ends.
196 backward: As a later micro-batch's backward ends, for a stage whose
197 FSDP holds its gradient outputs until then; ``None`` otherwise.
198 """
200 warm_up: float
201 backward: Optional[float] = None
204def _kept(option: LayerOption, in_flight: int) -> float:
205 """The bytes a layer running *option* keeps with *in_flight* micro-batches in flight."""
206 return option.memory(in_flight)
209def _share(in_flight: int) -> float:
210 """The share of what a layer keeps that stays as one of its *in_flight* micro-batches' backward ends."""
211 return (in_flight - 1) / in_flight
214def _time(option: LayerOption) -> float:
215 """A layer's forward and backward time under *option*."""
216 return option.forward_time + option.backward_time
219def micro_batches_in_flight(evaluator: EvaluatorV2) -> List[List[int]]:
220 """How many micro-batches each stage keeps in flight of each chunk's layers.
222 Counted as the memory model counts them, with the schedule's own formula
223 and at least one.
225 Args:
226 evaluator: The memory evaluator, set to the candidate's strategy.
228 Returns:
229 ``[stage][chunk]`` counts.
230 """
231 ccfg = evaluator.ccfg
232 count = evaluator.ctx.pp_micro_eval[ccfg.pp_sched]
233 less_memory = evaluator.ctx.vpp_less_mem
235 def _at(stage: int, chunk: int) -> int:
236 """The count of one stage's chunk, on a context of its own."""
237 where = SimpleNamespace(current_stage_id=stage, current_chunk_id=chunk, vpp_less_mem=less_memory)
238 return max(1, count(ccfg, where))
240 return [[_at(stage, chunk) for chunk in range(ccfg.vp)] for stage in range(ccfg.p)]
243def mode_recompute(mode: str, configured: Mapping[str, Any]) -> Optional[FrozenSet[str]]:
244 """The switches a mode recomputes; ``None`` for full recompute.
246 Args:
247 mode: One of :data:`MODES`.
248 configured: The switches the config sets, 1 to keep an op and 0 to
249 recompute it, which say what ``selective`` recomputes.
251 Returns:
252 The switches.
254 Raises:
255 ValueError: For a mode not in :data:`MODES`.
256 """
257 if mode == "off":
258 return frozenset()
259 if mode == "full":
260 return None
261 if mode == "selective":
262 return frozenset(name for name in SWITCHES if not int(bool(configured.get(name, 1))))
263 raise ValueError(f"unknown recompute mode {mode!r}; expected one of {', '.join(MODES)}")
266def mode_ranges(mode: str, configured: Mapping[str, Any]) -> Tuple[RecomputeRange, ...]:
267 """Every layer of a model running *mode*, as the recompute ranges an ExecSpec states.
269 Args:
270 mode: One of :data:`MODES`.
271 configured: The switches the selective mode sets, 1 to keep an op and
272 0 to recompute it.
274 Returns:
275 No range for a model that recomputes nothing, else one range over
276 every layer: full, or selective with each op's state.
277 """
278 recompute = mode_recompute(mode, configured)
279 if recompute is None:
280 return (RecomputeRange(option="full"),)
281 if not recompute:
282 return ()
283 return (RecomputeRange(option="selective",
284 ops={op: "recompute" if op in recompute else "keep" for op in RECOMPUTE_OPS}),)
287def census_prices(ccfg: Any) -> bool:
288 """Whether a census prices every layer kind of *ccfg*, HyperParallel's selective policy among what it measures."""
289 kinds = layer_kinds(ccfg)
290 if not kinds:
291 return False
292 if any(kind is None for kind in kinds):
293 return isinstance(getattr(ccfg, "kind_activations", None), KindActivations)
294 census = getattr(ccfg, "census", None)
295 census = census if isinstance(census, Mapping) else {}
296 return all(isinstance(census.get(kind.name), KindActivations) for kind in kinds)
299def whole_modes(
300 configs: Sequence[Any], modes: Sequence[str], selective: Optional[Mapping[str, int]]
301) -> Tuple[str, ...]:
302 """*modes*, less a runtime's own selective mode unless a census prices every kind of every one of *configs*.
304 The modes a model priced whole is offered, as :func:`choose_recompute`
305 offers them a model priced layer by layer: the formulas do not price the
306 matmuls HyperParallel's policy recomputes.
307 """
308 if selective is None or "selective" not in modes or all(census_prices(ccfg) for ccfg in configs):
309 return tuple(modes)
310 return tuple(mode for mode in modes if mode != "selective")
313def _own(node: LayerType, options: Sequence[LayerOption]) -> Optional[LayerOption]:
314 """The option a layer the partition gives *node* runs, if its kind's front has it."""
315 for option in options:
316 if node == LayerType.FULL_REC_LAYER and option.recompute is None:
317 return option
318 if node == LayerType.NOT_REC_LAYER and option.recompute == frozenset():
319 return option
320 if node == LayerType.SEL_REC_LAYER and "SLCT" in option.names:
321 return option
322 return None
325def _overhead_position(chunk: Sequence[LayerType], has_mtp: bool) -> Optional[int]:
326 """The position in a stage's last chunk of the layer whose working set warm-up ends on, if a body layer."""
327 if not chunk:
328 return None
329 if chunk[-1] not in _ENDS:
330 return len(chunk) - 1
331 if chunk[-1] == LayerType.OUTPUT_LAYER and has_mtp and len(chunk) > 1 and chunk[-2] not in _ENDS:
332 return len(chunk) - 2
333 return None
336def _body_layers(
337 evaluator: EvaluatorV2, fronts: Dict[Hashable, Tuple[LayerOption, ...]], counts: List[List[int]]
338) -> Optional[Tuple[List[List[_Layer]], List[List[_Layer]]]]:
339 """Each stage's body layers, and the layers each stage's warm-up ends on.
341 The memory model visits chunk after chunk, and within a chunk stage after
342 stage, which is the model's order of layers; it gives each body layer the
343 next kind in that order.
345 Returns:
346 ``(layers, ends)`` per stage, or ``None`` when some layer runs an
347 option its kind's front does not have.
348 """
349 ccfg = evaluator.ccfg
350 partition = ccfg.generate_partitions_vpp()
351 kinds = layer_kinds(ccfg)
352 layers = [[] for _ in range(ccfg.p)]
353 at = {}
354 index = 0
355 for chunk in range(ccfg.vp):
356 for stage in range(ccfg.p):
357 for position, node in enumerate(partition[stage][chunk]):
358 if node in _ENDS:
359 continue
360 key = (ccfg.model_name, kinds[index])
361 own = _own(node, fronts[key])
362 if own is None:
363 return None
364 layer = _Layer(index, key, counts[stage][chunk], own)
365 layers[stage].append(layer)
366 at[stage, chunk, position] = layer
367 index += 1
368 ends = [[] for _ in range(ccfg.p)]
369 if not ccfg.freeze:
370 last = ccfg.vp - 1
371 for stage in range(ccfg.p):
372 position = _overhead_position(partition[stage][last], ccfg.n_mtp > 0)
373 if position is not None:
374 ends[stage].append(at[stage, last, position])
375 return layers, ends
378def _plain(options: Sequence[LayerOption]) -> LayerOption:
379 """A kind's plain option."""
380 return next(option for option in options if option.recompute == frozenset())
383def _working(option: LayerOption, options: Sequence[LayerOption], first: bool = False) -> float:
384 """The working set a layer running *option* ends warm-up on, or with *first*, the backward on a stage's first layer.
386 What it keeps at one micro-batch, the plain layer for full recompute,
387 *options* being its kind's, and what its backward holds beyond.
388 """
389 run = option if option.recompute is not None else _plain(options)
390 return _kept(run, 1) + (option.first_working_extra if first else option.working_extra)
393def _charge_working_sets(
394 layers: List[List[_Layer]],
395 ends: Sequence[Sequence[_Layer]],
396 fronts: Dict[Hashable, Tuple[LayerOption, ...]],
397 by_mode: Optional[Dict[Hashable, Dict[str, LayerOption]]],
398) -> Tuple[Any, ...]:
399 """Give each layer that ends warm-up a front whose options keep its working set too.
401 Its kind's options, each with the working set it ends warm-up on added to
402 what it keeps once, so that the knapsack weighs the working set with the
403 option instead of keeping room for the heaviest.
405 Returns:
406 ``(layers, fronts, by_mode, own)``: the layers, those that end warm-up
407 on their new front; the fronts and the options by mode with the new
408 fronts added; and each added option's option of its kind.
409 """
410 fronts = dict(fronts)
411 by_mode = None if by_mode is None else dict(by_mode)
412 own = {}
413 ending = {layer.index for stage_ends in ends for layer in stage_ends}
414 for stage_layers in layers:
415 for position, layer in enumerate(stage_layers):
416 if layer.index not in ending:
417 continue
418 key = layer.key + (_ENDS_WARM_UP,)
419 options = fronts[layer.key]
420 charged = {
421 option: replace(option, memory_once=option.memory_once + _working(option, options))
422 for option in options
423 }
424 fronts[key] = tuple(charged.values())
425 if by_mode is not None:
426 by_mode[key] = {mode: charged[option] for mode, option in by_mode[layer.key].items()}
427 own.update({new: old for old, new in charged.items()})
428 stage_layers[position] = replace(layer, key=key, own=charged[layer.own])
429 return layers, fronts, by_mode, own
432def _kind_key(key: Hashable) -> Hashable:
433 """The key of a layer's kind's own front, for a layer on a front charged its working set."""
434 return key[:-1] if isinstance(key, tuple) and key and key[-1] == _ENDS_WARM_UP else key
437def _first_working(
438 layer: _Layer,
439 option: LayerOption,
440 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
441 own: Mapping[LayerOption, LayerOption],
442) -> float:
443 """The working set of a stage's first layer's backward, the last the stage runs, under *option*.
445 As for the layer that ends warm-up (:func:`_working`); for an offloaded
446 layer, the plain layer's, whose activations come back for it.
447 """
448 options = fronts[_kind_key(layer.key)]
449 option = own.get(option, option)
450 return _working(_plain(options) if option.link_bandwidth else option, options, first=True)
453def _backward_kept(
454 stage_layers: Sequence[_Layer],
455 chosen: Mapping[int, LayerOption],
456 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
457 own: Mapping[LayerOption, LayerOption],
458) -> float:
459 """What a stage's layers keep as a later micro-batch's backward ends, running *chosen*.
461 Each keeps what it keeps of the micro-batches still in flight, a share of
462 what it keeps as warm-up ends, and the first layer the working set of its
463 backward, the last the stage runs.
464 """
465 if not stage_layers:
466 return 0.0
467 kept = sum(
468 _kept(own.get(chosen[layer.index], chosen[layer.index]), layer.in_flight) * _share(layer.in_flight)
469 for layer in stage_layers
470 )
471 first = stage_layers[0]
472 return kept + _first_working(first, chosen[first.index], fronts, own)
475def _stage_memory(
476 stage_layers: Sequence[_Layer],
477 peaks: _Peaks,
478 chosen: Mapping[int, LayerOption],
479 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
480 own: Mapping[LayerOption, LayerOption],
481) -> float:
482 """A stage's peak in bytes, its layers running *chosen*: the higher of the points it can come at."""
483 memory = peaks.warm_up + sum(_kept(chosen[layer.index], layer.in_flight) for layer in stage_layers)
484 if peaks.backward is None:
485 return memory
486 return max(memory, peaks.backward + _backward_kept(stage_layers, chosen, fronts, own))
489def _groups(layers: Sequence[_Layer]) -> Tuple[Layers, ...]:
490 """Layers grouped by kind and micro-batches in flight, for the knapsack."""
491 groups: Dict[Tuple[Hashable, int], int] = {}
492 for layer in layers:
493 groups[layer.key, layer.in_flight] = groups.get((layer.key, layer.in_flight), 0) + 1
494 return tuple(Layers(key, count, flight) for (key, flight), count in groups.items())
497def _stage(layers: Sequence[_Layer], peak: float, capacity: float) -> Tuple[Stage, float]:
498 """A stage for the knapsack, and the bytes it keeps outside the choice as warm-up ends.
500 Args:
501 layers: The stage's body layers, the one that ends warm-up charged
502 its working set (:func:`_charge_working_sets`).
503 peak: The memory model's stage memory as warm-up ends, in MB.
504 capacity: The device's memory, in MB.
506 Returns:
507 ``(stage, fixed)``: *fixed* is what the stage keeps whatever its
508 layers run.
509 """
510 fixed = peak * MEGABYTE - sum(_kept(layer.own, layer.in_flight) for layer in layers)
511 return Stage(groups=_groups(layers), budget=capacity * MEGABYTE - fixed), fixed
514def _assign(layers: Sequence[_Layer], choices: Sequence[StageChoice]) -> Dict[int, LayerOption]:
515 """Each layer's option: a group's layers take its options in model order, the slower first."""
516 by_group: Dict[Tuple[Hashable, int], List[int]] = {}
517 for layer in layers:
518 by_group.setdefault((layer.key, layer.in_flight), []).append(layer.index)
519 chosen = {}
520 for choice in choices:
521 for assignment in choice.assignments:
522 waiting = by_group[assignment.layers.kind, assignment.layers.in_flight]
523 for _ in range(assignment.count):
524 chosen[waiting.pop(0)] = assignment.option
525 return chosen
528def _assign_counted(layers: Sequence[_Layer], chosen: Mapping[int, LayerOption]) -> Dict[int, LayerOption]:
529 """Assign a group's option counts to its layers in the order the normal optimizer uses.
531 The knapsack hands a group's options to its layers in model order, the
532 slower first (:func:`_assign`), so a plan built by changing one layer is
533 re-assigned the same way: two plans that run the same counts then name
534 the same layers whichever layer was changed.
536 Args:
537 layers: One group's layers, of one kind, stage and in-flight count.
538 chosen: Each layer's option, the group's among them.
540 Returns:
541 The group's layers and the option each runs.
542 """
543 group = Layers(layers[0].key, len(layers), layers[0].in_flight)
544 counts = Counter(chosen[layer.index] for layer in layers)
545 assignments = tuple(
546 Assignment(group, option, count)
547 for option, count in sorted(counts.items(), key=lambda item: -_time(item[0]))
548 )
549 choice = StageChoice(
550 assignments,
551 sum(count * _kept(option, group.in_flight) for option, count in counts.items()),
552 sum(count * _time(option) for option, count in counts.items()),
553 )
554 return _assign(layers, [choice])
557def _equivalent_cost(left: RecomputeChoice, right: RecomputeChoice) -> bool:
558 """Whether two plans of the same mode counts also have equal stage peaks and time saved."""
559 return all(
560 math.isclose(first, second, rel_tol=1e-12, abs_tol=1e-9)
561 for first, second in zip(left.stage_memory, right.stage_memory)
562 ) and all(
563 math.isclose(first, second, rel_tol=_ROUNDING, abs_tol=1e-9)
564 for first, second in zip(left.stage_savings, right.stage_savings)
565 )
568def _ranges(
569 layers: Sequence[_Layer], chosen: Dict[int, LayerOption], modes: Optional[Mapping[LayerOption, str]] = None
570) -> Tuple[LayerRange, ...]:
571 """The options as ranges of consecutive layers of one kind, in model order, with *modes* each option's mode."""
572 ranges = []
573 for layer in sorted(layers, key=lambda layer: layer.index):
574 option = chosen[layer.index]
575 kind = layer.key[1]
576 last = ranges[-1] if ranges else None
577 if last is not None and last.option == option and last.kind == kind and last.first + last.count == layer.index:
578 ranges[-1] = replace(last, count=last.count + 1)
579 else:
580 ranges.append(LayerRange(layer.index, 1, kind, option, None if modes is None else modes[option]))
581 return tuple(ranges)
584def _result(
585 layers: Sequence[Sequence[_Layer]],
586 peaks: Sequence[_Peaks],
587 chosen: Dict[int, LayerOption],
588 mode: Optional[str],
589 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
590 own: Mapping[LayerOption, LayerOption],
591 modes: Optional[Mapping[LayerOption, str]] = None,
592) -> RecomputeChoice:
593 """The choice of *chosen* options, with each stage's memory and time saved.
595 *own* maps an option charged a working set to its option of its kind,
596 which the ranges report; for a choice of a mode per layer, *modes* maps
597 each kind's option to its mode, which they report too.
598 """
599 every = [layer for stage_layers in layers for layer in stage_layers]
600 return RecomputeChoice(
601 ranges=_ranges(every, {index: own.get(option, option) for index, option in chosen.items()}, modes),
602 stage_memory=tuple(
603 _stage_memory(stage_layers, stage_peaks, chosen, fronts, own) / MEGABYTE
604 for stage_layers, stage_peaks in zip(layers, peaks)
605 ),
606 stage_savings=tuple(
607 sum(_time(layer.own) - _time(chosen[layer.index]) for layer in stage_layers) for stage_layers in layers
608 ),
609 mode=mode,
610 )
613def _one_mode(
614 layers: Sequence[Sequence[_Layer]],
615 peaks: Sequence[_Peaks],
616 by_mode: Dict[Hashable, Dict[str, LayerOption]],
617 modes: Sequence[str],
618 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
619 own: Mapping[LayerOption, LayerOption],
620 capacity: float,
621) -> Optional[Tuple[str, Dict[int, LayerOption]]]:
622 """The fastest of *modes* that every stage fits when every layer runs it, and each layer's option."""
623 fastest, best = math.inf, None
624 for mode in modes:
625 chosen = {layer.index: by_mode[layer.key][mode] for stage_layers in layers for layer in stage_layers}
626 fits = all(
627 _stage_memory(stage_layers, stage_peaks, chosen, fronts, own) <= capacity
628 for stage_layers, stage_peaks in zip(layers, peaks)
629 )
630 time = sum(_time(option) for option in chosen.values())
631 if fits and time < fastest:
632 fastest, best = time, (mode, chosen)
633 return best
636def _priced_modes(
637 modes: Optional[Sequence[str]], profiles: Mapping[Hashable, Any], selective: Optional[Mapping[str, int]]
638) -> Optional[Sequence[str]]:
639 """*modes*, less a runtime's own selective mode unless every kind's profile measures its setting whole."""
640 if modes is None or selective is None or "selective" not in modes:
641 return modes
642 setting = mode_recompute("selective", selective)
643 if all(setting in profile.whole for profile in profiles.values()):
644 return modes
645 return tuple(mode for mode in modes if mode != "selective")
648def _offered(
649 profiles: Dict[Hashable, Any],
650 configured: Mapping[str, Any],
651 modes: Optional[Sequence[str]],
652 selective: Optional[Mapping[str, int]] = None,
653) -> Tuple[Dict[Hashable, Tuple[LayerOption, ...]], Optional[Dict[Hashable, Dict[str, LayerOption]]]]:
654 """Each kind's options: its front, or with *modes*, the option of each mode the profiles can price.
656 The selective mode recomputes what *selective* sets, else the config's
657 switches.
659 Returns:
660 ``(options, by_mode)``: each kind's options, and with *modes* each
661 kind's option by mode.
662 """
663 if modes is None:
664 return {key: build_front(profile, configured) for key, profile in profiles.items()}, None
665 offered = MODES if "selective" in modes else tuple(mode for mode in MODES if mode != "selective")
666 switches = {mode: selective if mode == "selective" and selective is not None else configured for mode in offered}
667 by_mode = {
668 key: {mode: price_option(profile, mode_recompute(mode, switches[mode]), configured) for mode in offered}
669 for key, profile in profiles.items()
670 }
671 return {key: tuple(options.values()) for key, options in by_mode.items()}, by_mode
674@dataclass(frozen=True)
675class _Link:
676 """A host link in the performance estimate's time units.
678 Attributes:
679 per_byte: The time one byte takes over the link, one way.
680 overlap: The share of the compute time the copies may take.
681 cost_per_byte: The step time one byte moved costs, either way.
682 """
684 per_byte: float
685 overlap: float
686 cost_per_byte: float = 0.0
689def _link(link: HostLink, evaluator: EvaluatorV2) -> _Link:
690 """*link* in the estimate's units: FLOPs of forward and backward, times the precision's bytes, per device.
692 A link calibrated with a COMPUTE ratio converts with it instead
693 (:meth:`HostLink.units_per_second`), and the copies' cost with the
694 score's own ratio where one is stated
695 (:meth:`HostLink.score_units_per_second`).
696 """
697 bytes_p = evaluator.ccfg.bytes_p
698 return _Link(link.seconds_per_byte() * link.units_per_second(bytes_p), link.overlap,
699 link.cost_seconds_per_byte() * link.score_units_per_second(bytes_p))
702def _offloaded(option: LayerOption, share: float, cost_per_byte: float = 0.0) -> LayerOption:
703 """*option* with *share* of what it keeps per micro-batch moved to the host, so that it keeps the rest.
705 The copies cost the step *cost_per_byte* a byte, out in the forward and
706 back in the backward.
707 """
708 return replace(option, memory_per_micro_batch=option.memory_per_micro_batch - share, link_bandwidth=share,
709 forward_time=option.forward_time + share * cost_per_byte,
710 backward_time=option.backward_time + share * cost_per_byte, names=())
713def _offloadable(option: LayerOption, in_flight: int) -> float:
714 """The bytes a layer running *option* keeps per micro-batch, less any the two memories charge beyond them."""
715 return max(0.0, _kept(option, in_flight) - option.memory_once) / in_flight
718def _offload_shares(stage_layers: Sequence[_Layer], plain: Sequence[LayerOption], link: _Link) -> List[float]:
719 """The bytes per micro-batch each of a stage's first layers can move to the host running plain.
721 One stream copies what each layer keeps once the layer's forward ends, in
722 the order the layers run, and is done when the forward is, in the link's
723 share of the forward left: each layer moves what the window after it
724 leaves, all it keeps or part of it, once the earlier layers' copies take
725 their part. The layer that ends warm-up, and every layer after it, keep
726 their activations.
728 Returns:
729 A share for each layer before the one that ends warm-up, 0 for one
730 whose window is spent.
731 """
732 ending = next((position for position, layer in enumerate(stage_layers) if _ENDS_WARM_UP in layer.key),
733 len(stage_layers))
734 after = list(itertools.accumulate((option.forward_time for option in reversed(plain[1:])), initial=0.0))[::-1]
735 shares, room = [], math.inf
736 for position in range(ending):
737 room = min(room, link.overlap * after[position])
738 whole = _offloadable(plain[position], stage_layers[position].in_flight)
739 share = whole if link.per_byte == 0 else min(whole, room / link.per_byte)
740 room = max(0.0, room - share * link.per_byte)
741 shares.append(share)
742 return shares
745def _backward_can_bind(
746 stage_layers: Sequence[_Layer],
747 room: float,
748 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
749 own: Mapping[LayerOption, LayerOption],
750) -> bool:
751 """Whether some choice of a stage's layers keeps more than *room* bytes as a later micro-batch's backward ends."""
752 first = stage_layers[0]
753 most = max(
754 _kept(own.get(option, option), first.in_flight) * _share(first.in_flight)
755 + _first_working(first, option, fronts, own)
756 for option in fronts[first.key]
757 )
758 most += sum(
759 max(_kept(own.get(option, option), layer.in_flight) for option in fronts[layer.key]) * _share(layer.in_flight)
760 for layer in stage_layers[1:]
761 )
762 return most > room
765def _ends(stage_layers: Sequence[_Layer]) -> List[_Layer]:
766 """A stage's first layer and the layer that ends its warm-up, once each."""
767 first = stage_layers[0]
768 ending = next((layer for layer in stage_layers if _ENDS_WARM_UP in layer.key), first)
769 return [first] if ending.index == first.index else [first, ending]
772def _middle_budget(
773 ends: Sequence[_Layer],
774 chosen: Mapping[int, LayerOption],
775 budgets: Tuple[float, float],
776 share: float,
777 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
778 own: Mapping[LayerOption, LayerOption],
779) -> float:
780 """What a stage's other layers may keep as warm-up ends, its ends running *chosen*; negative if none fits.
782 *budgets* are the bytes all its layers may keep as warm-up ends and as a
783 later micro-batch's backward ends, and *share* the share the other layers
784 keep of the one at the other.
785 """
786 left = budgets[0] - sum(_kept(chosen[end.index], end.in_flight) for end in ends)
787 left_back = budgets[1] - _backward_kept(ends, chosen, fronts, own)
788 if left < 0 or left_back < 0:
789 return -1.0
790 return min(left, left_back / share) if share else left
793def _two_peaks(
794 stage_layers: Sequence[_Layer],
795 budget: float,
796 room: float,
797 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
798 own: Mapping[LayerOption, LayerOption],
799 bucket: float,
800) -> Optional[Dict[int, LayerOption]]:
801 """The fastest options of a stage's layers within *budget* bytes as warm-up ends and *room* as a backward ends.
803 What the layers keep as a later micro-batch's backward ends is a share of
804 what they keep as warm-up ends, but for the first layer, which adds its
805 working set, and the layer that ends warm-up, which leaves its own out.
806 So each pair of options of those two leaves the other layers one budget:
807 exact when they keep as many micro-batches in flight, at one chunk per
808 stage, and with the largest share of theirs otherwise, which fits.
809 """
810 ends = _ends(stage_layers)
811 middle = [layer for layer in stage_layers if layer not in ends]
812 share = max((_share(layer.in_flight) for layer in middle), default=0.0)
813 groups = _groups(middle)
814 within = least_times(groups, fronts, budget, bucket)
815 fastest, best, best_budget = math.inf, None, 0.0
816 for picks in itertools.product(*(fronts[end.key] for end in ends)):
817 chosen = {end.index: option for end, option in zip(ends, picks)}
818 limit = _middle_budget(ends, chosen, (budget, room), share, fronts, own)
819 time = within(limit) + sum(_time(option) for option in picks)
820 if time < fastest:
821 fastest, best, best_budget = time, chosen, limit
822 if best is None:
823 return None
824 choice = choose(groups, fronts, best_budget, bucket)
825 if choice is None:
826 return None
827 best.update(_assign(middle, [choice]))
828 return best
831def _choose_stage(
832 stage_layers: Sequence[_Layer],
833 stage: Stage,
834 peaks: _Peaks,
835 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
836 own: Mapping[LayerOption, LayerOption],
837 bucket: float,
838 capacity: float,
839) -> Optional[Dict[int, LayerOption]]:
840 """The fastest options of a stage's layers that keep it within *capacity* bytes wherever its peak comes."""
841 if peaks.backward is not None and stage_layers:
842 room = capacity - peaks.backward
843 if _backward_can_bind(stage_layers, room, fronts, own):
844 return _two_peaks(stage_layers, stage.budget, room, fronts, own, bucket)
845 choice = choose(stage.groups, fronts, stage.budget, bucket)
846 return None if choice is None else _assign(stage_layers, [choice])
849def _lightest_working(stage_layers: Sequence[_Layer], fronts: Mapping[Hashable, Tuple[LayerOption, ...]]) -> float:
850 """The least working set the layer that ends a stage's warm-up can end it on; 0 without one."""
851 ending = next((layer for layer in stage_layers if _ENDS_WARM_UP in layer.key), None)
852 if ending is None:
853 return 0.0
854 options = fronts[_kind_key(ending.key)]
855 return min(_working(option, options) for option in options)
858def _offload_budgets(
859 stage_layers: Sequence[_Layer],
860 stage: Stage,
861 peaks: _Peaks,
862 offloaded: Sequence[LayerOption],
863 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
864 own: Mapping[LayerOption, LayerOption],
865 sizes: Tuple[float, bool],
866) -> List[float]:
867 """For every count of a stage's first layers offloaded, what the other layers may keep as warm-up ends.
869 Less what the offloaded layers keep on the device and, held in transit,
870 the largest of their offloaded bytes. Where the stage's peak can also
871 come as a later micro-batch's backward ends, the other layers keep room
872 for it as though the layer that ends warm-up ran its lightest option
873 there, which fits.
875 Args:
876 offloaded: The options the first layers offload with, as many as
877 may offload.
878 sizes: ``(capacity, in_transit)``: the device's memory in bytes, and
879 whether the offloaded bytes are held in transit.
881 Returns:
882 A budget for every count, from none to every layer, negative for a
883 count beyond those that may offload.
884 """
885 capacity, in_transit = sizes
886 room = math.inf if peaks.backward is None else capacity - peaks.backward
887 share = _share(stage_layers[0].in_flight)
888 lightest = _lightest_working(stage_layers, fronts)
889 first_working = _first_working(stage_layers[0], offloaded[0], fronts, own)
890 budgets, held, largest = [stage.budget], 0.0, 0.0
891 for layer, option in zip(stage_layers, offloaded):
892 held += _kept(option, layer.in_flight)
893 largest = max(largest, option.link_bandwidth if in_transit else 0.0)
894 left_back = room - held * share - first_working - largest
895 rest_back = left_back / share + lightest if share else (math.inf if left_back >= 0 else -1.0)
896 budgets.append(min(stage.budget - held - largest, rest_back))
897 return budgets + [-1.0] * (len(stage_layers) + 1 - len(budgets))
900def _returns_fit(stage_layers: Sequence[_Layer], chosen: Mapping[int, LayerOption], count: int) -> bool:
901 """Whether each of a stage's *count* offloaded layers' bytes come back in room the backward has released.
903 A layer's bytes come back while the next layer's backward runs, beside
904 the next layer's own when it offloads too, once this micro-batch's
905 backward has released its share of what the layers after the next one
906 keep, *chosen* their options.
907 """
908 released = list(itertools.accumulate(
909 (_kept(chosen[layer.index], layer.in_flight) / layer.in_flight for layer in reversed(stage_layers)),
910 initial=0.0))[::-1]
911 for position in range(count):
912 back = sum(chosen[layer.index].link_bandwidth for layer in stage_layers[position:position + 2])
913 if back > released[min(position + 2, len(stage_layers))]:
914 return False
915 return True
918def _fastest_count(times: Sequence[float], rest: Sequence[float], most: int, fastest: float) -> int:
919 """How many of a stage's first layers to offload: the count, up to *most*, that runs fastest below *fastest*.
921 Args:
922 times: The time the first layers take offloaded, by their count.
923 rest: The time the other layers take in what each count leaves.
924 most: The most first layers the link reaches.
925 fastest: The time to beat, keeping every layer.
927 Returns:
928 The count, 0 where none beats *fastest*.
929 """
930 count, best = 0, fastest
931 for first in range(1, most + 1):
932 if times[first] + rest[first] < best:
933 count, best = first, times[first] + rest[first]
934 return count
937def _offload_choice(stage_layers: Sequence[_Layer], offloaded: Sequence[LayerOption], choice: Any,
938 in_transit: bool) -> Tuple[Dict[int, LayerOption], float]:
939 """Each layer's option, and what the offloaded layers hold in transit.
941 The stage's first layers offload with *offloaded*, one option each, and
942 the others run *choice*'s options. In transit, they hold the largest of
943 their offloaded bytes, where *in_transit*.
944 """
945 chosen = {layer.index: option for layer, option in zip(stage_layers, offloaded)}
946 chosen.update(_assign(stage_layers[len(offloaded):], [choice]))
947 transit = max(option.link_bandwidth for option in offloaded) if in_transit else 0.0
948 return chosen, transit
951def _offload_stage(
952 stage_layers: Sequence[_Layer],
953 stage: Stage,
954 peaks: _Peaks,
955 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
956 own: Mapping[LayerOption, LayerOption],
957 link: _Link,
958 bucket: float,
959 capacity: float,
960) -> Optional[Tuple[Dict[int, LayerOption], float]]:
961 """The fastest choice of a stage's layers when its first layers may offload.
963 Every count of first layers the link reaches is weighed against keeping
964 them, the rest of the layers choosing their options in what the
965 offloaded ones leave (:func:`_offload_budgets`): first with nothing in
966 transit, which holds where the backward releases room for every copy
967 back (:func:`_returns_fit`), else with the largest offloaded bytes held.
969 Returns:
970 ``(chosen, transit)``: each layer's option by index, and the memory
971 the offloaded layers keep in transit; ``None`` when the stage cannot
972 fit.
973 """
974 stay = _choose_stage(stage_layers, stage, peaks, fronts, own, bucket, capacity)
975 kept = None if stay is None else (stay, 0.0)
976 plain = [_plain(fronts[layer.key]) for layer in stage_layers]
977 shares = _offload_shares(stage_layers, plain, link)
978 most = next((position for position, share in enumerate(shares) if share <= 0), len(shares))
979 if not most:
980 return kept
981 offloaded = [_offloaded(option, share, link.cost_per_byte) for option, share in zip(plain, shares[:most])]
982 times = list(itertools.accumulate((_time(option) for option in offloaded), initial=0.0))
983 # Offload only where it saves time, more than rounding.
984 fastest = math.inf if stay is None else sum(_time(option) for option in stay.values()) * (1 - _ROUNDING)
985 for in_transit in (False, True):
986 budgets = _offload_budgets(stage_layers, stage, peaks, offloaded, fronts, own, (capacity, in_transit))
987 rest = suffix_times([Layers(layer.key, 1, layer.in_flight) for layer in stage_layers], fronts, budgets,
988 bucket)
989 count = _fastest_count(times, rest, most, fastest)
990 if not count:
991 return kept
992 choice = choose(_groups(stage_layers[count:]), fronts, budgets[count], bucket)
993 if choice is None:
994 return None
995 chosen, transit = _offload_choice(stage_layers, offloaded[:count], choice, in_transit)
996 if in_transit or _returns_fit(stage_layers, chosen, count):
997 held = _Peaks(peaks.warm_up + transit, None if peaks.backward is None else peaks.backward + transit)
998 return _trim_offload(stage_layers[:count], chosen, plain, held, (stage_layers, fronts, own),
999 (capacity, link.cost_per_byte)), transit
1000 return kept
1003def _trim_offload(
1004 offloading: Sequence[_Layer],
1005 chosen: Dict[int, LayerOption],
1006 plain: Sequence[LayerOption],
1007 held: _Peaks,
1008 stage: Tuple[Sequence[_Layer], Mapping[Hashable, Tuple[LayerOption, ...]], Mapping[LayerOption, LayerOption]],
1009 sizes: Tuple[float, float],
1010) -> Dict[int, LayerOption]:
1011 """*chosen* with no more offloaded than the stage needs, where the copies cost time.
1013 What the stage still has under the device's memory once its layers run
1014 their choice would be offloaded for nothing but the copies' cost, so the
1015 last offloaded layers keep it instead, a layer that keeps all it holds
1016 running plain.
1018 Args:
1019 offloading: The stage's layers that offload, in order.
1020 chosen: Each layer's option by index, those layers' offloading.
1021 plain: Each of the stage's layers' plain option, in order.
1022 held: What the stage keeps whatever its layers run, transit included.
1023 stage: ``(stage_layers, fronts, own)``, as for :func:`_stage_memory`.
1024 sizes: ``(capacity, cost_per_byte)``: the device's memory in bytes,
1025 and the step time one byte moved costs.
1027 Returns:
1028 The trimmed choice; *chosen* itself where the copies cost nothing.
1029 """
1030 capacity, cost_per_byte = sizes
1031 if cost_per_byte <= 0:
1032 return chosen
1033 stage_layers, fronts, own = stage
1034 # Trimmed to the byte, a stage would fill a budget of a fraction of a megabyte to that fraction, and the
1035 # search's result, reported in whole megabytes, would read over the budget its own fit check accepted.
1036 capacity = math.floor(capacity / MEGABYTE) * MEGABYTE
1037 slack = capacity - _stage_memory(stage_layers, held, chosen, fronts, own)
1038 chosen = dict(chosen)
1039 for position in range(len(offloading) - 1, -1, -1):
1040 if slack <= 0:
1041 break
1042 layer = offloading[position]
1043 option = chosen[layer.index]
1044 give = min(option.link_bandwidth, slack / layer.in_flight)
1045 share = option.link_bandwidth - give
1046 chosen[layer.index] = _offloaded(plain[position], share, cost_per_byte) if share > 0 else plain[position]
1047 slack -= give * layer.in_flight
1048 return chosen
1051def _stages(
1052 evaluator: EvaluatorV2,
1053 layers: Sequence[Sequence[_Layer]],
1054 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
1055 own: Mapping[LayerOption, LayerOption],
1056) -> Tuple[List[Stage], List[_Peaks]]:
1057 """Every stage for the knapsack, and what each keeps whatever its layers run, at each point its peak can come."""
1058 capacity = evaluator.ccfg.device_capacity.to_mb().size
1059 insights = evaluator.estimate_peak_insight()
1060 stages, peaks = [], []
1061 for stage_layers, insight, (warm_up, backward) in zip(layers, insights, evaluator.peak_points):
1062 stage, fixed = _stage(stage_layers, insight["Static"] + warm_up, capacity)
1063 stages.append(stage)
1064 second = None
1065 if backward:
1066 configured = {layer.index: layer.own for layer in stage_layers}
1067 second = (insight["Static"] + backward) * MEGABYTE - _backward_kept(stage_layers, configured, fronts, own)
1068 peaks.append(_Peaks(fixed, second))
1069 return stages, peaks
1072def choose_recompute(
1073 evaluator: EvaluatorV2,
1074 device_type: Any,
1075 ccfg: Optional[CustomConfig] = None,
1076 bucket: float = MEGABYTE,
1077 modes: Optional[Sequence[str]] = None,
1078 link: Optional[HostLink] = None,
1079 selective: Optional[Mapping[str, int]] = None,
1080 per_layer: bool = False,
1081) -> Optional[RecomputeChoice]:
1082 """The fastest recompute option of every layer that fits, at the evaluator's current strategy.
1084 Args:
1085 evaluator: The memory evaluator, set to the candidate's strategy.
1086 device_type: The device the times are priced on.
1087 ccfg: Estimator options; the search's defaults when omitted.
1088 bucket: The knapsack's memory granularity, in bytes.
1089 modes: For a runtime that runs every layer one way, the modes it
1090 runs, of :data:`MODES`: the fastest that fits is chosen for every
1091 layer. Without them, each layer gets its own option from its
1092 kind's front.
1093 link: For a choice per layer at one chunk per stage, the host link
1094 each stage's first layers may offload all or part of what they
1095 keep over; ``None`` keeps every layer's activations on the
1096 device.
1097 selective: With *modes*, the switches the runtime's own selective
1098 mode sets, 1 to keep an op and 0 to recompute it, where they are
1099 not the config's, as HyperParallel's trainer runs its policy
1100 whatever its train yaml's mode. The mode is offered where every
1101 kind's profile measures that setting whole, as a census prices
1102 it; the formulas do not price the matmuls the policy recomputes.
1103 per_layer: With *modes*, give each layer the fastest of them that
1104 fits rather than one for every layer, for a runtime that runs a
1105 mode per layer, as HyperParallel's trainer does with
1106 ``activation_checkpoint.layers``. Each range then names its mode
1107 (:func:`trainer_plan`). With *link*, and off among *modes*, each
1108 stage's first layers may also run off and offload all or part of
1109 what they keep, reporting the mode off: priced as offload should
1110 run, for a runtime that offloads them, which the trainer does
1111 not yet.
1113 Returns:
1114 The choice, or ``None`` when there is none to make: a multimodal
1115 model, a schedule other than those in :data:`SCHEDULES`, a layer
1116 that runs an option its kind cannot offer, or a stage that does not
1117 fit even with the lightest option or mode.
1118 """
1119 config = evaluator.ccfg
1120 if config.multimodal or config.pp_sched not in SCHEDULES:
1121 return None
1122 counts = micro_batches_in_flight(evaluator)
1123 profiles = layer_profiles(evaluator, device_type, ccfg, most_in_flight=max(max(row) for row in counts),
1124 each_switch=modes is None or ("selective" in modes and selective is None),
1125 in_flight=[count for row in counts for count in row])
1126 modes = _priced_modes(modes, profiles, selective)
1127 fronts, by_mode = _offered(profiles, configured_switches(evaluator), modes, selective)
1128 found = _body_layers(evaluator, fronts, counts)
1129 if found is None:
1130 return None
1131 layers, fronts, by_mode, own = _charge_working_sets(*found, fronts, by_mode)
1132 stages, peaks = _stages(evaluator, layers, fronts, own)
1133 capacity = config.device_capacity.to_mb().size * MEGABYTE
1134 if by_mode is not None:
1135 if per_layer:
1136 option_modes = _option_modes(by_mode)
1137 if link is not None and config.vp == 1 and "off" in modes:
1138 return _offload_result(layers, stages, peaks, fronts, own, _link(link, evaluator), bucket, capacity,
1139 option_modes)
1140 return _each_layer(layers, stages, peaks, fronts, own, (bucket, capacity), option_modes)
1141 one = _one_mode(layers, peaks, by_mode, modes, fronts, own, capacity)
1142 return None if one is None else _result(layers, peaks, one[1], one[0], fronts, own)
1143 if link is not None and config.vp == 1:
1144 return _offload_result(layers, stages, peaks, fronts, own, _link(link, evaluator), bucket, capacity)
1145 return _each_layer(layers, stages, peaks, fronts, own, (bucket, capacity))
1148class RecomputeAnalysis:
1149 """The plans one more recomputed layer reaches at a fixed strategy, with the layer prices kept.
1151 The layer options, the fixed stage memory and the peak constraints are
1152 priced once for the strategy, as :func:`choose_recompute` prices them,
1153 and every alternative reuses them. The alternatives are returned
1154 unscored: what each would cost the step is the caller's estimate to
1155 make, and it is the cost model's estimate, not a measured figure.
1156 """
1158 def __init__(
1159 self,
1160 evaluator: EvaluatorV2,
1161 device_type: Any,
1162 modes: Sequence[str],
1163 selective: Optional[Mapping[str, int]] = None,
1164 dimensions: Tuple[Tuple[str, int], ...] = (),
1165 ) -> None:
1166 """Price the layer options and the fixed stage memory once for one strategy.
1168 Args:
1169 evaluator: The memory evaluator, set to the fixed strategy.
1170 device_type: The device the times are priced on.
1171 modes: The named recompute modes a layer may run, which must
1172 include full for a layer to be turned up to it.
1173 selective: The switches the runtime's own selective mode sets,
1174 as :func:`choose_recompute` takes them.
1175 dimensions: The strategy's parallel degrees, which every
1176 alternative reports.
1178 Raises:
1179 ValueError: No modes, or a mode that is not one of :data:`MODES`.
1180 """
1181 if not modes or set(modes) - set(MODES):
1182 raise ValueError(f"recompute analysis requires modes from {', '.join(MODES)}")
1183 self._config = evaluator.ccfg
1184 self._device_type = device_type
1185 self._dimensions = dimensions
1186 self._layers: List[List[_Layer]] = []
1187 self._fronts: Dict[Hashable, Tuple[LayerOption, ...]] = {}
1188 self._by_mode: Dict[Hashable, Dict[str, LayerOption]] = {}
1189 self._own: Dict[LayerOption, LayerOption] = {}
1190 self._groups: List[List[_Layer]] = []
1191 self._peaks: List[_Peaks] = []
1192 self._capacity = self._config.device_capacity.to_mb().size
1193 if self._config.multimodal or self._config.pp_sched not in SCHEDULES or "full" not in modes:
1194 return
1195 self._price(evaluator, modes, selective)
1197 @property
1198 def config(self) -> Any:
1199 """The config the layer options were priced on, which a caller scores an alternative with."""
1200 return self._config
1202 @property
1203 def dimensions(self) -> Tuple[Tuple[str, int], ...]:
1204 """The parallel degrees the analysis is fixed at."""
1205 return self._dimensions
1207 def _price(
1208 self, evaluator: EvaluatorV2, modes: Sequence[str], selective: Optional[Mapping[str, int]]
1209 ) -> None:
1210 """Price every kind's options and each stage's fixed memory, then group the equivalent layers."""
1211 counts = micro_batches_in_flight(evaluator)
1212 profiles = layer_profiles(
1213 evaluator, self._device_type, most_in_flight=max(max(row) for row in counts),
1214 each_switch="selective" in modes and selective is None,
1215 in_flight=[count for row in counts for count in row],
1216 )
1217 modes = _priced_modes(modes, profiles, selective)
1218 fronts, by_mode = _offered(profiles, configured_switches(evaluator), modes, selective)
1219 found = _body_layers(evaluator, fronts, counts)
1220 if found is None:
1221 return
1222 self._layers, self._fronts, self._by_mode, self._own = _charge_working_sets(*found, fronts, by_mode)
1223 _, self._peaks = _stages(evaluator, self._layers, self._fronts, self._own)
1224 self._groups = _equivalent_groups(self._layers, self._peaks)
1226 def alternatives(self, current: RecomputeChoice) -> List[RecomputeChoice]:
1227 """The distinct plans that fully recompute one more layer of *current* and still fit.
1229 For each group of equivalent layers and each mode its layers start
1230 in, one layer is turned up to full and the group's new counts are
1231 assigned to its layers as the normal optimizer assigns them. Index
1232 swaps within a group give one plan; different kinds, stages,
1233 in-flight counts and peak boundaries stay distinct. Every plan is
1234 weighed against the strategy's own device capacity.
1236 Args:
1237 current: A plan of a named mode per layer, every body layer in it.
1239 Returns:
1240 The fitting plans, without scores, each reporting the analysis's
1241 degrees; empty where no group can be turned up, where this
1242 strategy has no plan to extend, or where *current* runs one mode
1243 for every layer.
1245 Raises:
1246 ValueError: *current* was priced at other degrees, leaves a body
1247 layer out or names one twice, runs a mode this analysis does
1248 not price, or offloads activations.
1249 """
1250 if current.dimensions and current.dimensions != self._dimensions:
1251 raise ValueError("recompute alternatives must keep the same parallel dimensions")
1252 if not self._layers or current.mode is not None:
1253 return []
1254 defined = self._defined(current)
1255 chosen = self._chosen(defined)
1256 option_modes = _option_modes(self._by_mode)
1257 results: List[RecomputeChoice] = []
1258 seen: Dict[Tuple[FrozenSet[Any], ...], List[RecomputeChoice]] = {}
1259 for group in self._groups:
1260 for result, counts in self._turned_up(group, defined, chosen, option_modes):
1261 if any(_equivalent_cost(result, previous) for previous in seen.get(counts, ())):
1262 continue
1263 seen.setdefault(counts, []).append(result)
1264 results.append(replace(result, dimensions=self._dimensions))
1265 return results
1267 def _defined(self, current: RecomputeChoice) -> Dict[int, LayerRange]:
1268 """Each body layer's range in *current*, refusing a plan that does not cover them once each."""
1269 defined: Dict[int, LayerRange] = {}
1270 for item in current.ranges:
1271 if item.count <= 0 or item.first < 0 or item.option.link_bandwidth:
1272 raise ValueError("recompute analysis requires positive layer ranges without activation offload")
1273 for index in range(item.first, item.first + item.count):
1274 if index in defined:
1275 raise ValueError("recompute layer ranges must not overlap")
1276 defined[index] = item
1277 every = {layer.index for stage in self._layers for layer in stage}
1278 if set(defined) != every:
1279 raise ValueError("the recompute plan must define every body layer exactly once")
1280 return defined
1282 def _chosen(self, defined: Mapping[int, LayerRange]) -> Dict[int, LayerOption]:
1283 """Each body layer's option under *defined*, refusing a kind or a mode this analysis does not price."""
1284 chosen = {}
1285 for stage in self._layers:
1286 for layer in stage:
1287 item = defined[layer.index]
1288 if item.kind != layer.key[1] or item.mode not in self._by_mode[layer.key]:
1289 raise ValueError("the recompute plan must use the fixed model's kinds and allowed modes")
1290 chosen[layer.index] = self._by_mode[layer.key][item.mode]
1291 return chosen
1293 def _turned_up(
1294 self,
1295 group: Sequence[_Layer],
1296 defined: Mapping[int, LayerRange],
1297 chosen: Mapping[int, LayerOption],
1298 option_modes: Mapping[LayerOption, str],
1299 ) -> Iterator[Tuple[RecomputeChoice, Tuple[FrozenSet[Any], ...]]]:
1300 """Yield each fitting plan that turns one of *group*'s layers up to full, with its mode counts."""
1301 started: set = set()
1302 for layer in group:
1303 mode = defined[layer.index].mode
1304 if mode == "full" or mode in started:
1305 continue
1306 started.add(mode)
1307 changed = dict(chosen)
1308 changed[layer.index] = self._by_mode[layer.key]["full"]
1309 changed.update(_assign_counted(group, changed))
1310 result = _result(self._layers, self._peaks, changed, None, self._fronts, self._own, option_modes)
1311 if result.memory > self._capacity:
1312 continue
1313 yield result, self._mode_counts(changed, option_modes)
1315 def _mode_counts(
1316 self, chosen: Mapping[int, LayerOption], option_modes: Mapping[LayerOption, str]
1317 ) -> Tuple[FrozenSet[Any], ...]:
1318 """Each stage's count of kinds and modes, the key plans of the same shape share."""
1319 return tuple(
1320 frozenset(Counter((layer.key[1], option_modes[chosen[layer.index]]) for layer in stage).items())
1321 for stage in self._layers
1322 )
1325def _equivalent_groups(layers: Sequence[Sequence[_Layer]], peaks: Sequence[_Peaks]) -> List[List[_Layer]]:
1326 """The layers that are interchangeable: one kind, stage and in-flight count, either side of a peak boundary.
1328 A stage whose peak can also come as a later backward ends keeps its
1329 first layer apart, since that layer's option is weighed at that peak too.
1330 """
1331 groups: List[List[_Layer]] = []
1332 for stage_layers, stage_peaks in zip(layers, peaks):
1333 grouped: Dict[Tuple[Hashable, int, bool], List[_Layer]] = {}
1334 for layer in stage_layers:
1335 at_boundary = stage_peaks.backward is not None and layer.index == stage_layers[0].index
1336 grouped.setdefault((layer.key, layer.in_flight, at_boundary), []).append(layer)
1337 groups.extend(grouped.values())
1338 return groups
1341def _option_modes(by_mode: Mapping[Hashable, Mapping[str, LayerOption]]) -> Dict[LayerOption, str]:
1342 """Each option of every kind's modes, by the mode it runs."""
1343 return {option: mode for options in by_mode.values() for mode, option in options.items()}
1346def _each_layer(
1347 layers: Sequence[Sequence[_Layer]],
1348 stages: Sequence[Stage],
1349 peaks: Sequence[_Peaks],
1350 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
1351 own: Mapping[LayerOption, LayerOption],
1352 sizes: Tuple[float, float],
1353 modes: Optional[Mapping[LayerOption, str]] = None,
1354) -> Optional[RecomputeChoice]:
1355 """Each layer's fastest option of its kind's that fits, stage by stage.
1357 Args:
1358 sizes: ``(bucket, capacity)``, the knapsack's granularity and the
1359 device's memory, in bytes.
1360 modes: For a choice of a mode per layer, each option's mode.
1362 Returns:
1363 The choice, or ``None`` when some stage does not fit.
1364 """
1365 bucket, capacity = sizes
1366 chosen: Dict[int, LayerOption] = {}
1367 for stage_layers, stage, stage_peaks in zip(layers, stages, peaks):
1368 stage_chosen = _choose_stage(stage_layers, stage, stage_peaks, fronts, own, bucket, capacity)
1369 if stage_chosen is None:
1370 return None
1371 chosen.update(stage_chosen)
1372 return _result(layers, peaks, chosen, None, fronts, own, modes)
1375def _offload_result(
1376 layers: Sequence[Sequence[_Layer]],
1377 stages: Sequence[Stage],
1378 peaks: Sequence[_Peaks],
1379 fronts: Mapping[Hashable, Tuple[LayerOption, ...]],
1380 own: Mapping[LayerOption, LayerOption],
1381 link: _Link,
1382 bucket: float,
1383 capacity: float,
1384 modes: Optional[Mapping[LayerOption, str]] = None,
1385) -> Optional[RecomputeChoice]:
1386 """The choice per layer, each stage's first layers offloading where that is faster.
1388 Args:
1389 modes: For a choice of a mode per layer, each option's mode. An
1390 offloaded layer runs its kind's plain option, and reports that
1391 option's mode.
1392 """
1393 chosen: Dict[int, LayerOption] = {}
1394 held = []
1395 for stage_layers, stage, stage_peaks in zip(layers, stages, peaks):
1396 found = _offload_stage(stage_layers, stage, stage_peaks, fronts, own, link, bucket, capacity)
1397 if found is None:
1398 return None
1399 chosen.update(found[0])
1400 transit = found[1]
1401 held.append(_Peaks(stage_peaks.warm_up + transit,
1402 None if stage_peaks.backward is None else stage_peaks.backward + transit))
1403 if modes is not None:
1404 modes = dict(modes)
1405 for layer in (layer for stage_layers in layers for layer in stage_layers):
1406 if chosen[layer.index].link_bandwidth:
1407 plain = _plain(fronts[layer.key])
1408 modes[chosen[layer.index]] = modes[own.get(plain, plain)]
1409 return _result(layers, held, chosen, None, fronts, own, modes)
1412def offloaded_share(option: LayerOption) -> float:
1413 """The share of what a layer running *option* keeps per micro-batch that it moves to the host; 0 without offload."""
1414 if not option.link_bandwidth:
1415 return 0.0
1416 return option.link_bandwidth / (option.link_bandwidth + max(0.0, option.memory_per_micro_batch))
1419def option_label(option: LayerOption) -> str:
1420 """How an option reads in the search's output."""
1421 if option.recompute is None:
1422 label = "full recompute"
1423 elif not option.recompute:
1424 label = "no recompute"
1425 else:
1426 label = "recompute " + "+".join(name for name in SWITCHES if name in option.recompute)
1427 if not option.link_bandwidth:
1428 return label
1429 return f"{label}, {offloaded_share(option):.0%} offloaded to the host"
1432def option_record(option: LayerOption) -> Any:
1433 """An option as plain data: ``"none"``, ``"full"``, or the switches it recomputes, in switch order."""
1434 if option.recompute is None:
1435 return "full"
1436 if not option.recompute:
1437 return "none"
1438 return [name for name in SWITCHES if name in option.recompute]
1441def to_records(choice: RecomputeChoice) -> List[Dict[str, Any]]:
1442 """The choice's ranges as plain data, for a result file.
1444 Returns:
1445 One ``{"first", "count", "kind", "recompute"}`` per range, in model
1446 order, with ``"offload": True`` and ``"offloaded_bytes"``, what each
1447 of its layers moves per micro-batch, for a range that offloads, and
1448 ``"mode"`` for one a choice of a mode per layer gave its mode;
1449 ``kind`` is the kind's name, or ``None``.
1450 """
1451 records = []
1452 for item in choice.ranges:
1453 record = {
1454 "first": item.first,
1455 "count": item.count,
1456 "kind": item.kind.name if item.kind is not None else None,
1457 "recompute": option_record(item.option),
1458 }
1459 if item.option.link_bandwidth:
1460 record["offload"] = True
1461 record["offloaded_bytes"] = item.option.link_bandwidth
1462 if item.mode is not None:
1463 record["mode"] = item.mode
1464 records.append(record)
1465 return records
1468def trainer_plan(choice: RecomputeChoice) -> Tuple[str, Dict[str, str]]:
1469 """The choice as HyperParallel's trainer states it: ``activation_checkpoint.mode`` and ``.layers``.
1471 The trainer runs every layer its mode but those its plan names, by index
1472 or by an inclusive range ``"first-last"``. The mode is the one most
1473 layers run of those that recompute, full on a tie, or off where none
1474 recomputes; the plan names the other layers, consecutive layers of one
1475 mode in one entry whatever their kinds.
1477 The mode is a mode that recomputes, rather than off with the recomputed
1478 layers named, because the trainer accepts an entry equal to the mode: a
1479 runner that states the mode itself, as a sweep forcing its own full
1480 default does, then still runs the layers this plan states off. Named the
1481 other way round, the same run would recompute every layer and invert the
1482 choice without a word.
1484 The indices are the choice's, which count the body layers in model order,
1485 MTP's after the decoder's: the indices of the trainer's blocks, for a
1486 model without MTP layers.
1488 Args:
1489 choice: A choice of one mode for every layer, or of a mode per layer
1490 (:func:`choose_recompute` with *per_layer*).
1492 Returns:
1493 ``(mode, layers)``, *layers* empty when every layer runs *mode*.
1495 Raises:
1496 ValueError: For a choice of options no mode stands for, each layer's
1497 own from its kind's front.
1498 """
1499 if choice.mode is not None:
1500 return choice.mode, {}
1501 if any(item.mode is None for item in choice.ranges):
1502 raise ValueError("the trainer runs a mode per layer, and this choice gives some layers options that are not "
1503 "one; choose among its modes with per_layer")
1504 counts = Counter()
1505 for item in choice.ranges:
1506 counts[item.mode] += item.count
1507 recomputing = [mode for mode in ("full", "selective") if counts[mode]]
1508 if not recomputing:
1509 return "off", {}
1510 mode = max(recomputing, key=lambda name: counts[name])
1511 runs: List[List[Any]] = []
1512 for item in sorted(choice.ranges, key=lambda item: item.first):
1513 if runs and runs[-1][2] == item.mode and runs[-1][1] + 1 == item.first:
1514 runs[-1][1] = item.first + item.count - 1
1515 else:
1516 runs.append([item.first, item.first + item.count - 1, item.mode])
1517 return mode, {
1518 (str(first) if first == last else f"{first}-{last}"): run_mode
1519 for first, last, run_mode in runs
1520 if run_mode != mode
1521 }
1524# Where the installed trainer configures activation checkpointing: the module,
1525# and the dataclass whose fields are the keys ``activation_checkpoint`` takes.
1526_TRAINER_CHECKPOINT_CONFIG = ("hyper_parallel.trainer.config.parallelism", "ActivationCheckpointConfig")
1529def trainer_runs_layer_plans() -> bool:
1530 """Whether the installed trainer runs a recompute mode per layer: its ``activation_checkpoint`` takes ``layers``.
1532 A trainer whose ``ActivationCheckpointConfig`` declares no ``layers``
1533 field refuses a yaml that states one, so a plan per layer
1534 (:func:`trainer_plan`) is then advice the trainer cannot run: it runs
1535 every layer the plan's mode. The trainer's configuration loads only
1536 when asked, as importing it costs about as much as the cost model.
1538 Returns:
1539 True where the dataclass declares ``layers``; False where it does
1540 not, or where this tree has no such trainer.
1541 """
1542 module, name = _TRAINER_CHECKPOINT_CONFIG
1543 try:
1544 config = getattr(importlib.import_module(module), name)
1545 except (ImportError, AttributeError):
1546 return False
1547 return is_dataclass(config) and any(field.name == "layers" for field in fields(config))
1550def describe(choice: RecomputeChoice) -> str:
1551 """The choice as one line per range of layers, after its mode when one was chosen for every layer.
1553 A range a choice of a mode per layer gave its mode reads as the mode.
1554 """
1555 lines = [] if choice.mode is None else [f"every layer: {choice.mode}"]
1556 for item in choice.ranges:
1557 last = item.first + item.count - 1
1558 layers = f"layer {item.first}" if item.count == 1 else f"layers {item.first}-{last}"
1559 kind = f" ({item.kind.name})" if item.kind is not None else ""
1560 lines.append(f"{layers}{kind}: {option_label(item.option) if item.mode is None else item.mode}")
1561 return "\n".join(lines)