Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / common / arch_hooks.py: 91%
183 statements
« 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"""Custom variables per model (expert knowledge)"""
16import math
17from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config
18from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
19from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
22class CWrap:
23 """Temporary evaluator-like instance"""
25 def __init__(self, e) -> None:
26 self.ccfg = e
28 def set_ccfg(self, hook):
29 """Apply a hook to the wrapped config object."""
30 return hook(self.ccfg)
32 def get_model_name(self):
33 """Return model name from the wrapped config."""
34 return self.ccfg.model_name
36 def reset(self, e):
37 """Replace the wrapped config object."""
38 self.ccfg = e
40 def __getattr__(self, attr):
41 if attr not in self.__dict__:
42 return lambda *args, **kwargs: None
43 return self.__dict__[attr]
45 def set_strategy(self, **kwargs):
46 """Forward strategy updates to the wrapped config."""
47 self.ccfg.set_strategy(**kwargs)
49 def get_strategy(self):
50 """Return strategy from the wrapped config."""
51 return self.ccfg.get_strategy()
54def custom_default_transformer(ccfg):
55 """base"""
56 ccfg.n_attMM = 4 # num attention matmul
57 ccfg.n_attBMM = 2 # num attention batch matmul
58 ccfg.n_attParamCast = (
59 ccfg.n_attMM if not ccfg.has_op else 0
60 ) # num attention parameters cast
61 ccfg.n_ffMM = 3 # num feedforward matmul
62 ccfg.n_ffBMM = 0 # num feedforward batch matmul
63 ccfg.n_ffParamCast = (
64 ccfg.n_ffMM if not ccfg.has_op else 0
65 ) # num feedforward parameters cast
66 ccfg.n_softmax = 1 # num softmax
67 ccfg.n_dropout = 0 # num dropout
68 ccfg.n_normOp = 2 # num normalization
69 ccfg.n_gather = 4 # num gather (TP)
70 ccfg.bytes_grad = 4 if ccfg.p > 1 else 0 # gradients
71 ccfg.bytes_os = 4 # optimizer states
72 ccfg.bytes_dropout = 0 # dropout mask
73 ccfg.bytes_norm = 4 # normalization input
76def custom_llama2(ccfg):
77 """llama2"""
78 custom_default_transformer(ccfg)
79 ccfg.n_gather = 4 # num gather (TP)
80 ccfg.bytes_grad = 2 # gradients
83def custom_mixtral(ccfg):
84 """mixtral"""
85 ccfg.n_attMM = 4 # num attention matmul
86 ccfg.n_attBMM = 2 # num attention batch matmul
87 ccfg.n_attParamCast = (
88 ccfg.n_attMM if not ccfg.has_op else 0
89 ) # num attention parameters cast
90 ccfg.n_ffMM = 0 # num feedforward matmul
91 ccfg.n_ffBMM = 3 # num feedforward batch matmul
92 ccfg.n_ffParamCast = (
93 ccfg.n_ffMM if not ccfg.has_op else 0
94 ) # num feedforward parameters cast
95 ccfg.n_softmax = 2 # num softmax
96 ccfg.n_dropout = 0 # num dropout
97 ccfg.n_normOp = 5 # num normalization
98 ccfg.n_gather = 4 # num gather (TP)
99 ccfg.bytes_grad = 2 if ccfg.p > 1 else 0 # gradients
100 ccfg.bytes_os = 4 # optimizer states
101 ccfg.bytes_dropout = 0 # dropout mask
102 ccfg.bytes_norm = 4 # normalization input
103 ccfg.hff = ccfg.hff_exp
106def custom_t5(ccfg):
107 """t5"""
109 # Encoder + Decoder
110 def encode(c):
111 c.n_attMM = 4 # num attention matmul
112 c.n_attBMM = 1 # num attention batch matmul
113 c.n_attParamCast = (
114 c.n_attMM if not c.has_op else 0
115 ) # num attention parameters cast
116 c.n_ffMM = 2 # num feedforward matmul
117 c.n_ffBMM = 0 # num feedforward batch matmul
118 c.n_ffParamCast = (
119 c.n_ffMM if not c.has_op else 0
120 ) # num feedforward parameters cast
121 c.n_softmax = 2 # num softmax
122 c.n_dropout = 5 # num dropout
123 c.n_normOp = 2 # num normalization
124 c.n_gather = 4 # num gather (TP)
125 c.bytes_grad = 4 if c.p > 1 else 0 # gradients
126 c.bytes_os = 4 # optimizer states
127 c.bytes_dropout = 1 # dropout mask
128 c.bytes_norm = 4 # normalization input
130 def decode(c):
131 c.n_attMM = 8 # num attention matmul
132 c.n_attBMM = 2 # num attention batch matmul
133 c.n_attParamCast = (
134 c.n_attMM if not c.has_op else 0
135 ) # num attention parameters cast
136 c.n_ffMM = 2 # num feedforward matmul
137 c.n_ffBMM = 0 # num feedforward batch matmul
138 c.n_ffParamCast = (
139 c.n_ffMM if not c.has_op else 0
140 ) # num feedforward parameters cast
141 c.n_softmax = 4 # num softmax
142 c.n_dropout = 7 # num dropout
143 c.n_normOp = 3 # num normalization
144 c.n_gather = 6 # num gather (TP)
145 c.bytes_grad = 4 if c.p > 1 else 0 # gradients
146 c.bytes_os = 4 # optimizer states
147 c.bytes_dropout = 1 # dropout mask
148 c.bytes_norm = 4 # normalization input
150 def hook_encode(e):
151 if isinstance(e, CostModelConfig):
152 e = CWrap(e)
153 e.set_ccfg(encode)
155 def hook_decode(e):
156 if isinstance(e, CostModelConfig):
157 e = CWrap(e)
158 e.set_ccfg(decode)
160 ccfg.layer_custom_config = [
161 (ccfg.n_lay // 2, hook_encode),
162 (ccfg.n_lay // 2, hook_decode),
163 ]
166def custom_pangualpha(ccfg):
167 """pangualpha"""
168 ccfg.n_attMM = 4 # num attention matmul
169 ccfg.n_attBMM = 1 # num attention batch matmul
170 ccfg.n_attParamCast = (
171 ccfg.n_attMM if not ccfg.has_op else 0
172 ) # num attention parameters cast
173 ccfg.n_ffMM = 2 # num feedforward matmul
174 ccfg.n_ffBMM = 0 # num feedforward batch matmul
175 ccfg.n_ffParamCast = (
176 ccfg.n_ffMM if not ccfg.has_op else 0
177 ) # num feedforward parameters cast
178 ccfg.n_softmax = 2 # num softmax
179 ccfg.n_dropout = 5 # num dropout
180 ccfg.n_normOp = 4 # num normalization
181 ccfg.n_gather = 4 # num gather (TP)
182 ccfg.bytes_grad = 4 if ccfg.p > 1 else 0 # gradients
183 ccfg.bytes_os = 4 # optimizer states
184 ccfg.bytes_dropout = 1 # dropout mask
185 ccfg.bytes_norm = 4 # normalization input
188def custom_deepseek3(ccfg):
189 """deepseekv3"""
190 saved = Config({})
191 if ccfg.config_format == "yaml":
192 saved.hff = (
193 ccfg.config.model.model_config.intermediate_size
194 if ccfg.config.model.model_config.intermediate_size
195 else ccfg.parser.init_hff()
196 )
197 elif ccfg.config_format == "json":
198 saved.hff = ccfg.ffn_hidden_size
199 else:
200 saved.hff = ccfg.specs.inter_dim
201 if not saved.hff:
202 saved.hff = ccfg.specs.hidden_dim
203 if not saved.hff:
204 saved.hff = ccfg.h
205 saved.n_chosen_exp = ccfg.n_chosen_exp
206 saved.n_exp = ccfg.n_exp
207 saved.n_shared_exp = ccfg.n_shared_exp
208 saved.ep = ccfg.ep
209 custom_default_transformer(ccfg)
210 ccfg.dh = 128
212 def dense(c):
213 c.hff = saved.hff
214 c.n_chosen_exp = 1
215 c.n_exp = 1
216 c.n_shared_exp = 0
218 def moe(c):
219 c.hff = c.hff_exp
220 c.n_chosen_exp = saved.n_chosen_exp
221 c.n_exp = saved.n_exp
222 c.n_shared_exp = saved.n_shared_exp
224 def hook_dense(e):
225 if isinstance(e, CostModelConfig):
226 e = CWrap(e)
227 # e.ccfg.ep = 1
228 e.set_ccfg(dense)
229 e.ccfg.ep = 1
230 # e.set_strategy(ep=1)
232 def hook_moe(e):
233 if isinstance(e, CostModelConfig):
234 e = CWrap(e)
235 # e.ccfg.ep = saved.ep
236 e.set_ccfg(moe)
237 e.ccfg.ep = saved.ep
238 # e.set_strategy(ep=saved.ep)
240 n_moe = ccfg.n_lay - ccfg.k_1st_dense
241 ccfg.layer_custom_config = [
242 (ccfg.k_1st_dense, hook_dense),
243 (n_moe, hook_moe),
244 (ccfg.n_mtp, hook_moe if n_moe > 0 else hook_dense),
245 ]
248def custom_qwen(ccfg):
249 """qwen2"""
250 custom_default_transformer(ccfg)
251 # if "72b" in ccfg.model_name :
252 # ccfg.s = ccfg.s * 3/4
253 ccfg.shard_recompute_input = ccfg.t
254 ccfg.shard_output_activ = ccfg.t
255 # ccfg.bytes_grad = 4
258def custom_cm(ccfg):
259 """llama moe"""
260 shard_p_os_exp = ccfg.shard_p_os_exp_partial
261 shard_p_os_non_exp_partial = math.gcd(ccfg.n_exp, ccfg.shard_p_os_non_exp)
262 shard_embed = ccfg.t
263 custom_deepseek3(ccfg)
265 def custom_shard(c):
266 c.shard_p_os_exp = shard_p_os_exp
267 c.shard_p_os_non_exp_partial = shard_p_os_non_exp_partial
268 c.shard_embed = shard_embed
270 for idx, f in enumerate(ccfg.layer_custom_config):
272 def wrap_hook(e, f=f):
273 if isinstance(e, CostModelConfig):
274 e = CWrap(e)
275 f[1](e)
276 e.set_ccfg(custom_shard)
278 ccfg.layer_custom_config[idx] = (f[0], wrap_hook)
280 def num_params_norm_cm(c, _):
281 return c.n_normOp * 2 * c.h + 0.5 * c.n_attMM * c.dh
283 ccfg.overwrite_eval_functions["num_params_norm"] = num_params_norm_cm
286def check_and_apply_custom_hook(e):
287 """routing hooks"""
288 if isinstance(e, CostModelConfig):
289 e = CWrap(e)
290 map_modelname_custom = {
291 "llama2": custom_llama2,
292 "mixtral": custom_mixtral,
293 "t5": custom_t5,
294 "pangualpha": custom_pangualpha,
295 "deepseek": custom_deepseek3,
296 "qwen": custom_qwen,
297 "cm": custom_cm,
298 }
299 for k, v in map_modelname_custom.items():
300 if k in e.get_model_name().lower():
301 e.set_ccfg(v)
302 return
303 logger.warning(
304 "Hook not defined for: %s. Default one is chosen", e.get_model_name()
305 )
306 e.set_ccfg(custom_default_transformer)