« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 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"""HyperParallel native train.yaml parser (Hyper V2).
17Parses the HyperParallel YAML configuration format
18(``examples/qwen3_5_0_8b_base/train.yaml``) and populates
19a :class:`CostModelConfig` for memory estimation.
21Model hyperparameters are resolved by importing the model module's
22``_build_config`` function directly (by naming convention) rather than
23through the ``ModelSpec`` registry. This avoids coupling the
24memory-estimation pipeline to the model registration interface.
26Expected YAML structure::
28 model:
29 name: qwen3_5 # model spec key
30 config_overrides:
31 num_hidden_layers: 4
32 hidden_size: 3584
33 ...
35 train:
36 accelerator:
37 dp_shard: 4 # FSDP shard degree
38 gradient_checkpointing:
39 activation_checkpoint: none # none | full | selective
40 global_batch_size: 4
41 micro_batch_size: 1
43 data:
44 max_seq_len: 64
45"""
46# pylint: disable=too-many-locals,too-many-statements,too-many-branches
47import logging
48from typing import Any, Dict
50from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config, YamlObject
51from hyper_parallel.auto_parallel.sapp_nd.nd.common.framework_parsers._cost_model_parser import _CostModelParser
52from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory
54logger = logging.getLogger(__name__)
57class CostModelParserHyperV2(_CostModelParser):
58 """Parser for HyperParallel native train.yaml configuration format.
60 This parser replaces the placeholder ``CostModelParserHyperparallel``
61 which was written for an older TorchTitan TOML format. It reads the
62 current Hyper YAML schema (``model.name`` + ``train.accelerator``)
63 and resolves model parameters by importing ``_build_config`` directly
64 from the model module. When the model module is not available it
65 falls back to reading hyperparameters from ``config_overrides``.
66 """
68 def parse(self) -> None:
69 """Main parsing entry point."""
70 self.ccfg.config_format = "yaml"
71 self.ccfg.multimodal = False
72 self.ccfg.mm_ccfgs = None
73 self.ccfg.mm_order = None
75 # Resolve model hyperparameters via Hyper's own config pipeline
76 self._resolve_model_config_pipeline()
78 # --- Parallelism ---
79 self._parse_parallelism()
81 # --- Batch ---
82 self._parse_batch()
84 # --- Feature flags ---
85 self._parse_feature_flags()
87 # Flash attention factor
88 if self.ccfg.has_fa and self.ccfg.a > 0:
89 self.ccfg.s_fa = self.ccfg.s / self.ccfg.a
90 else:
91 self.ccfg.s_fa = self.ccfg.s
93 # --- Recompute ---
94 self._parse_recompute()
96 # --- n_s_split ---
97 self.ccfg.n_s_split = 1
99 # --- Bytes ---
100 self._init_bytes()
102 # --- Post-processing ---
103 self._init_moe_strategy()
104 self.config_optimizer_shard(self.ccfg)
105 self.config_comm_flag(self.ccfg)
106 self._init_shard()
107 self.ccfg.layer_custom_config = [(self.ccfg.n_lay + self.ccfg.n_mtp, None)]
108 self.ccfg.offset = 0
109 self.ccfg.overwrite_eval_functions = {}
111 def _resolve_model_config_pipeline(self):
112 """Resolve model hyperparameters to populate ``ccfg``.
114 Tries to import ``_build_config`` directly from the model module.
115 On failure, falls back to reading hyperparameters from
116 ``config_overrides``. This avoids coupling the ModelSpec
117 registration interface to the memory-estimation pipeline.
118 """
119 try:
120 self._resolve_via_direct_import()
121 except ValueError:
122 logger.debug(
123 "Direct _build_config import failed; "
124 "falling back to config_overrides"
125 )
126 self._resolve_from_config_overrides()
127 except TypeError as exc:
128 raise ValueError(
129 f"Failed to build HyperTrainerConfig: {exc}"
130 ) from exc
132 def _resolve_via_direct_import(self):
133 """Import ``_build_config`` directly from the model module.
135 Each model module may expose an internal ``_build_config(cfg)``
136 function (by naming convention) that returns a model-specific
137 Config object. The parser imports it directly rather than
138 going through the ``ModelSpec`` registry.
140 Raises:
141 ValueError: If the model module or ``_build_config`` is not found.
142 """
143 # pylint: disable=import-outside-toplevel
144 from hyper_parallel.trainer.config import (
145 _instantiate_recursive, HyperTrainerConfig,
146 )
148 config_dict = self._config_to_flat_dict(self.config)
149 trainer_cfg = _instantiate_recursive(HyperTrainerConfig, config_dict)
150 self.ccfg.model_name = trainer_cfg.model.name
152 import importlib # pylint: disable=import-outside-toplevel
153 try:
154 mod = importlib.import_module(
155 f"hyper_parallel.models.{trainer_cfg.model.name}")
156 except ImportError as exc:
157 raise ValueError(
158 f"Model module '{trainer_cfg.model.name}' not found: {exc}"
159 ) from exc
161 build_config_fn = getattr(mod, "_build_config", None)
162 if build_config_fn is None:
163 raise ValueError(
164 f"Model module '{trainer_cfg.model.name}' has no "
165 f"_build_config function. The parser falls back to "
166 f"config_overrides."
167 )
169 model_config = build_config_fn(trainer_cfg)
170 self._map_model_config_to_ccfg(model_config)
172 def _map_model_config_to_ccfg(self, model_config) -> None:
173 """Map model-specific Config object fields to ``ccfg``."""
174 self.ccfg.h = int(model_config.hidden_size)
175 self.ccfg.n_lay = int(model_config.num_hidden_layers)
176 self.ccfg.a = int(model_config.num_attention_heads)
177 self.ccfg.hff = int(model_config.intermediate_size)
178 self.ccfg.v = int(model_config.vocab_size)
179 self.ccfg.s = int(model_config.max_position_embeddings)
181 self.ccfg.n_kv = int(getattr(model_config, "num_key_value_heads", 0))
182 if not self.ccfg.n_kv:
183 self.ccfg.n_kv = self.ccfg.a
184 self.ccfg.dh = self.ccfg.h / self.ccfg.a
185 self.ccfg.dc_kv = int(getattr(model_config, "kv_lora_rank", 0))
186 self.ccfg.dc_q = int(getattr(model_config, "q_lora_rank", 0))
187 self.ccfg.dhr = int(getattr(model_config, "qk_rope_head_dim", 0))
189 self.ccfg.n_exp = 1
190 self.ccfg.n_chosen_exp = 1
191 self.ccfg.n_shared_exp = 0
192 self.ccfg.hff_exp = self.ccfg.hff
193 self.ccfg.cap_fact = 1
194 self.ccfg.t_exp = self.ccfg.t
195 self.ccfg.d_exp = self.ccfg.d
196 self.ccfg.gmm = False
197 self.ccfg.k_1st_dense = 0
198 num_exp = int(getattr(model_config, "num_experts", 1))
199 if num_exp > 1:
200 self.ccfg.n_exp = max(1, num_exp)
201 self.ccfg.n_chosen_exp = max(
202 1, int(getattr(model_config, "num_experts_per_tok", 1)))
203 self.ccfg.n_shared_exp = int(
204 getattr(model_config, "n_shared_experts",
205 getattr(model_config, "num_shared_experts", 0)))
206 moe_inter = int(getattr(model_config, "moe_intermediate_size", 0))
207 if moe_inter:
208 self.ccfg.hff_exp = moe_inter
209 self.ccfg.k_1st_dense = int(
210 getattr(model_config, "first_k_dense_replace", 0))
211 self.ccfg.gmm = True
213 self.ccfg.n_mtp = int(getattr(model_config, "mtp_depth", 0))
214 self.ccfg.is_mtp_in_offset = bool(self.ccfg.n_mtp)
215 self.ccfg.multiple_of = int(getattr(model_config, "multiple_of", 256))
216 self.ccfg.fdm = float(getattr(model_config, "ffn_dim_multiplier", 1.0))
218 self._resolve_device_capacity()
220 def _resolve_from_config_overrides(self) -> None:
221 """Populate ``ccfg`` directly from ``config_overrides``.
223 Used when the model is not registered in ``ModelSpec``.
224 """
225 model_raw = self._get_cfg_attr(self.config, "model", Config({}))
226 overrides = self._get_cfg_attr(model_raw, "config_overrides", Config({}))
227 data_raw = self._get_cfg_attr(self.config, "data", Config({}))
229 self.ccfg.model_name = str(
230 self._get_cfg_attr(model_raw, "name", "custom"))
231 self.ccfg.h = int(self._get_cfg_attr(overrides, "hidden_size", 0))
232 self.ccfg.n_lay = int(self._get_cfg_attr(overrides, "num_hidden_layers", 0))
233 self.ccfg.a = int(self._get_cfg_attr(overrides, "num_attention_heads", 0))
234 self.ccfg.hff = int(self._get_cfg_attr(overrides, "intermediate_size", 0))
235 self.ccfg.v = int(self._get_cfg_attr(overrides, "vocab_size", 0))
237 # seq_len: data.max_seq_len > overrides > default
238 self.ccfg.s = int(
239 self._get_cfg_attr(data_raw, "max_seq_len", 0)
240 or self._get_cfg_attr(overrides, "max_position_embeddings", 0)
241 or self._get_cfg_attr(overrides, "seq_length", 0)
242 or 4096
243 )
245 self.ccfg.n_kv = int(
246 self._get_cfg_attr(overrides, "num_key_value_heads", 0))
247 if not self.ccfg.n_kv:
248 self.ccfg.n_kv = self.ccfg.a
249 self.ccfg.dh = self.ccfg.h / self.ccfg.a if self.ccfg.a else 0
250 self.ccfg.dc_kv = int(
251 self._get_cfg_attr(overrides, "kv_lora_rank", 0))
252 self.ccfg.dc_q = int(
253 self._get_cfg_attr(overrides, "q_lora_rank", 0))
254 self.ccfg.dhr = int(
255 self._get_cfg_attr(overrides, "qk_rope_head_dim", 0))
257 # MoE
258 self.ccfg.n_exp = 1
259 self.ccfg.n_chosen_exp = 1
260 self.ccfg.n_shared_exp = 0
261 self.ccfg.hff_exp = self.ccfg.hff
262 self.ccfg.cap_fact = 1
263 self.ccfg.t_exp = self.ccfg.t
264 self.ccfg.d_exp = self.ccfg.d
265 self.ccfg.gmm = False
266 self.ccfg.k_1st_dense = 0
267 num_exp = int(self._get_cfg_attr(overrides, "num_experts", 1))
268 if num_exp > 1:
269 self.ccfg.n_exp = max(1, num_exp)
270 self.ccfg.n_chosen_exp = max(
271 1, int(self._get_cfg_attr(overrides, "num_experts_per_tok", 1)))
272 self.ccfg.n_shared_exp = int(
273 self._get_cfg_attr(overrides, "num_shared_experts", 0))
274 moe_inter = int(
275 self._get_cfg_attr(overrides, "moe_intermediate_size", 0))
276 if moe_inter:
277 self.ccfg.hff_exp = moe_inter
278 self.ccfg.k_1st_dense = int(
279 self._get_cfg_attr(overrides, "first_k_dense_replace", 0))
280 self.ccfg.gmm = True
282 self.ccfg.n_mtp = int(self._get_cfg_attr(overrides, "mtp_depth", 0))
283 self.ccfg.is_mtp_in_offset = bool(self.ccfg.n_mtp)
284 self.ccfg.multiple_of = int(
285 self._get_cfg_attr(overrides, "multiple_of", 256))
286 self.ccfg.fdm = float(
287 self._get_cfg_attr(overrides, "ffn_dim_multiplier", 1.0))
289 self._resolve_device_capacity()
291 def _resolve_device_capacity(self) -> None:
292 """Set device capacity from config or default (64 GB)."""
293 ctx = self._get_cfg_attr(self.config, "context", Config({}))
294 device_mem_str = ctx.__dict__.get("max_device_memory", None) if isinstance(ctx, (Config, YamlObject)) else None
295 if device_mem_str:
296 self.ccfg.device_capacity = Memory.from_string(str(device_mem_str))
297 else:
298 self.ccfg.device_capacity = Memory.from_string("64GB")
300 # ── Private helpers ───────────────────────────────────────────────
302 @staticmethod
303 def _get_cfg_attr(cfg: Any, attr: str, default: Any = None) -> Any:
304 """Get an attribute from ``Config`` / ``YamlObject`` safely.
306 ``YamlObject.__getattr__`` returns ``0`` for missing attributes
307 instead of raising ``AttributeError``, which breaks Python's
308 ``getattr(obj, attr, default)`` fallback protocol. This helper
309 checks ``__dict__`` directly.
310 """
311 if isinstance(cfg, (Config, YamlObject)):
312 return cfg.__dict__.get(attr, default)
313 return getattr(cfg, attr, default)
315 @staticmethod
316 def _config_to_flat_dict(cfg: Any) -> Dict[str, Any]:
317 """Recursively convert a ``Config`` or ``YamlObject`` to a flat dict."""
318 if isinstance(cfg, (Config, YamlObject)):
319 return {k: CostModelParserHyperV2._config_to_flat_dict(v)
320 for k, v in cfg.__dict__.items()
321 if not k.startswith("_")}
322 if isinstance(cfg, (int, float, str, bool)):
323 return cfg # type: ignore[return-value]
324 if isinstance(cfg, list):
325 return [CostModelParserHyperV2._config_to_flat_dict(i) for i in cfg]
326 return cfg
328 @staticmethod
329 def _bytes_from_dtype(dtype_str: Any) -> int:
330 """Parse a dtype string (e.g. ``\"float32\"``) to byte size.
332 Returns ``4`` for float32, ``2`` for bfloat16/float16, etc.
333 Defaults to ``4`` when parsing fails.
334 """
335 import re # pylint: disable=import-outside-toplevel
336 dtype_str = str(dtype_str)
337 m = re.search(r"(\d+)", dtype_str)
338 if m:
339 return max(2, int(m.group(1)) // 8)
340 return 4
342 def _parse_parallelism(self):
343 """Extract parallelism settings from ``train.accelerator``."""
344 train_raw = self._get_cfg_attr(self.config, "train", Config({}))
345 accel = self._get_cfg_attr(train_raw, "accelerator", Config({}))
347 dp_shard = int(self._get_cfg_attr(accel, "dp_shard", 1) or 1)
348 dp_replicate = int(self._get_cfg_attr(accel, "dp_replicate", 1) or 1)
349 tp = int(self._get_cfg_attr(accel, "tp_degree", 1) or 1)
350 pp = int(self._get_cfg_attr(accel, "pipeline_parallel_degree", 1) or 1)
351 cp = int(self._get_cfg_attr(accel, "context_parallel_degree", 1) or 1)
352 ep = int(self._get_cfg_attr(accel, "expert_parallel_degree", 1) or 1)
353 etp = int(self._get_cfg_attr(accel, "expert_tensor_parallel_degree", 0) or 0)
354 ep = max(ep, 1)
356 # FSDP: d = replicate * shard (when shard > 1)
357 self.ccfg.d = max(1, dp_replicate * dp_shard)
358 self.ccfg.t = max(1, tp)
359 self.ccfg.p = max(1, pp)
360 self.ccfg.cp = max(1, cp)
361 self.ccfg.ep = max(1, ep)
362 self.ccfg.sp = self.ccfg.t # Sequence parallel factor
363 self.ccfg.etp = etp
364 self.ccfg.vp = max(1, int(
365 self._get_cfg_attr(accel, "pp_interleave_num", 1) or 1
366 ))
367 use_sp = bool(self._get_cfg_attr(accel, "use_seq_parallel", False))
368 self.ccfg.sp = self.ccfg.t if use_sp else 1
369 self.ccfg.pp_sched = str(
370 self._get_cfg_attr(accel, "pipeline_scheduler", "1f1b")
371 )
373 # Optimizer parallel sharding
374 self.ccfg.has_op = bool(self._get_cfg_attr(accel,
375 "enable_parallel_optimizer",
376 True))
377 self.ccfg.op_weight_shard = max(1, int(
378 self._get_cfg_attr(accel, "optimizer_weight_shard_size", 0)
379 ) or (self.ccfg.d * self.ccfg.t))
380 self.ccfg.has_grad_shard = bool(self._get_cfg_attr(accel,
381 "gradient_accumulation_shard",
382 False))
383 self.ccfg.os_max_shard = (
384 self.ccfg.op_weight_shard if self.ccfg.op_weight_shard >= 1
385 else self.ccfg.d * self.ccfg.t
386 )
388 def _parse_batch(self):
389 """Extract batch settings from ``train`` section."""
390 train_raw = self._get_cfg_attr(self.config, "train", Config({}))
391 self.ccfg.b = max(1, int(self._get_cfg_attr(train_raw, "micro_batch_size", 1) or 1))
392 m = int(self._get_cfg_attr(train_raw, "micro_batch_num", 0) or 0)
393 if m > 0:
394 self.ccfg.m = m
395 else:
396 self.ccfg.m = self.ccfg.p
397 gbs = int(self._get_cfg_attr(train_raw, "global_batch_size", 0) or 0)
398 if gbs > 0:
399 self.ccfg.gbs = gbs
400 else:
401 self.ccfg.gbs = self.ccfg.b * self.ccfg.d * self.ccfg.m
403 def _parse_feature_flags(self):
404 """Set training feature flags."""
405 self.ccfg.has_fa = True
406 self.ccfg.vocab_emb_dp = True
407 self.ccfg.tie_emb_out = False
408 self.ccfg.freeze = False
409 train_raw = self._get_cfg_attr(self.config, "train", Config({}))
410 optimizer = self._get_cfg_attr(train_raw, "optimizer", Config({}))
411 max_grad_norm = float(
412 self._get_cfg_attr(optimizer, "max_grad_norm", 0.0) or 0.0
413 )
414 self.ccfg.has_clip = max_grad_norm > 0
415 self.ccfg.vp_less_mem = False
416 self.ccfg.cp_algo = "colossalai_cp"
418 def _parse_recompute(self):
419 """Parse recompute mode from ``train.gradient_checkpointing``."""
420 train_raw = self._get_cfg_attr(self.config, "train", Config({}))
421 gc = self._get_cfg_attr(train_raw, "gradient_checkpointing", Config({}))
422 ac_mode = str(self._get_cfg_attr(gc, "activation_checkpoint", "none"))
423 self.ccfg.full_rec = ac_mode == "full"
424 self.ccfg.sel_rec = ac_mode == "selective"
425 self.ccfg.rec_op = Config({
426 "attBMM": 1,
427 "headCast": 1,
428 "dropout": 1,
429 "softmax": 1,
430 "normOp": 1,
431 "gather": 1,
432 "ffAct": 1,
433 })
435 def _init_bytes(self):
436 """Set FP byte sizes from dtype fields in the model section."""
437 model_raw = self._get_cfg_attr(self.config, "model", Config({}))
438 self.ccfg.bytes_p = self._bytes_from_dtype(
439 self._get_cfg_attr(model_raw, "param_init_type", "float32"))
440 self.ccfg.bytes_compute = self._bytes_from_dtype(
441 self._get_cfg_attr(model_raw, "compute_dtype", "bfloat16"))
442 self.ccfg.bytes_softmax = self._bytes_from_dtype(
443 self._get_cfg_attr(model_raw, "softmax_compute_type", "float32"))
444 self.ccfg.bytes_grad = 4
445 self.ccfg.bytes_os = 4
446 self.ccfg.bytes_norm = 4
448 def _init_moe_strategy(self):
449 """Initialize MoE strategy variables via base helper."""
450 self.config_dp_tp_exp(self.ccfg)
452 def _init_shard(self):
453 """Initialize sharding variables.
455 Note: ``shard_output_activ`` and ``shard_recompute_input`` are set to
456 ``ccfg.t`` directly instead of ``True`` (which would evaluate to 1).
457 This ensures activation memory is correctly divided by the tensor
458 parallel degree, matching the behavior of the MF parser. The
459 ``custom_qwen`` arch hook (in ``arch_hooks.py``) would also set these
460 to ``ccfg.t`` via ``check_and_apply_custom_hook``, but relying on the
461 hook is fragile — it may fail silently when the model name convention
462 changes or when the ``set_ccfg`` mechanism is bypassed.
463 """
464 self.ccfg.shard_embed = self.ccfg.t * self.ccfg.d
465 self.ccfg.shard_output_activ = self.ccfg.t
466 self.ccfg.shard_recompute_input = self.ccfg.t
467 self.ccfg.is_shard_mtp_param = True