« 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"""parser child class"""
16import ast
17from pathlib import Path
18import importlib.util
19import time
20import sys
21import os
22from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config
23from hyper_parallel.auto_parallel.sapp_nd.nd.common.framework_parsers._cost_model_parser import _CostModelParser
24from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory
25from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
28class CostModelParserHyperparallel(_CostModelParser):
29 """parser class for HyperParallel format"""
31 def parse(self):
32 """Parse a HyperParallel TOML configuration."""
33 if self.ccfg.source_code:
34 specs = self.__parse_init_code()
35 self.ccfg.specs = specs
36 self.__parse_toml()
37 else:
38 now = time.time()
39 home_path = os.path.expanduser("~")
40 if home_path not in sys.path:
41 sys.path.append(home_path)
42 spec_torch = importlib.util.find_spec("torch")
43 spec_torchtitan = importlib.util.find_spec("torchtitan")
44 if spec_torch is not None and spec_torchtitan is not None:
45 # existing torchtitan package
46 spec_path = spec_torchtitan.submodule_search_locations[0]
47 if spec_path not in sys.path:
48 sys.path.append(spec_path)
49 try:
50 logger.info(
51 "found torchtitan package from homedir: %s",
52 spec_path,
53 )
54 logger.info("importing getter from torchitan...")
55 os.environ["ASCEND_SLOG_PRINT_TO_STDOUT"] = "0"
56 # pylint: disable=import-outside-toplevel
57 import torchtitan.protocols.train_spec as train_spec_module
58 logger.info(
59 "import time: %s sec",
60 round(time.time() - now, 2),
61 )
62 train_spec = train_spec_module.get_train_spec(
63 self.config.model.name
64 )
65 model_args = train_spec.model_args[self.config.model.flavor]
66 # Convert recursively to Config
67 # pprint.pprint(train_spec.model_args)
68 # pprint.pprint(model_args)
69 # model = train_spec.model_cls(model_args)
70 # print(model)
71 # print(dir(model))
72 # print("NAMED PARAMETERS")
73 # for k,_ in model.named_parameters():
74 # print(k)
75 specs = self.__obj_to_config(model_args)
76 self.ccfg.specs = specs
77 self.__parse_toml()
78 return
79 except ModuleNotFoundError:
80 pass
81 raise AttributeError(
82 "Hyperparallel config: Could not find torch/torchtitan package. "
83 "Please specify argument --code-path with __init__.py path "
84 )
86 def __obj_to_config(self, obj):
87 """convert to Config"""
88 if obj and not isinstance(obj, (str, int, float, bool, list)):
89 res = Config({})
90 for k, v in obj.__dict__.items():
91 setattr(res, k, self.__obj_to_config(v))
92 return res
93 return obj
95 def __parse_init_code(self):
96 """parse __init__.py through ast"""
97 def parse_args(res, arg):
98 """parse config from target flavor (var = dict: flavor -> config)"""
99 for kw in arg.keywords:
100 if not isinstance(kw.value, ast.Call):
101 res[kw.arg] = ast.literal_eval(kw.value)
102 else:
103 res[kw.arg] = {}
104 parse_args(res[kw.arg], kw.value)
106 path = Path(self.ccfg.source_code)
107 source_code = path.read_text(encoding="utf-8")
108 tree = ast.parse(source_code, filename=path.name)
109 # fetch get_spec() AST
110 tree_get_spec = next(
111 node
112 for node in ast.walk(tree)
113 if isinstance(node, ast.FunctionDef) and node.name == "get_train_spec"
114 )
115 # fetch model_args variable AST
116 var_model_args = next(
117 k.value.id
118 for k in tree_get_spec.body[0].value.keywords
119 if k.arg == "model_args"
120 )
121 var_tree = None
122 for node in ast.walk(tree):
123 if isinstance(node, ast.Assign) and node.targets[0].id == var_model_args:
124 var_tree = node
125 break
126 # fetch dict of hyperparameters according to flavor
127 params = {}
128 for k, v in zip(var_tree.value.keys, var_tree.value.values):
129 if k.value == self.config.model.flavor:
130 parse_args(params, v)
131 break
132 return Config(params)
134 def __parse_toml(self):
135 """main parsing order"""
136 self.ccfg.model_name = self.config.model.name
137 self.ccfg.config_format = "toml"
138 self.ccfg.multimodal = False
139 self.ccfg.device_capacity = Memory.from_string("56GB") # important
140 self.ccfg.mm_ccfgs = None
141 self.ccfg.mm_order = None
142 self.__parse_feature_flag()
143 self.__parse_hyperparam()
144 self.__parse_strat()
145 self.__parse_moe()
146 self.config_optimizer_shard(self.ccfg) # need to adapt FSDP
147 self.config_comm_flag(self.ccfg)
148 self.__parse_batch()
149 self.__init_shard()
150 self.__init_bytes()
151 self.ccfg.n_mtp = 0
152 self.ccfg.layer_custom_config = [(self.ccfg.n_lay, None)]
153 self.ccfg.overwrite_eval_functions = {}
155 def __parse_strat(self):
156 """strategy vars"""
157 self.ccfg.d = max(
158 1,
159 self.config.parallelism.data_parallel_replicate_degree
160 * self.config.parallelism.data_parallel_shard_degree,
161 ) # need correction
162 self.ccfg.t = max(1, self.config.parallelism.tensor_parallel_degree)
163 self.ccfg.p = max(1, self.config.parallelism.pipeline_parallel_degree)
164 self.ccfg.cp = max(1, self.config.parallelism.context_parallel_degree)
165 self.ccfg.ep = max(1, self.config.parallelism.expert_parallel_degree)
166 self.ccfg.sp = self.ccfg.t
167 self.ccfg.vp = 1
168 self.ccfg.op_weight_shard = (
169 self.config.parallelism.data_parallel_shard_degree * self.ccfg.t
170 )
171 self.ccfg.os_max_shard = (
172 self.ccfg.op_weight_shard if self.ccfg.op_weight_shard >= 1
173 else self.ccfg.d * self.ccfg.t
174 ) # need correction
175 self.ccfg.offset = 0 # important
176 self.ccfg.full_rec = self.config.activation_checkpoint.mode == "full"
177 self.ccfg.sel_rec = self.config.activation_checkpoint.mode == "selective"
178 self.ccfg.pp_sched = self.config.parallelism.pipeline_parallel_schedule
179 if self.ccfg.pp_sched:
180 # From Pytorch code
181 schedule_map = {
182 "1F1B": "1f1b",
183 "Interleaved1F1B": "1f1b",
184 "GPipe": "gpipe",
185 "FlexibleInterleaved1F1B": "1f1b", # ??
186 "LoopedBFS": "1f1b",
187 "InterleavedZeroBubble": "1f1b",
188 "ScheduleZBVZeroBubble": "zero_bubble_v",
189 "PipelineScheduleSingle": None,
190 "PipelineScheduleMulti": None,
191 }
192 if "Interleaved" in self.ccfg.pp_sched and self.ccfg.p > 1:
193 self.ccfg.vp = 2
194 if self.ccfg.pp_sched in schedule_map:
195 self.ccfg.pp_sched = schedule_map[self.ccfg.pp_sched]
196 else:
197 logger.warning(
198 "Unsupported pipeline_parallel_schedule '%s'. "
199 "Defaulting to 1f1b.",
200 self.ccfg.pp_sched,
201 )
202 self.ccfg.pp_sched = "1f1b"
203 else:
204 self.ccfg.pp_sched = "1f1b"
205 self.ccfg.emb_out_in_offset = True
206 self.ccfg.n_s_split = 1
207 self.ccfg.cp_algo = "colossalai_cp"
208 self.ccfg.rec_op = Config({
209 "attBMM": 1,
210 "headCast": 1,
211 "dropout": 1,
212 "softmax": 1,
213 "normOp": 1,
214 "gather": 1,
215 "ffAct": 1,
216 })
217 self.ccfg.pp_partition = None
219 def __parse_hyperparam(self):
220 """hyperparameter vars"""
221 self.ccfg.multiple_of = max(1, self.ccfg.specs.multiple_of)
222 self.ccfg.fdm = max(1, self.ccfg.specs.ffn_dim_multiplier)
223 self.ccfg.h = self.ccfg.specs.dim
224 self.ccfg.hff = self.ccfg.specs.inter_dim
225 if not self.ccfg.hff:
226 self.ccfg.hff = self.ccfg.specs.hidden_dim
227 if not self.ccfg.hff:
228 if "llama" in self.ccfg.model_name:
229 self.ccfg.hff = self.init_hff()
230 else:
231 self.ccfg.hff = self.ccfg.h
232 self.ccfg.v = self.ccfg.specs.vocab_size
233 self.ccfg.s = self.config.training.seq_len
234 self.ccfg.a = self.ccfg.specs.n_heads
235 self.ccfg.s_fa = (
236 (self.ccfg.s / self.ccfg.a) if self.ccfg.has_fa else self.ccfg.s
237 )
238 self.ccfg.n_lay = self.ccfg.specs.n_layers
239 self.ccfg.n_kv = self.ccfg.specs.n_kv_heads
240 if not self.ccfg.n_kv:
241 self.ccfg.n_kv = self.ccfg.a
242 self.ccfg.dh = self.ccfg.h / self.ccfg.a
243 self.ccfg.dc_kv = self.ccfg.specs.kv_lora_rank
244 self.ccfg.dc_q = self.ccfg.specs.q_lora_rank
245 self.ccfg.dhr = self.ccfg.specs.qk_rope_head_dim
246 self.ccfg.k_1st_dense = self.ccfg.specs.n_dense_layers
247 self.ccfg.is_mtp_in_offset = True
249 def __parse_moe(self):
250 """MoE vars"""
251 self.ccfg.hff_exp = (
252 self.ccfg.specs.moe_inter_dim if self.ccfg.specs.moe_inter_dim
253 else self.ccfg.hff
254 )
255 if (
256 not hasattr(self.ccfg.specs, "moe_enabled")
257 or self.ccfg.specs.moe_enabled
258 ):
259 if self.ccfg.specs.moe_args:
260 self.ccfg.n_exp = self.ccfg.specs.moe_args.num_experts
261 self.ccfg.n_chosen_exp = self.ccfg.specs.moe_args.top_k
262 self.ccfg.n_shared_exp = (
263 self.ccfg.specs.moe_args.num_shared_experts
264 )
265 else:
266 self.ccfg.n_exp = 1
267 self.ccfg.n_chosen_exp = 1
268 self.ccfg.n_shared_exp = 0
269 self.ccfg.cap_fact = 1 # Assuming
270 self.ccfg.etp = self.config.parallelism.expert_tensor_parallel_degree
271 self.config_dp_tp_exp(self.ccfg) # need verification in code
273 def __parse_feature_flag(self):
274 """training feature vars"""
275 self.ccfg.has_op = True # Assuming
276 self.ccfg.has_grad_shard = True # Assuming FSDP
277 self.ccfg.freeze = False
278 self.ccfg.has_fa = True # Assuming
279 self.ccfg.vp_less_mem = False
280 self.ccfg.has_clip = False
281 self.ccfg.gmm = True
282 self.ccfg.vocab_emb_dp = True
283 self.ccfg.tie_emb_out = self.ccfg.specs.enable_weight_tying
285 def __parse_batch(self):
286 """batch related vars"""
287 self.ccfg.b = self.config.training.local_batch_size
288 self.ccfg.m = self.ccfg.p
289 self.ccfg.gbs = self.ccfg.b * self.ccfg.d * self.ccfg.m
291 def __init_shard(self):
292 """sharding vars"""
293 self.ccfg.shard_embed = self.ccfg.t
294 self.ccfg.shard_output_activ = True
295 self.ccfg.shard_recompute_input = True
296 self.ccfg.is_shard_mtp_param = True
298 def __init_bytes(self):
299 """fp bytes vars"""
300 self.ccfg.bytes_p = 4
301 self.ccfg.bytes_compute = 2
302 self.ccfg.bytes_softmax = 4
303 self.ccfg.bytes_grad = 4
304 self.ccfg.bytes_os = 4
305 self.ccfg.bytes_norm = 4