« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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.sapp_nd.nd.common.layer_type import LayerType
20if TYPE_CHECKING:
21 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
22 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
25class EvalMTP:
26 """MTP"""
28 @staticmethod
29 def num_params_mtp(ccfg: CostModelConfig, _) -> float:
30 """Param count (MTP)"""
31 # Linear + Norm
32 return 2 * ccfg.h * ccfg.h + 4 * ccfg.h
34 @staticmethod
35 def stat_mtp_p(ccfg: CostModelConfig, ctx: Context) -> float:
36 """static mem for model param (MTP)"""
37 if not ccfg.n_mtp:
38 return 0
39 extra_param_size = EvalMTP.num_params_mtp(ccfg, ctx)
40 b_p = ccfg.bytes_p
41 if ccfg.is_shard_mtp_param:
42 b_p /= ccfg.shard_p_os_non_exp_partial
43 extra = ccfg.n_mtp * extra_param_size * b_p
44 # Shared Head
45 ctx.current_node = LayerType.EMBEDDING_LAYER
46 head = ccfg.n_mtp * ctx.eval.stat.p(ccfg, ctx)
47 # Shared Tail
48 ctx.current_node = LayerType.OUTPUT_LAYER
49 tail = ccfg.n_mtp * EvalTailSingle.stat_output_single_p(ccfg, ctx)
50 return head + extra + tail
52 @staticmethod
53 def stat_mtp_os(ccfg: CostModelConfig, ctx: Context) -> float:
54 """static mem for optimizer states (MTP)"""
55 if not ccfg.n_mtp or ctx.swap_os:
56 return 0
57 extra_param_size = EvalMTP.num_params_mtp(ccfg, ctx)
58 b_os = 2 * ccfg.bytes_os
59 if ccfg.is_shard_mtp_param:
60 b_os /= ccfg.shard_p_os_non_exp_partial
61 extra = ccfg.n_mtp * extra_param_size * b_os
62 # Shared Head
63 ctx.current_node = LayerType.EMBEDDING_LAYER
64 head = ccfg.n_mtp * ctx.eval.stat.os(ccfg, ctx)
65 # Shared Tail
66 ctx.current_node = LayerType.OUTPUT_LAYER
67 tail = ccfg.n_mtp * EvalTailSingle.stat_output_single_os(ccfg, ctx)
68 return head + extra + tail
70 @staticmethod
71 def stat_mtp_grad(ccfg: CostModelConfig, ctx: Context) -> float:
72 """static mem for gradients (MTP)"""
73 if not ccfg.n_mtp:
74 return 0
75 extra_param_size = EvalMTP.num_params_mtp(ccfg, ctx)
76 b_grad = ccfg.bytes_grad
77 if ccfg.is_shard_mtp_param:
78 b_grad /= ccfg.shard_grad_non_exp
79 extra = ccfg.n_mtp * extra_param_size * b_grad
80 # Shared Head
81 ctx.current_node = LayerType.EMBEDDING_LAYER
82 head = ccfg.n_mtp * ctx.eval.stat.grad(ccfg, ctx)
83 # Shared Tail
84 ctx.current_node = LayerType.OUTPUT_LAYER
85 tail = ccfg.n_mtp * EvalTailSingle.stat_output_single_grad(ccfg, ctx)
86 return head + extra + tail
88 @staticmethod
89 def activ_mtp(ccfg: CostModelConfig, ctx: Context) -> float:
90 """activation mem (MTP)"""
91 if not ccfg.n_mtp:
92 return 0
93 micro_factor = ctx.micro_factor
94 res = micro_factor * ccfg.n_mtp * ccfg.bytes_compute
95 res *= ccfg.s * ccfg.b * 3 * ccfg.h
96 # Shared Head
97 ctx.current_node = LayerType.EMBEDDING_LAYER
98 res += ccfg.n_mtp * ctx.eval.dyn.activation(ccfg, ctx)
99 # Shared Tail
100 ctx.current_node = LayerType.OUTPUT_LAYER
101 res += ccfg.n_mtp * EvalTailSingle.activ_out_single(ccfg, ctx)
102 return res
104 @staticmethod
105 def comm_mtp(ccfg: CostModelConfig, ctx: Context) -> float:
106 """communication mem (MTP)"""
107 if not ccfg.n_mtp:
108 return 0
109 mtp_dp_comm_size = 0
110 param_size = EvalMTP.num_params_mtp(ccfg, ctx)
111 mtp_dp_comm_size += (
112 ccfg.comm_d_non_exp * ccfg.n_mtp * param_size / (ccfg.t * ccfg.cp)
113 )
114 # Shared Head
115 ctx.current_node = LayerType.EMBEDDING_LAYER
116 param_size = ctx.eval.num_p(ccfg, ctx)
117 mtp_dp_comm_size += (
118 ccfg.comm_d_non_exp * ccfg.n_mtp * param_size / (ccfg.t * ccfg.cp)
119 )
120 mtp_dp_comm_size += ccfg.n_mtp * ctx.eval.dyn.comm.dp(ccfg, ctx)
121 # Shared Tail
122 ctx.current_node = LayerType.OUTPUT_LAYER
123 param_size = ctx.eval.num_p(ccfg, ctx)
124 mtp_dp_comm_size += (
125 ccfg.comm_d_non_exp * ccfg.n_mtp * param_size / (ccfg.t * ccfg.cp)
126 )
127 mtp_dp_comm_size += ccfg.n_mtp * EvalTailSingle.comm_out_single(
128 ccfg, ctx
129 )
130 return mtp_dp_comm_size
133class EvalTailSingle:
134 """Single tail layer formulas class"""
136 @staticmethod
137 def stat_output_single_p(ccfg: CostModelConfig, ctx: Context) -> float:
138 """static mem for model param (lmhead)"""
139 param_size = ctx.eval.num_p(ccfg, ctx)
140 b_p = ccfg.bytes_p
141 b_p /= ccfg.shard_p_os_non_exp_partial
142 return param_size * b_p
144 @staticmethod
145 def stat_output_single_os(ccfg: CostModelConfig, ctx: Context) -> float:
146 """static mem for optim state (lmhead)"""
147 if ctx.swap_os:
148 return 0
149 param_size = ctx.eval.num_p(ccfg, ctx)
150 b_os = 2 * ccfg.bytes_os
151 b_os /= ccfg.shard_p_os_non_exp_partial
152 return param_size * b_os
154 @staticmethod
155 def stat_output_single_grad(ccfg: CostModelConfig, ctx: Context) -> float:
156 """static mem for gradient (lmhead)"""
157 param_size = ctx.eval.num_p(ccfg, ctx)
158 b_grad = ccfg.bytes_grad
159 b_grad /= ccfg.shard_grad_non_exp
160 return param_size * b_grad
162 @staticmethod
163 def activ_out_single(ccfg: CostModelConfig, ctx: Context) -> float:
164 """activation mem (lmhead)"""
165 micro_factor = ctx.micro_factor
166 last_norm = ccfg.s * ccfg.b * ccfg.bytes_norm * ccfg.h
167 lm_head = ccfg.s * ccfg.b * ccfg.bytes_compute * ccfg.v
168 activ_size = last_norm + lm_head
169 activ_size /= ccfg.shard_output_activ
170 return micro_factor * activ_size
172 @staticmethod
173 def comm_out_single(ccfg: CostModelConfig, ctx: Context) -> float:
174 """communicaiton mem (lmhead)"""
175 return (
176 ccfg.comm_d_non_exp
177 * ctx.eval.num_p(ccfg, ctx)
178 / (ccfg.t * ccfg.cp)
179 )
182class EvalTail:
183 """Single tail layer formulas class"""
185 @staticmethod
186 def num_params_output(ccfg: CostModelConfig, _) -> float:
187 """Parameters count (lmhead)"""
188 return ccfg.h * ccfg.v + ccfg.v
190 @staticmethod
191 def stat_output_p(ccfg: CostModelConfig, ctx: Context) -> float:
192 """total model param"""
193 return sum(
194 [
195 EvalTailSingle.stat_output_single_p(ccfg, ctx),
196 EvalMTP.stat_mtp_p(ccfg, ctx),
197 ]
198 )
200 @staticmethod
201 def stat_output_os(ccfg: CostModelConfig, ctx: Context) -> float:
202 """total optim state"""
203 return sum(
204 [
205 EvalTailSingle.stat_output_single_os(ccfg, ctx),
206 EvalMTP.stat_mtp_os(ccfg, ctx),
207 ]
208 )
210 @staticmethod
211 def stat_output_grad(ccfg: CostModelConfig, ctx: Context) -> float:
212 """total gradients"""
213 return sum(
214 [
215 EvalTailSingle.stat_output_single_grad(ccfg, ctx),
216 EvalMTP.stat_mtp_grad(ccfg, ctx),
217 ]
218 )
220 @staticmethod
221 def activ_output(ccfg: CostModelConfig, ctx: Context) -> float:
222 """total activations"""
223 return sum(
224 [
225 EvalTailSingle.activ_out_single(ccfg, ctx),
226 EvalMTP.activ_mtp(ccfg, ctx),
227 ]
228 )
230 @staticmethod
231 def comm_output(ccfg: CostModelConfig, ctx: Context) -> float:
232 """total communications"""
233 return sum(
234 [
235 EvalTailSingle.comm_out_single(ccfg, ctx),
236 EvalMTP.comm_mtp(ccfg, ctx),
237 ]
238 )