« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-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"""Body module"""
16from __future__ import annotations
17from typing import TYPE_CHECKING, NamedTuple
18from hyper_parallel.auto_parallel._model_spec import KindActivations
19from hyper_parallel.auto_parallel._op_records import load_op_records
20from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
21from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.utils import EvalUtils
22from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.comm import EvalLayerComm
23from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.layer_block import EvalAttn
24from hyper_parallel.auto_parallel.sapp_nd.nd.common.cp_types import (
25 CPMemoryBreakdown,
26 CPAlgo,
27 _resolve_cp_algo,
28)
29from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import (
30 detect_attention_type,
31 compute_kv_dim,
32)
33from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import runs_hyper_selective
34from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
36if TYPE_CHECKING:
37 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
38 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
39 from typing import Optional, Tuple
41# Bytes per element of the attention tensors that CP shards.
42_CP_ATTENTION_SCORES_BYTES = 4
43_CP_SOFTMAX_OUTPUTS_BYTES = 4
44_CP_DROPOUT_MASK_BYTES = 1
45_CP_KV_BYTES = 2 * 2 # fp16 key and value
48class _CPAttentionTerms(NamedTuple):
49 """Attention memory of one layer under CP, and what CP saves."""
51 kv_cache_memory: float
52 attention_scores_memory: float
53 softmax_outputs_memory: float
54 dropout_mask_memory: float
55 s2_reduction: float
56 kv_reduction: float
59def _ulysses_cp_attention_terms(
60 s: float, b: float, cp: float, kv_dim: float, a_per_rank: float
61) -> _CPAttentionTerms:
62 """Ulysses CP: each rank holds all s tokens but a/(t*cp) heads."""
63 a_per_cp_rank = a_per_rank / cp
64 kv_dim_per_cp_rank = kv_dim / cp
65 s2_items_no_cp = (
66 (_CP_ATTENTION_SCORES_BYTES + _CP_SOFTMAX_OUTPUTS_BYTES + _CP_DROPOUT_MASK_BYTES)
67 * s * s * b * a_per_rank
68 )
69 s2_items_with_cp = (
70 (_CP_ATTENTION_SCORES_BYTES + _CP_SOFTMAX_OUTPUTS_BYTES + _CP_DROPOUT_MASK_BYTES)
71 * s * s * b * a_per_cp_rank
72 )
73 return _CPAttentionTerms(
74 kv_cache_memory=_CP_KV_BYTES * s * b * kv_dim_per_cp_rank,
75 attention_scores_memory=_CP_ATTENTION_SCORES_BYTES * s * s * b * a_per_cp_rank,
76 softmax_outputs_memory=_CP_SOFTMAX_OUTPUTS_BYTES * s * s * b * a_per_cp_rank,
77 dropout_mask_memory=_CP_DROPOUT_MASK_BYTES * s * s * b * a_per_cp_rank,
78 s2_reduction=s2_items_no_cp - s2_items_with_cp,
79 kv_reduction=_CP_KV_BYTES * s * b * kv_dim * ((cp - 1) / cp),
80 )
83def _ring_cp_attention_terms(
84 s: float, b: float, cp: float, kv_dim: float, a_per_rank: float
85) -> _CPAttentionTerms:
86 """Ring CP: each rank holds s/cp tokens and a/t heads, KV is all-gathered."""
87 return _CPAttentionTerms(
88 kv_cache_memory=_CP_KV_BYTES * (s / cp) * b * kv_dim,
89 attention_scores_memory=_CP_ATTENTION_SCORES_BYTES * (s / cp) * s * b * a_per_rank,
90 softmax_outputs_memory=_CP_SOFTMAX_OUTPUTS_BYTES * (s / cp) * s * b * a_per_rank,
91 dropout_mask_memory=_CP_DROPOUT_MASK_BYTES * (s / cp) * s * b * a_per_rank,
92 s2_reduction=(
93 (_CP_ATTENTION_SCORES_BYTES + _CP_SOFTMAX_OUTPUTS_BYTES + _CP_DROPOUT_MASK_BYTES)
94 * s * s * b * a_per_rank * ((cp - 1) / cp)
95 ),
96 kv_reduction=_CP_KV_BYTES * s * b * kv_dim * ((cp - 1) / cp),
97 )
100class EvalBody:
101 """Body layer formulas class"""
103 @staticmethod
104 def num_params_layer(
105 ccfg: CostModelConfig, ctx: Context
106 ) -> Tuple[float, float, float]:
107 """Parameters count.
109 Returns a 3-tuple (non_exp, routed, shared):
110 - non_exp: attention + norm params (and dense FFN if n_exp==1,
111 the router otherwise)
112 - routed: routed expert params (0 if n_exp==1)
113 - shared: shared expert params (0 if n_shared_exp==0 or no pointer)
114 """
115 non_exp = ctx.attn_num_p(ccfg, ctx) + ctx.norm_num_p(ccfg, ctx)
116 routed = 0.0
117 shared = 0.0
118 if ccfg.n_exp == 1:
119 non_exp += ctx.ffn_num_p(ccfg, ctx)
120 else:
121 if ctx.ffn_router_num_p is not None:
122 non_exp += ctx.ffn_router_num_p(ccfg, ctx)
123 if ctx.ffn_routed_num_p is not None:
124 routed = ctx.ffn_routed_num_p(ccfg, ctx)
125 if ctx.ffn_shared_num_p is not None:
126 shared = ctx.ffn_shared_num_p(ccfg, ctx)
127 return (non_exp, routed, shared)
129 @staticmethod
130 def stat_p_layer(ccfg: CostModelConfig, ctx: Context) -> float:
131 """model param"""
132 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx)
133 # Routed experts: EP sharding
134 routed_mem = routed_p / ccfg.ep * ccfg.bytes_p / ccfg.shard_p_os_exp
135 # Shared experts: partial DP sharding
136 shared_mem = shared_p * ccfg.bytes_p / ccfg.shard_p_os_exp_partial
137 # Non expert
138 non_exp_mem = non_exp_p * ccfg.bytes_p / ccfg.shard_p_os_non_exp_partial
139 return non_exp_mem + routed_mem + shared_mem
141 @staticmethod
142 def stat_os_layer(ccfg: CostModelConfig, ctx: Context) -> float:
143 """optim state"""
144 if ctx.swap_os:
145 return 0
146 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx)
147 # Routed experts
148 routed_mem = routed_p / ccfg.ep * ccfg.bytes_optim / ccfg.shard_p_os_exp
149 # Shared experts
150 shared_mem = shared_p * ccfg.bytes_optim / ccfg.shard_p_os_exp_partial
151 # Non expert
152 non_exp_mem = non_exp_p * ccfg.bytes_optim / ccfg.shard_p_os_non_exp_partial
153 return non_exp_mem + routed_mem + shared_mem
155 @staticmethod
156 def stat_grad_layer(ccfg: CostModelConfig, ctx: Context) -> float:
157 """gradients"""
158 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx)
159 # Routed experts
160 routed_mem = routed_p / ccfg.ep * ccfg.bytes_grad / ccfg.shard_grad_exp
161 # Shared experts: use shard_grad_exp_partial (independent of os sharding)
162 shared_mem = shared_p * ccfg.bytes_grad / ccfg.shard_grad_exp_partial
163 # Non expert
164 non_exp_mem = non_exp_p * ccfg.bytes_grad / ccfg.shard_grad_non_exp
165 return non_exp_mem + routed_mem + shared_mem
167 @staticmethod
168 def reduced_grad_layer(ccfg: CostModelConfig, ctx: Context) -> Tuple[float, float]:
169 """The layer's gradients FSDP reduce-scatters, whole as its backward computes them and sharded."""
170 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx)
171 return EvalUtils.reduced_grads(ccfg, (
172 (non_exp_p, ccfg.t, ccfg.shard_grad_non_exp),
173 (routed_p / ccfg.ep, ccfg.t_exp, ccfg.shard_grad_exp),
174 (shared_p, ccfg.t_exp, ccfg.shard_grad_exp_partial),
175 ))
177 # No recompute and select recompute
179 @staticmethod
180 def layer_activ(ccfg: CostModelConfig, ctx: Context) -> float:
181 """activations"""
182 census = getattr(ccfg, "kind_activations", None)
183 if isinstance(census, KindActivations) and (
184 ctx.current_node != LayerType.SEL_REC_LAYER or census.ops is not None
185 or census.selective is not None and runs_hyper_selective(ccfg, EvalUtils.switches(ccfg, ctx))):
186 return EvalBody.census_activ(ccfg, ctx, census)
187 attn_size = sum(
188 [
189 ctx.attn_qkv_activ(ccfg, ctx),
190 ctx.attn_score_activ(ccfg, ctx),
191 ctx.attn_proj_activ(ccfg, ctx),
192 ]
193 )
194 if ccfg.n_exp == 1:
195 ffn_size = ctx.ffn_activ(ccfg, ctx)
196 else:
197 ffn_size = ctx.ffn_moe_activ(ccfg, ctx)
198 norm_size = ctx.norm_activ(ccfg, ctx)
199 return attn_size + ffn_size + norm_size
201 @staticmethod
202 def census_activ(ccfg: CostModelConfig, ctx: Context, census: KindActivations) -> float:
203 """A layer's activations, as the census of its kind measured them.
205 What the layer keeps between its passes, or its backward's working
206 set (:meth:`EvalUtils.census_bytes`), per token of a CP rank's share
207 of the sequence: TP splits one part, sequence parallelism the
208 other, and CP that gathers keys and values leaves them whole
209 (:meth:`EvalAttn.kv_shards`). A selective layer keeps what
210 :meth:`census_kept` says, and its backward holds what the plain
211 layer's does.
212 """
213 tokens = ctx.micro_factor * ccfg.s * ccfg.b / max(1, ccfg.cp)
214 held = census.working / max(1, ccfg.sp) + census.working_tp / max(1, ccfg.t)
215 # A census counts a rank's share of the keys and values; where CP
216 # gathers them the attention keeps the rest of the sequence's too.
217 gathered = EvalAttn.gathered_kv_bytes(ccfg)
218 if ctx.current_node == LayerType.SEL_REC_LAYER:
219 return EvalUtils.census_bytes(ctx, tokens, EvalBody.census_kept(ccfg, ctx, census), held + gathered)
220 kept = census.saved / max(1, ccfg.sp) + census.saved_tp / max(1, ccfg.t)
221 return EvalUtils.census_bytes(ctx, tokens, kept + gathered, held + gathered)
223 @staticmethod
224 def census_kept(ccfg: CostModelConfig, ctx: Context, census: KindActivations) -> float:
225 """What a selective layer of a kind *census* prices keeps, bytes per token of a CP rank's share.
227 Under HyperParallel's selective checkpointing, whose switches it has
228 (:func:`runs_hyper_selective`), what the census measured under it;
229 the keys and values CP gathers it gathers again to recompute. Under
230 other switches, what the census's records per op state for each op
231 they keep (:func:`load_op_records`): an op no switch names, and the
232 census's ``other``, whatever they say; the keys and values CP
233 gathers where the attention's batched matmuls are kept. Split as
234 what the plain layer keeps (:meth:`census_activ`).
235 """
236 if census.selective is not None and runs_hyper_selective(ccfg, EvalUtils.switches(ccfg, ctx)):
237 return census.selective / max(1, ccfg.sp) + census.selective_tp / max(1, ccfg.t)
238 records = load_op_records().ops
240 def keeps(op: str) -> bool:
241 """Whether the layer keeps what the census states for *op*."""
242 switch = records[op].switch if op in records else None
243 return switch is None or bool(EvalUtils.switch(ccfg, ctx, switch))
245 kept = sum(size for op, size in census.ops.items() if keeps(op)) / max(1, ccfg.sp)
246 kept += sum(size for op, size in census.ops_tp.items() if keeps(op)) / max(1, ccfg.t)
247 return kept + (EvalAttn.gathered_kv_bytes(ccfg) if keeps("attBMM") else 0)
249 @staticmethod
250 def census_working_terms(ccfg: CostModelConfig, ctx: Context) -> Optional[Tuple[float, float]]:
251 """What a selective layer priced by its census's records per op keeps, and what its backward holds.
253 In bytes, at one micro-batch: the two terms of its working set as
254 warm-up ends (:meth:`EvalUtils.census_bytes`), which is not the sum
255 of its ops' where the layer keeps more than its backward holds.
256 ``None`` for a layer its census does not price so: plain, fully
257 recomputed, under HyperParallel's policy, or without a census's
258 records per op.
259 """
260 census = getattr(ccfg, "kind_activations", None)
261 if (not isinstance(census, KindActivations) or census.ops is None
262 or ctx.current_node != LayerType.SEL_REC_LAYER
263 or census.selective is not None and runs_hyper_selective(ccfg, EvalUtils.switches(ccfg, ctx))):
264 return None
265 tokens = ccfg.s * ccfg.b / max(1, ccfg.cp)
266 held = census.working / max(1, ccfg.sp) + census.working_tp / max(1, ccfg.t)
267 return tokens * EvalBody.census_kept(ccfg, ctx, census), tokens * (held + EvalAttn.gathered_kv_bytes(ccfg))
269 # Full recompute
271 @staticmethod
272 def fullrec_layer_activ(ccfg: CostModelConfig, ctx: Context) -> float:
273 """activations"""
274 micro_factor = ctx.micro_factor
275 forward_activation = (
276 micro_factor * ccfg.bytes_compute * ccfg.s * ccfg.b * ccfg.h
277 )
278 forward_activation /= ccfg.shard_recompute_input
279 return forward_activation
281 @staticmethod
282 def fullrec_layer_activ_gradclip(
283 ccfg: CostModelConfig, ctx: Context
284 ) -> float:
285 """special case with gradient clipping"""
286 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx)
287 grad_clip_mem = (
288 non_exp_p
289 + routed_p / ccfg.ep * ccfg.bytes_os / ccfg.shard_p_os_exp
290 + shared_p * ccfg.bytes_os / ccfg.shard_p_os_exp_partial
291 )
292 grad_clip_mem *= ccfg.bytes_os / ccfg.shard_p_os_non_exp_partial
293 grad_clip_mem *= int(ccfg.has_clip)
294 forward_activation = EvalBody.fullrec_layer_activ(ccfg, ctx)
295 dp_comm_size = ctx.eval.dyn.comm.dp(ccfg, ctx)
296 if forward_activation + dp_comm_size > grad_clip_mem:
297 return forward_activation
298 logger.debug(
299 "gradient clipping %s > %s",
300 EvalUtils.mb(grad_clip_mem),
301 EvalUtils.mb(forward_activation + dp_comm_size),
302 )
303 return grad_clip_mem
305 @staticmethod
306 def fullrec_layer_comm_gradclip(
307 ccfg: CostModelConfig, ctx: Context
308 ) -> float:
309 """special case with gradient clipping"""
310 if EvalBody.fullrec_layer_activ_gradclip(ccfg, ctx) > 0:
311 return ctx.eval.dyn.comm.dp(ccfg, ctx)
312 return 0
314 @staticmethod
315 def act_cp_layer(
316 ccfg: CostModelConfig,
317 ctx: Context
318 ) -> CPMemoryBreakdown:
319 """Estimate CP activation memory impact for one transformer layer.
321 Ring CP (colossalai_cp / hybrid_cp):
322 Each rank holds s/cp tokens (Q sharded along seq) and a/t heads;
323 KV is all-gathered, so only one S dim of the S² score tensor
324 is divided.
325 KV cache: (s/cp) × b × kv_dim_per_rank
326 Attn scores: (s/cp) × s × b × (a/t)
328 Ulysses CP:
329 Each rank holds all s tokens but a/(t*cp) heads.
330 KV cache: s × b × kv_dim_per_rank / cp
331 Attn scores: s × s × b × (a/(t*cp))
332 """
333 s, b = ccfg.s, ccfg.b
334 a = ccfg.a
335 t = max(1, ccfg.t)
337 if a <= 0:
338 raise ValueError(f"Number of attention heads must be positive, got {a}")
340 cp = ccfg.cp
342 if cp <= 0:
343 raise ValueError(f"CP degree must be positive, got {cp}")
345 attention_type = detect_attention_type(ccfg)
346 cp_algo = _resolve_cp_algo(ccfg)
348 kv_dim = compute_kv_dim(ccfg)
349 a_per_rank = a / t
351 if cp_algo == CPAlgo.ULYSSES_CP:
352 terms = _ulysses_cp_attention_terms(s, b, cp, kv_dim, a_per_rank)
353 else:
354 terms = _ring_cp_attention_terms(s, b, cp, kv_dim, a_per_rank)
356 comm_buffer = EvalLayerComm.cp_comm_buffer(ccfg, ctx)
358 total_memory = (
359 terms.kv_cache_memory + terms.attention_scores_memory +
360 terms.softmax_outputs_memory + terms.dropout_mask_memory + comm_buffer
361 )
362 total_reduction = terms.s2_reduction + terms.kv_reduction - comm_buffer
364 return CPMemoryBreakdown(
365 kv_cache_memory=terms.kv_cache_memory,
366 attention_scores_memory=terms.attention_scores_memory,
367 softmax_outputs_memory=terms.softmax_outputs_memory,
368 dropout_mask_memory=terms.dropout_mask_memory,
369 comm_buffer_memory=comm_buffer,
370 kv_reduction=terms.kv_reduction,
371 s2_reduction=terms.s2_reduction,
372 total_reduction=total_reduction,
373 total_memory=total_memory,
374 cp_degree=int(cp),
375 seq_len=int(s),
376 attention_type=attention_type,
377 cp_algo=cp_algo,
378 )