Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / recompute / profile.py: 100%
44 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"""What each recompute switch of a layer kind saves and costs, measured alone.
17A selective layer's memory and backward time add up over its switches. Each
18switch gates its own activation term in the memory model and its own op in
19the recompute estimate, and only ``gather`` acts on the communication
20buffers, both on the memory kept and on the volume sent again. So a layer that
21recomputes a set of ops costs the plain layer plus what each of those ops
22costs alone, and a :class:`SwitchProfile`, nine measurements, prices all 128
23settings of the seven switches; an MLA layer's, a tenth, all 256 with the
24switch of its up-projections.
26A census of a layer kind (IR phase 5) prices its plain layer, and a layer
27running HyperParallel's selective policy, by what it measured, and every
28other setting by what its records per op state for the ops the setting
29keeps, or by the formulas where it states none. The settings then add up
30from the layer selective with every op kept, and the policy's is measured
31whole. Where the census prices them per op, a setting's working set as
32warm-up ends is clamped where the layer keeps more than its backward holds,
33so it adds up before the clamp, which :meth:`SwitchProfile.selective`
34applies to each setting.
35"""
36from dataclasses import dataclass, field, replace
37from typing import FrozenSet, Iterable, Mapping, Optional, Tuple
39# The recompute switches; 1 keeps an op's activation, 0 recomputes the op.
40SWITCHES = ("attBMM", "headCast", "dropout", "softmax", "normOp", "gather", "ffAct", "attUp")
41# The switches a layer kind may keep nothing under, an MLA layer's
42# up-projections: a front weighs them only where they change what the kind
43# costs, so that every other kind's settings stay the seven switches'.
44OPTIONAL = ("attUp",)
47@dataclass(frozen=True)
48class Cost:
49 """What running a layer one way costs.
51 Attributes:
52 memory_per_micro_batch: Bytes kept for each micro-batch in flight: the
53 activations, and the communication buffers that grow with the
54 micro-batches in flight.
55 memory_once: Bytes of communication buffers kept once, however many
56 micro-batches are in flight.
57 backward_time: The backward time, recompute included.
58 excess: At each of its profile's :attr:`~SwitchProfile.counts`, the
59 bytes the two memories charge beyond what the layer keeps with
60 that many micro-batches in flight. The split is exact at one
61 micro-batch and at the most any stage keeps, and a buffer that
62 does not grow can hide one that does in between.
63 working: What the working set of the layer's backward holds beyond
64 what the layer keeps at one micro-batch, as the memory model
65 charges it to the layer that ends warm-up: under FSDP that
66 reshards, two layers' gathered parameters, its own and the
67 next's, in the buffers the gathers take, and where a census
68 prices the layer, what its backward holds beyond what the stage
69 keeps for it already. For full recompute, beyond what the plain
70 layer keeps: its backward runs the plain layer, on activations
71 the stage does not keep.
72 first_working: The same for the backward a stage runs last, its
73 first layer's, as a micro-batch's backward ends: one layer's
74 gathered parameters, with none left to prefetch, and all a
75 census says the backward holds, as the stage keeps no
76 activations of that micro-batch for it.
77 census_kept: Where a census's records per op price the layer, what
78 it keeps of them at one micro-batch; *working* is then stated
79 before the clamp (:attr:`SwitchProfile.census_held`).
80 """
82 memory_per_micro_batch: float
83 memory_once: float
84 backward_time: float
85 excess: Tuple[float, ...] = ()
86 working: float = 0.0
87 first_working: float = 0.0
88 census_kept: float = 0.0
90 def __add__(self, other: "Cost") -> "Cost":
91 """The sum, cost by cost."""
92 return Cost(
93 *(mine + theirs for mine, theirs in zip(self.values(), other.values())),
94 excess=tuple(mine + theirs for mine, theirs in zip(self.excess, other.excess)),
95 working=self.working + other.working,
96 first_working=self.first_working + other.first_working,
97 census_kept=self.census_kept + other.census_kept,
98 )
100 def __sub__(self, other: "Cost") -> "Cost":
101 """The difference, cost by cost."""
102 return Cost(
103 *(mine - theirs for mine, theirs in zip(self.values(), other.values())),
104 excess=tuple(mine - theirs for mine, theirs in zip(self.excess, other.excess)),
105 working=self.working - other.working,
106 first_working=self.first_working - other.first_working,
107 census_kept=self.census_kept - other.census_kept,
108 )
110 def values(self) -> Tuple[float, float, float]:
111 """``(memory per micro-batch, memory once, backward time)``."""
112 return self.memory_per_micro_batch, self.memory_once, self.backward_time
115@dataclass(frozen=True)
116class SwitchProfile:
117 """A layer kind's costs: plain, with each op alone recomputed, and fully recomputed.
119 Attributes:
120 forward_time: The layer's forward time, whatever it recomputes.
121 plain: The layer keeping every op.
122 alone: Per switch, the layer recomputing that op alone.
123 full: The layer fully recomputed.
124 counts: The counts of micro-batches in flight, between one and the
125 most any stage keeps, at which each cost states its excess.
126 selective_base: The layer selective with every op kept, which the
127 settings add up from where a census prices the plain layer and
128 not a selective one; the plain layer when omitted.
129 whole: The settings measured whole, which do not add up over their
130 switches: HyperParallel's selective policy, where a census
131 prices it.
132 census_held: Where a census's records per op price the selective
133 settings, what the layer's backward holds at one micro-batch:
134 as warm-up ends, it holds that beyond what the layer keeps, and
135 no less than nothing.
136 """
138 forward_time: float
139 plain: Cost
140 alone: Mapping[str, Cost]
141 full: Cost
142 counts: Tuple[int, ...] = ()
143 selective_base: Optional[Cost] = None
144 whole: Mapping[FrozenSet[str], Cost] = field(default_factory=dict)
145 census_held: Optional[float] = None
147 def acting(self) -> Tuple[str, ...]:
148 """The switches a setting is made of: every one but an optional one that changes nothing."""
149 base = self.plain if self.selective_base is None else self.selective_base
150 return tuple(name for name in SWITCHES
151 if name not in OPTIONAL or (name in self.alone and self.alone[name] != base))
153 def selective(self, recompute: Iterable[str]) -> Cost:
154 """The cost of recomputing the ops *recompute* names.
156 The plain layer's for none, a setting's own where it is measured
157 whole, and otherwise the base's, plus what each op costs alone
158 beyond it; where a census's records per op price it, its working
159 set as warm-up ends clamped as the memory model clamps it.
160 """
161 chosen = frozenset(recompute)
162 if chosen in self.whole:
163 return self.whole[chosen]
164 base = self.plain if self.selective_base is None or not chosen else self.selective_base
165 cost = base
166 for name in SWITCHES:
167 if name in chosen:
168 cost = cost + (self.alone[name] - base)
169 if chosen and self.census_held is not None:
170 cost = replace(cost, working=cost.working + max(0.0, cost.census_kept - self.census_held))
171 return cost