« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025 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"""Tail submodule"""
16from __future__ import annotations
17from typing import TYPE_CHECKING
18from hyper_parallel.auto_parallel._model_spec import KindActivations
19from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
20from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import param_cp
21from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.utils import EvalUtils
23if TYPE_CHECKING:
24 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
25 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
28class EvalMTP:
29 """MTP"""
31 @staticmethod
32 def num_params_mtp(ccfg: CostModelConfig, _) -> float:
33 """Param count (MTP)"""
34 # Linear + Norm
35 return 2 * ccfg.h * ccfg.h + 4 * ccfg.h
37 @staticmethod
38 def stat_mtp_p(ccfg: CostModelConfig, ctx: Context) -> float:
39 """static mem for model param (MTP)"""
40 if not ccfg.n_mtp:
41 return 0
42 extra_param_size = EvalMTP.num_params_mtp(ccfg, ctx)
43 b_p = ccfg.bytes_p
44 if ccfg.is_shard_mtp_param:
45 b_p /= ccfg.shard_p_os_non_exp_partial
46 extra = ccfg.n_mtp * extra_param_size * b_p
47 # Shared Head
48 ctx.current_node = LayerType.EMBEDDING_LAYER
49 head = ccfg.n_mtp * ctx.eval.stat.p(ccfg, ctx)
50 # Shared Tail
51 ctx.current_node = LayerType.OUTPUT_LAYER
52 tail = ccfg.n_mtp * EvalTailSingle.stat_output_single_p(ccfg, ctx)
53 return head + extra + tail
55 @staticmethod
56 def stat_mtp_os(ccfg: CostModelConfig, ctx: Context) -> float:
57 """static mem for optimizer states (MTP)"""
58 if not ccfg.n_mtp or ctx.swap_os:
59 return 0
60 extra_param_size = EvalMTP.num_params_mtp(ccfg, ctx)
61 b_os = ccfg.bytes_optim
62 if ccfg.is_shard_mtp_param:
63 b_os /= ccfg.shard_p_os_non_exp_partial
64 extra = ccfg.n_mtp * extra_param_size * b_os
65 # Shared Head
66 ctx.current_node = LayerType.EMBEDDING_LAYER
67 head = ccfg.n_mtp * ctx.eval.stat.os(ccfg, ctx)
68 # Shared Tail
69 ctx.current_node = LayerType.OUTPUT_LAYER
70 tail = ccfg.n_mtp * EvalTailSingle.stat_output_single_os(ccfg, ctx)
71 return head + extra + tail
73 @staticmethod
74 def stat_mtp_grad(ccfg: CostModelConfig, ctx: Context) -> float:
75 """static mem for gradients (MTP)"""
76 if not ccfg.n_mtp:
77 return 0
78 extra_param_size = EvalMTP.num_params_mtp(ccfg, ctx)
79 b_grad = ccfg.bytes_grad
80 if ccfg.is_shard_mtp_param:
81 b_grad /= ccfg.shard_grad_non_exp
82 extra = ccfg.n_mtp * extra_param_size * b_grad
83 # Shared Head
84 ctx.current_node = LayerType.EMBEDDING_LAYER
85 head = ccfg.n_mtp * ctx.eval.stat.grad(ccfg, ctx)
86 # Shared Tail
87 ctx.current_node = LayerType.OUTPUT_LAYER
88 tail = ccfg.n_mtp * EvalTailSingle.stat_output_single_grad(ccfg, ctx)
89 return head + extra + tail
91 @staticmethod
92 def activ_mtp(ccfg: CostModelConfig, ctx: Context) -> float:
93 """activation mem (MTP)"""
94 if not ccfg.n_mtp:
95 return 0
96 micro_factor = ctx.micro_factor
97 res = micro_factor * ccfg.n_mtp * ccfg.bytes_compute
98 # A CP rank runs the layer on its own chunk of the sequence.
99 res *= ccfg.s * ccfg.b * 3 * ccfg.h / max(1, ccfg.cp)
100 # Shared Head
101 ctx.current_node = LayerType.EMBEDDING_LAYER
102 res += ccfg.n_mtp * ctx.eval.dyn.activation(ccfg, ctx)
103 # Shared Tail
104 ctx.current_node = LayerType.OUTPUT_LAYER
105 res += ccfg.n_mtp * EvalTailSingle.activ_out_single(ccfg, ctx)
106 return res
108 @staticmethod
109 def comm_mtp(ccfg: CostModelConfig, ctx: Context) -> float:
110 """communication mem (MTP)"""
111 if not ccfg.n_mtp:
112 return 0
113 mtp_dp_comm_size = 0
114 param_size = EvalMTP.num_params_mtp(ccfg, ctx)
115 mtp_dp_comm_size += (
116 ccfg.comm_d_non_exp * ccfg.n_mtp * param_size / (ccfg.t * param_cp(ccfg))
117 )
118 # Shared Head
119 ctx.current_node = LayerType.EMBEDDING_LAYER
120 param_size = ctx.eval.num_p(ccfg, ctx)
121 mtp_dp_comm_size += (
122 ccfg.comm_d_non_exp * ccfg.n_mtp * param_size / (ccfg.t * param_cp(ccfg))
123 )
124 mtp_dp_comm_size += ccfg.n_mtp * ctx.eval.dyn.comm.dp(ccfg, ctx)
125 # Shared Tail
126 ctx.current_node = LayerType.OUTPUT_LAYER
127 param_size = ctx.eval.num_p(ccfg, ctx)
128 mtp_dp_comm_size += (
129 ccfg.comm_d_non_exp * ccfg.n_mtp * param_size / (ccfg.t * param_cp(ccfg))
130 )
131 mtp_dp_comm_size += ccfg.n_mtp * EvalTailSingle.comm_out_single(
132 ccfg, ctx
133 )
134 return mtp_dp_comm_size
137class EvalTailSingle:
138 """Single tail layer formulas class"""
140 @staticmethod
141 def stat_output_single_p(ccfg: CostModelConfig, ctx: Context) -> float:
142 """static mem for model param (lmhead)"""
143 param_size = ctx.eval.num_p(ccfg, ctx)
144 b_p = ccfg.bytes_p
145 b_p /= ccfg.shard_p_os_non_exp_partial
146 return param_size * b_p
148 @staticmethod
149 def stat_output_single_os(ccfg: CostModelConfig, ctx: Context) -> float:
150 """static mem for optim state (lmhead)"""
151 if ctx.swap_os:
152 return 0
153 param_size = ctx.eval.num_p(ccfg, ctx)
154 b_os = ccfg.bytes_optim_table
155 b_os /= ccfg.shard_p_os_non_exp_partial
156 return param_size * b_os
158 @staticmethod
159 def stat_output_single_grad(ccfg: CostModelConfig, ctx: Context) -> float:
160 """static mem for gradient (lmhead)"""
161 param_size = ctx.eval.num_p(ccfg, ctx)
162 b_grad = ccfg.bytes_grad
163 b_grad /= ccfg.shard_grad_non_exp
164 return param_size * b_grad
166 @staticmethod
167 def activ_out_single(ccfg: CostModelConfig, ctx: Context) -> float:
168 """activation mem (lmhead)"""
169 census = getattr(ccfg, "output_census", None)
170 if isinstance(census, KindActivations):
171 return EvalTailSingle.census_out_single(ccfg, ctx, census)
172 micro_factor = ctx.micro_factor
173 last_norm = ccfg.s * ccfg.b * ccfg.bytes_norm * ccfg.h
174 lm_head = ccfg.s * ccfg.b * ccfg.bytes_compute * ccfg.v
175 activ_size = last_norm + lm_head
176 # A CP rank computes the logits of its own chunk of the sequence: the
177 # output layer gathers none across CP.
178 activ_size /= ccfg.shard_output_activ * max(1, ccfg.cp)
179 return micro_factor * activ_size
181 @staticmethod
182 def census_out_single(ccfg: CostModelConfig, ctx: Context, census: KindActivations) -> float:
183 """The output layer's activations, as its census measured them (lmhead).
185 What the layer keeps between its passes, or its backward's working
186 set (:meth:`EvalUtils.census_bytes`), per token of a CP rank's share
187 of the sequence. Sequence parallelism splits the part the
188 vocabulary does not scale, and TP the part it does where the loss
189 runs on logits sharded over it (``shards_logits``); otherwise every
190 TP rank holds them whole.
191 """
192 tokens = ctx.micro_factor * ccfg.s * ccfg.b / max(1, ccfg.cp)
193 vocab_shards = max(1, ccfg.t) if ccfg.shards_logits else 1
194 kept = census.saved / max(1, ccfg.sp) + census.saved_tp / vocab_shards
195 held = census.working / max(1, ccfg.sp) + census.working_tp / vocab_shards
196 return EvalUtils.census_bytes(ctx, tokens, kept, held)
198 @staticmethod
199 def comm_out_single(ccfg: CostModelConfig, ctx: Context) -> float:
200 """communicaiton mem (lmhead)"""
201 return (
202 ccfg.comm_d_non_exp
203 * ctx.eval.num_p(ccfg, ctx)
204 / (ccfg.t * param_cp(ccfg))
205 )
207 @staticmethod
208 def reduced_grad_out_single(ccfg: CostModelConfig, ctx: Context) -> tuple:
209 """The output table's gradient FSDP reduce-scatters, whole and sharded (lmhead)."""
210 return EvalUtils.reduced_grads(ccfg, ((ctx.eval.num_p(ccfg, ctx), ccfg.t, ccfg.shard_grad_non_exp),))
213class EvalTail:
214 """Single tail layer formulas class"""
216 @staticmethod
217 def num_params_output(ccfg: CostModelConfig, _) -> float:
218 """Parameters count: the output table, and the final norm of a model that states its norms.
220 One that does not is counted a bias per vocabulary entry instead,
221 the formulas' convention.
222 """
223 if getattr(ccfg, "layer_norms", None) is None:
224 return ccfg.h * ccfg.v + ccfg.v
225 return ccfg.h * ccfg.v + (2 if getattr(ccfg, "norm_bias", None) else 1) * ccfg.h
227 @staticmethod
228 def stat_output_p(ccfg: CostModelConfig, ctx: Context) -> float:
229 """total model param"""
230 return sum(
231 [
232 EvalTailSingle.stat_output_single_p(ccfg, ctx),
233 EvalMTP.stat_mtp_p(ccfg, ctx),
234 ]
235 )
237 @staticmethod
238 def stat_output_os(ccfg: CostModelConfig, ctx: Context) -> float:
239 """total optim state"""
240 return sum(
241 [
242 EvalTailSingle.stat_output_single_os(ccfg, ctx),
243 EvalMTP.stat_mtp_os(ccfg, ctx),
244 ]
245 )
247 @staticmethod
248 def stat_output_grad(ccfg: CostModelConfig, ctx: Context) -> float:
249 """total gradients"""
250 return sum(
251 [
252 EvalTailSingle.stat_output_single_grad(ccfg, ctx),
253 EvalMTP.stat_mtp_grad(ccfg, ctx),
254 ]
255 )
257 @staticmethod
258 def activ_output(ccfg: CostModelConfig, ctx: Context) -> float:
259 """total activations"""
260 return sum(
261 [
262 EvalTailSingle.activ_out_single(ccfg, ctx),
263 EvalMTP.activ_mtp(ccfg, ctx),
264 ]
265 )
267 @staticmethod
268 def comm_output(ccfg: CostModelConfig, ctx: Context) -> float:
269 """total communications"""
270 return sum(
271 [
272 EvalTailSingle.comm_out_single(ccfg, ctx),
273 EvalMTP.comm_mtp(ccfg, ctx),
274 ]
275 )