« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +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
18from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
19from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.utils import EvalUtils
20from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.comm import EvalLayerComm
21from hyper_parallel.auto_parallel.sapp_nd.nd.common.cp_types import (
22 CPMemoryBreakdown,
23 CPAlgo,
24 _resolve_cp_algo,
25)
26from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import (
27 detect_attention_type,
28 compute_kv_dim,
29)
31if TYPE_CHECKING:
32 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
33 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
34 from typing import Tuple
37class EvalBody:
38 """Body layer formulas class"""
40 @staticmethod
41 def num_params_layer(
42 ccfg: CostModelConfig, ctx: Context
43 ) -> Tuple[float, float, float]:
44 """Parameters count.
46 Returns a 3-tuple (non_exp, routed, shared):
47 - non_exp: attention + norm params (and dense FFN if n_exp==1)
48 - routed: routed expert params (0 if n_exp==1)
49 - shared: shared expert params (0 if n_shared_exp==0 or no pointer)
50 """
51 non_exp = ctx.attn_num_p(ccfg, ctx) + ctx.norm_num_p(ccfg, ctx)
52 routed = 0.0
53 shared = 0.0
54 if ccfg.n_exp == 1:
55 non_exp += ctx.ffn_num_p(ccfg, ctx)
56 else:
57 if ctx.ffn_routed_num_p is not None:
58 routed = ctx.ffn_routed_num_p(ccfg, ctx)
59 if ctx.ffn_shared_num_p is not None:
60 shared = ctx.ffn_shared_num_p(ccfg, ctx)
61 return (non_exp, routed, shared)
63 @staticmethod
64 def stat_p_layer(ccfg: CostModelConfig, ctx: Context) -> float:
65 """model param"""
66 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx)
67 # Routed experts: EP sharding
68 routed_mem = routed_p / ccfg.ep * ccfg.bytes_p / ccfg.shard_p_os_exp
69 # Shared experts: partial DP sharding
70 shared_mem = shared_p * ccfg.bytes_p / ccfg.shard_p_os_exp_partial
71 # Non expert
72 non_exp_mem = non_exp_p * ccfg.bytes_p / ccfg.shard_p_os_non_exp_partial
73 return non_exp_mem + routed_mem + shared_mem
75 @staticmethod
76 def stat_os_layer(ccfg: CostModelConfig, ctx: Context) -> float:
77 """optim state"""
78 if ctx.swap_os:
79 return 0
80 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx)
81 # Routed experts
82 routed_mem = routed_p / ccfg.ep * 2 * ccfg.bytes_os / ccfg.shard_p_os_exp
83 # Shared experts
84 shared_mem = shared_p * 2 * ccfg.bytes_os / ccfg.shard_p_os_exp_partial
85 # Non expert
86 non_exp_mem = non_exp_p * 2 * ccfg.bytes_os / ccfg.shard_p_os_non_exp_partial
87 return non_exp_mem + routed_mem + shared_mem
89 @staticmethod
90 def stat_grad_layer(ccfg: CostModelConfig, ctx: Context) -> float:
91 """gradients"""
92 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx)
93 # Routed experts
94 routed_mem = routed_p / ccfg.ep * ccfg.bytes_grad / ccfg.shard_grad_exp
95 # Shared experts: use shard_grad_exp_partial (independent of os sharding)
96 shared_mem = shared_p * ccfg.bytes_grad / ccfg.shard_grad_exp_partial
97 # Non expert
98 non_exp_mem = non_exp_p * ccfg.bytes_grad / ccfg.shard_grad_non_exp
99 return non_exp_mem + routed_mem + shared_mem
101 # No recompute and select recompute
103 @staticmethod
104 def layer_activ(ccfg: CostModelConfig, ctx: Context) -> float:
105 """activations"""
106 attn_size = sum(
107 [
108 ctx.attn_qkv_activ(ccfg, ctx),
109 ctx.attn_score_activ(ccfg, ctx),
110 ctx.attn_proj_activ(ccfg, ctx),
111 ]
112 )
113 if ccfg.n_exp == 1:
114 ffn_size = ctx.ffn_activ(ccfg, ctx)
115 else:
116 ffn_size = ctx.ffn_moe_activ(ccfg, ctx)
117 norm_size = ctx.norm_activ(ccfg, ctx)
118 return attn_size + ffn_size + norm_size
120 # Full recompute
122 @staticmethod
123 def fullrec_layer_activ(ccfg: CostModelConfig, ctx: Context) -> float:
124 """activations"""
125 micro_factor = ctx.micro_factor
126 forward_activation = (
127 micro_factor * ccfg.bytes_compute * ccfg.s * ccfg.b * ccfg.h
128 )
129 forward_activation /= ccfg.shard_recompute_input
130 return forward_activation
132 @staticmethod
133 def fullrec_layer_activ_gradclip(
134 ccfg: CostModelConfig, ctx: Context
135 ) -> float:
136 """special case with gradient clipping"""
137 non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx)
138 grad_clip_mem = (
139 non_exp_p
140 + routed_p / ccfg.ep * ccfg.bytes_os / ccfg.shard_p_os_exp
141 + shared_p * ccfg.bytes_os / ccfg.shard_p_os_exp_partial
142 )
143 grad_clip_mem *= ccfg.bytes_os / ccfg.shard_p_os_non_exp_partial
144 grad_clip_mem *= int(ccfg.has_clip)
145 forward_activation = EvalBody.fullrec_layer_activ(ccfg, ctx)
146 dp_comm_size = ctx.eval.dyn.comm.dp(ccfg, ctx)
147 if forward_activation + dp_comm_size > grad_clip_mem:
148 return forward_activation
149 logger.debug(
150 "gradient clipping %s > %s",
151 EvalUtils.mb(grad_clip_mem),
152 EvalUtils.mb(forward_activation + dp_comm_size),
153 )
154 return grad_clip_mem
156 @staticmethod
157 def fullrec_layer_comm_gradclip(
158 ccfg: CostModelConfig, ctx: Context
159 ) -> float:
160 """special case with gradient clipping"""
161 if EvalBody.fullrec_layer_activ_gradclip(ccfg, ctx) > 0:
162 return ctx.eval.dyn.comm.dp(ccfg, ctx)
163 return 0
165 @staticmethod
166 def act_cp_layer(
167 ccfg: CostModelConfig,
168 ctx: Context
169 ) -> CPMemoryBreakdown:
170 """Estimate CP activation memory impact for one transformer layer.
172 Ring CP (colossalai_cp / hybrid_cp):
173 Each rank holds s/cp tokens (Q sharded along seq) and a/t heads;
174 KV is all-gathered, so only one S dim of the S² score tensor
175 is divided.
176 KV cache: (s/cp) × b × kv_dim_per_rank
177 Attn scores: (s/cp) × s × b × (a/t)
179 Ulysses CP:
180 Each rank holds all s tokens but a/(t*cp) heads.
181 KV cache: s × b × kv_dim_per_rank / cp
182 Attn scores: s × s × b × (a/(t*cp))
183 """
184 s, b = ccfg.s, ccfg.b
185 a = ccfg.a
186 t = max(1, ccfg.t)
188 if a <= 0:
189 raise ValueError(f"Number of attention heads must be positive, got {a}")
191 cp = ccfg.cp
193 if cp <= 0:
194 raise ValueError(f"CP degree must be positive, got {cp}")
196 attention_type = detect_attention_type(ccfg)
197 cp_algo = _resolve_cp_algo(ccfg)
199 attention_scores_bytes = 4
200 softmax_outputs_bytes = 4
201 dropout_mask_bytes = 1
202 fp16_bytes = 2
203 kv_bytes = fp16_bytes * 2
205 kv_dim = compute_kv_dim(ccfg)
206 a_per_rank = a / t
208 if cp_algo == CPAlgo.ULYSSES_CP:
209 a_per_cp_rank = a_per_rank / cp
210 kv_dim_per_cp_rank = kv_dim / cp
212 kv_cache_memory = kv_bytes * s * b * kv_dim_per_cp_rank
213 attention_scores_memory = attention_scores_bytes * s * s * b * a_per_cp_rank
214 softmax_outputs_memory = softmax_outputs_bytes * s * s * b * a_per_cp_rank
215 dropout_mask_memory = dropout_mask_bytes * s * s * b * a_per_cp_rank
217 s2_items_no_cp = (
218 (attention_scores_bytes + softmax_outputs_bytes + dropout_mask_bytes)
219 * s * s * b * a_per_rank
220 )
221 s2_items_with_cp = (
222 (attention_scores_bytes + softmax_outputs_bytes + dropout_mask_bytes)
223 * s * s * b * a_per_cp_rank
224 )
225 s2_items_reduction = s2_items_no_cp - s2_items_with_cp
226 kv_reduction = kv_bytes * s * b * kv_dim * ((cp - 1) / cp)
227 else:
228 kv_cache_memory = kv_bytes * (s / cp) * b * kv_dim
229 attention_scores_memory = attention_scores_bytes * (s / cp) * s * b * a_per_rank
230 softmax_outputs_memory = softmax_outputs_bytes * (s / cp) * s * b * a_per_rank
231 dropout_mask_memory = dropout_mask_bytes * (s / cp) * s * b * a_per_rank
233 s2_items_reduction = (
234 (attention_scores_bytes + softmax_outputs_bytes + dropout_mask_bytes)
235 * s * s * b * a_per_rank * ((cp - 1) / cp)
236 )
237 kv_reduction = kv_bytes * s * b * kv_dim * ((cp - 1) / cp)
239 comm_buffer = EvalLayerComm.cp_comm_buffer(ccfg, ctx)
241 total_memory = (
242 kv_cache_memory + attention_scores_memory +
243 softmax_outputs_memory + dropout_mask_memory + comm_buffer
244 )
245 total_reduction = s2_items_reduction + kv_reduction - comm_buffer
247 return CPMemoryBreakdown(
248 kv_cache_memory=kv_cache_memory,
249 attention_scores_memory=attention_scores_memory,
250 softmax_outputs_memory=softmax_outputs_memory,
251 dropout_mask_memory=dropout_mask_memory,
252 comm_buffer_memory=comm_buffer,
253 kv_reduction=kv_reduction,
254 s2_reduction=s2_items_reduction,
255 total_reduction=total_reduction,
256 total_memory=total_memory,
257 cp_degree=int(cp),
258 seq_len=int(s),
259 attention_type=attention_type,
260 cp_algo=cp_algo,
261 )