« 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"""Head submodule"""
16from __future__ import annotations
17from typing import TYPE_CHECKING
19from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import param_cp
21if TYPE_CHECKING:
22 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
23 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
26class EvalHead:
27 """Head layer formulas class"""
29 @staticmethod
30 def num_params_embed(ccfg: CostModelConfig, _) -> float:
31 """Parameter count"""
32 return ccfg.h * ccfg.v
34 @staticmethod
35 def shares_output_table(ccfg: CostModelConfig) -> bool:
36 """Whether the embedding's table is the output layer's, priced there.
38 A tied table is one only where both ends run on one stage; under
39 pipeline parallelism the first and the last stage each hold a copy.
40 """
41 return bool(ccfg.tie_emb_out) and ccfg.p <= 1
43 @staticmethod
44 def stat_embed_p(ccfg: CostModelConfig, ctx: Context) -> float:
45 """model param"""
46 if EvalHead.shares_output_table(ccfg):
47 return 0
48 param_size = ctx.eval.num_p(ccfg, ctx)
49 param_size /= ccfg.shard_embed
50 b_p = ccfg.bytes_p
51 b_p /= param_cp(ccfg)
52 return param_size * b_p
54 @staticmethod
55 def stat_embed_os(ccfg: CostModelConfig, ctx: Context) -> float:
56 """optim state"""
57 if ctx.swap_os or EvalHead.shares_output_table(ccfg):
58 return 0
59 param_size = ctx.eval.num_p(ccfg, ctx)
60 param_size /= ccfg.shard_embed
61 b_os = ccfg.bytes_optim_table
62 b_os /= param_cp(ccfg)
63 return param_size * b_os
65 @staticmethod
66 def stat_embed_grad(ccfg: CostModelConfig, ctx: Context) -> float:
67 """gradient"""
68 if EvalHead.shares_output_table(ccfg):
69 return 0
70 param_size = ctx.eval.num_p(ccfg, ctx)
71 param_size /= ccfg.shard_embed
72 b_grad = ccfg.bytes_grad
73 b_grad /= param_cp(ccfg)
74 return param_size * b_grad
76 @staticmethod
77 def dp_comm_embed(ccfg: CostModelConfig, ctx: Context) -> float:
78 """DP Communication size: the table as the layer gathers it to compute."""
79 return (
80 ccfg.comm_d_non_exp
81 * ctx.eval.num_p(ccfg, ctx)
82 / (ccfg.shard_embed / (ccfg.gather_embed or 1) * param_cp(ccfg))
83 )
85 @staticmethod
86 def tp_comm_embed(ccfg: CostModelConfig, _) -> float:
87 """TP Communication size.
89 The embedding is never recomputed, so it keeps its buffer whatever
90 the gather switch says: a switch acts only in a selective layer.
91 """
92 return (
93 ccfg.comm_t
94 * ccfg.s
95 * ccfg.h
96 * ccfg.b
97 * (ccfg.t - 1)
98 / (ccfg.t * ccfg.cp)
99 )
101 @staticmethod
102 def activ_embed(ccfg: CostModelConfig, ctx: Context) -> float:
103 """activations"""
104 micro_factor = ctx.micro_factor
105 activ_size = micro_factor * ccfg.bytes_compute / (ccfg.t * ccfg.cp)
106 activ_size *= ccfg.s * ccfg.b * ccfg.h
107 return activ_size