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-08-04 05:18 +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 = int(ccfg.hff)
193 elif ccfg.config_format == "json":
194 saved.hff = ccfg.ffn_hidden_size
195 else:
196 saved.hff = ccfg.specs.inter_dim
197 if not saved.hff:
198 saved.hff = ccfg.specs.hidden_dim
199 if not saved.hff:
200 saved.hff = ccfg.h
201 saved.n_chosen_exp = ccfg.n_chosen_exp
202 saved.n_exp = ccfg.n_exp
203 saved.n_shared_exp = ccfg.n_shared_exp
204 saved.ep = ccfg.ep
205 custom_default_transformer(ccfg)
206 ccfg.dh = 128
208 def dense(c):
209 c.hff = saved.hff
210 c.n_chosen_exp = 1
211 c.n_exp = 1
212 c.n_shared_exp = 0
214 def moe(c):
215 c.hff = c.hff_exp
216 c.n_chosen_exp = saved.n_chosen_exp
217 c.n_exp = saved.n_exp
218 c.n_shared_exp = saved.n_shared_exp
220 def hook_dense(e):
221 if isinstance(e, CostModelConfig):
222 e = CWrap(e)
223 # e.ccfg.ep = 1
224 e.set_ccfg(dense)
225 e.ccfg.ep = 1
226 # e.set_strategy(ep=1)
228 def hook_moe(e):
229 if isinstance(e, CostModelConfig):
230 e = CWrap(e)
231 # e.ccfg.ep = saved.ep
232 e.set_ccfg(moe)
233 e.ccfg.ep = saved.ep
234 # e.set_strategy(ep=saved.ep)
236 n_moe = ccfg.n_lay - ccfg.k_1st_dense
237 ccfg.layer_custom_config = [
238 (ccfg.k_1st_dense, hook_dense),
239 (n_moe, hook_moe),
240 (ccfg.n_mtp, hook_moe if n_moe > 0 else hook_dense),
241 ]
244def custom_qwen(ccfg):
245 """qwen2"""
246 custom_default_transformer(ccfg)
247 # if "72b" in ccfg.model_name :
248 # ccfg.s = ccfg.s * 3/4
249 ccfg.shard_recompute_input = ccfg.t
250 ccfg.shard_output_activ = ccfg.t
251 # ccfg.bytes_grad = 4
254def custom_cm(ccfg):
255 """llama moe"""
256 shard_p_os_exp = ccfg.shard_p_os_exp_partial
257 shard_p_os_non_exp_partial = math.gcd(ccfg.n_exp, ccfg.shard_p_os_non_exp)
258 shard_embed = ccfg.t
259 custom_deepseek3(ccfg)
261 def custom_shard(c):
262 c.shard_p_os_exp = shard_p_os_exp
263 c.shard_p_os_non_exp_partial = shard_p_os_non_exp_partial
264 c.shard_embed = shard_embed
266 for idx, f in enumerate(ccfg.layer_custom_config):
268 def wrap_hook(e, f=f):
269 if isinstance(e, CostModelConfig):
270 e = CWrap(e)
271 f[1](e)
272 e.set_ccfg(custom_shard)
274 ccfg.layer_custom_config[idx] = (f[0], wrap_hook)
276 def num_params_norm_cm(c, _):
277 return c.n_normOp * 2 * c.h + 0.5 * c.n_attMM * c.dh
279 ccfg.overwrite_eval_functions["num_params_norm"] = num_params_norm_cm
282def check_and_apply_custom_hook(e):
283 """routing hooks"""
284 if isinstance(e, CostModelConfig):
285 e = CWrap(e)
286 map_modelname_custom = {
287 "llama2": custom_llama2,
288 "mixtral": custom_mixtral,
289 "t5": custom_t5,
290 "pangualpha": custom_pangualpha,
291 "deepseek": custom_deepseek3,
292 "qwen": custom_qwen,
293 "cm": custom_cm,
294 }
295 for k, v in map_modelname_custom.items():
296 if k in e.get_model_name().lower():
297 e.set_ccfg(v)
298 return
299 logger.warning(
300 "Hook not defined for: %s. Default one is chosen", e.get_model_name()
301 )
302 e.set_ccfg(custom_default_transformer)