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