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"""cost model parser module"""
16from __future__ import annotations
17from typing import TYPE_CHECKING
18
19import math
20from abc import ABC
21from abc import abstractmethod
22
23if TYPE_CHECKING:
24 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import _CostModVar
25
26
27class _CostModelParser(ABC):
28 """abstract parser class"""
29
30 def __init__(self, ccfg: _CostModVar):
31 self.ccfg = ccfg
32 self.config = ccfg.config
33
34 @abstractmethod
35 def parse(self):
36 """Parse the cost model configuration and populate the cost model variables.
37
38 Subclasses must implement this method to read framework-specific
39 configuration values into the shared _CostModVar instance.
40 """
41
42 def config_optimizer_shard(self, ccfg):
43 """OP related variables"""
44 # Non expert params
45 ccfg.shard_p_os_non_exp_partial = (
46 ccfg.os_max_shard if ccfg.has_op else ccfg.t
47 ) * ccfg.cp
48 ccfg.shard_p_os_non_exp = (
49 (ccfg.d if ccfg.has_op else 1) * ccfg.cp * ccfg.t
50 )
51 ccfg.shard_grad_non_exp = (
52 ccfg.shard_p_os_non_exp if ccfg.has_grad_shard else ccfg.t
53 )
54
55 # Expert params
56 ccfg.shard_p_os_exp_partial = math.gcd(
57 ccfg.n_exp,
58 (ccfg.os_max_shard if ccfg.has_op else 1) * ccfg.t_exp,
59 )
60 ccfg.shard_p_os_exp = (
61 (ccfg.d_exp if ccfg.has_op else 1) * ccfg.cp * ccfg.t_exp
62 )
63 ccfg.shard_grad_exp = (
64 ccfg.shard_p_os_exp
65 if ccfg.has_grad_shard
66 else ccfg.t_exp
67 )
68 ccfg.shard_grad_exp_partial = (
69 ccfg.shard_p_os_exp_partial
70 if ccfg.has_grad_shard
71 else ccfg.t_exp
72 )
73
74 # def config_op_level(self, ccfg, strategy):
75 # def full_partial():
76 # return Config({"full":0, "partial":0})
77 # def exp_or_not():
78 # return Config({
79 # "non_exp":full_partial(),
80 # "exp":full_partial()
81 # })
82 # ccfg.op = Config({
83 # "p":exp_or_not(),
84 # "os":exp_or_not(),
85 # "grad"exp_or_not()
86 # })
87 # shard_strat = {
88 # "grad":0, #zero 1
89 # "os+grad":0, #zero 2
90 # "p+os+grad":0, # zero 3
91 # "p+os":0 # zero2 mindspore
92 # }
93 # shard_strat[strategy]
94
95 def init_hff(self):
96 """MindFormers format for FFn hidden size"""
97 # Assuming following 3 variables are already parsed
98 hidden_size = self.ccfg.h
99 ffn_dim_multiplier = self.ccfg.fdm
100 multiple_of = self.ccfg.multiple_of
101 hff = 4 * hidden_size
102 if ffn_dim_multiplier:
103 hff = int((ffn_dim_multiplier + 0.01) * hff)
104 hff = int(2 * hff / 3)
105 hff = multiple_of * ((hff + multiple_of - 1) // multiple_of)
106 return hff
107
108 def config_comm_flag(self, ccfg):
109 """comm flag variables"""
110 ccfg.comm_d_non_exp = (
111 0
112 if ((ccfg.d == 1) or not ccfg.has_op)
113 else (2 if not ccfg.has_grad_shard else 3)
114 ) # data parallel comm factor
115 ccfg.comm_d_exp = (
116 0
117 if ((ccfg.d_exp == 1) or not ccfg.has_op)
118 else (2 if not ccfg.has_grad_shard else 3)
119 ) # data parallel comm factor
120 ccfg.comm_t = float(ccfg.t > 1) # tensor parallel comm factor
121 ccfg.comm_ep = float(
122 ccfg.ep > 1 or ccfg.n_exp > 1
123 ) # expert parallel comm factor
124 ccfg.comm_cp = float(ccfg.cp > 1) # context parallel comm factor
125
126 def config_dp_tp_exp(self, ccfg):
127 """MoE strategy variables"""
128 if ccfg.etp:
129 ccfg.t_exp = ccfg.etp
130 # d * t = inner dp * outer dp * etp
131 # inner dp = EP, outer dp = the rest
132 ccfg.d_exp = ccfg.d * ccfg.t * ccfg.cp // ccfg.t_exp // ccfg.ep
133 else:
134 ccfg.t_exp = ccfg.t
135 if ccfg.d >= ccfg.ep:
136 ccfg.d_exp = ccfg.d // ccfg.ep
137 else:
138 ccfg.d_exp = ccfg.d * ccfg.t // ccfg.ep
139 if ccfg.t_exp * ccfg.ep > ccfg.d * ccfg.t:
140 ccfg.t_exp = 1
141
142 exp_group1_invalid = ccfg.d_exp < 1 or ccfg.t_exp < 1
143 exp_group2_invalid = ccfg.hff_exp < 1 or ccfg.n_exp < 1
144 if exp_group1_invalid or exp_group2_invalid:
145 raise TypeError(
146 f"MoE parsing error: d_exp({ccfg.d_exp})/t_exp({ccfg.t_exp})/"
147 f"hff_exp({ccfg.hff_exp})/n_exp({ccfg.n_exp})/"
148 f"DP = {ccfg.d}, TP = {ccfg.t}, EP = {ccfg.ep}/"
149 )