Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / config.py: 99%
341 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-22 04:23 +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"""Training configuration schema — strict three-tier (model/data/train).
17Top-level keys are exactly ``model``, ``data`` and ``train`` and nothing else;
18anything outside this three-tier schema is rejected by the parser.
20YAML shape::
22 model:
23 name: qwen3_5
24 weights_path: /path/to/weights
25 data:
26 type: hf_datasets
27 train_path: /path/to/data
28 train:
29 max_steps: 100
30 micro_batch_size: 1
31 global_batch_size: 8
32 seed: 42
33 init_device: meta
34 optimizer:
35 type: adamw
36 lr: 1.0e-4
37 accelerator:
38 dp_shard: 8
39 tp: 1
40 mixed_precision:
41 enabled: true
42 param_dtype: bfloat16
43 checkpoint:
44 output_dir: outputs/run1
45 ...
46"""
47import argparse
48import difflib
49import logging
50import os
51from dataclasses import dataclass, field, fields, is_dataclass
52from typing import Any, Dict, List, Optional, Type, TypeVar, Union, get_args, get_origin
54import yaml
56logger = logging.getLogger(__name__)
58T = TypeVar("T")
60_BOOL_TRUE_STRINGS = frozenset(("true", "yes", "y", "on", "1", "t"))
61_BOOL_FALSE_STRINGS = frozenset(("false", "no", "n", "off", "0", "f"))
63# ============================================================================
64# model:
65# ============================================================================
68@dataclass
69class ModelConfig:
70 """``model.*`` — model identity, weights, and architecture overrides.
72 Only universal transformer fields are typed here. Anything model-
73 specific (mRoPE section split, MoE expert geometry, linear-attention
74 head counts, ``layer_types`` ...) goes through ``config_overrides`` —
75 a free-form ``dict`` that is merged into the underlying model
76 constructor by the model's ``build_model_fn``.
77 """
78 name: str = "qwen3_5"
79 weights_path: Optional[str] = None
80 tokenizer_path: Optional[str] = None
81 freeze_modules: Optional[list] = None
82 tp_plan: Optional[dict] = None
83 cp_modules: Optional[list] = None
84 ep_modules: Optional[list] = None
85 # Visual Encoder local DP/CP overrides for VL models.
86 vision_parallel: Optional[dict] = None
87 # Universal transformer architecture overrides.
88 num_hidden_layers: Optional[int] = None
89 hidden_size: Optional[int] = None
90 intermediate_size: Optional[int] = None
91 num_attention_heads: Optional[int] = None
92 num_key_value_heads: Optional[int] = None
93 vocab_size: Optional[int] = None
94 max_position_embeddings: Optional[int] = None
95 # Free-form per-model overrides handed to ``build_model_fn``.
96 config_overrides: Optional[dict] = None
99def _coerce_vision_parallel_bool(value: Any) -> bool:
100 """Normalize free-form ``model.vision_parallel`` bool fields."""
101 if isinstance(value, bool):
102 return value
103 if isinstance(value, int) and value in (0, 1):
104 return bool(value)
105 return _string_to_bool(value)
108def get_vision_parallel_config(model_cfg: ModelConfig) -> Dict[str, Any]:
109 """Return visual Encoder local parallel config from ``model.*``.
111 ``model.vision_parallel`` is the preferred schema field. The earlier
112 ``model.config_overrides.vision_parallel`` location remains supported for
113 compatibility with validation configs. Bool-like fields are normalized here
114 so trainer/model code can consume a parsed config instead of re-parsing YAML
115 strings at each use site.
116 """
117 vision_parallel_cfg = getattr(model_cfg, "vision_parallel", None)
118 if isinstance(vision_parallel_cfg, dict):
119 vision_parallel = vision_parallel_cfg
120 else:
121 config_overrides = getattr(model_cfg, "config_overrides", None)
122 vision_parallel = {}
123 if isinstance(config_overrides, dict):
124 legacy_vision_parallel = config_overrides.get("vision_parallel", {})
125 if isinstance(legacy_vision_parallel, dict):
126 vision_parallel = legacy_vision_parallel
128 normalized = dict(vision_parallel)
129 for key in ("reuse_dp_shard_mesh", "share_samples_across_dp", "async_cp"):
130 if key in normalized:
131 normalized[key] = _coerce_vision_parallel_bool(normalized[key])
132 return normalized
134# ============================================================================
135# data:
136# ============================================================================
139@dataclass
140class DataConfig:
141 """``data.*`` — dataset, tokenizer/processor, sampler, batch shape.
143 - ``streaming``: only ``False`` supported today.
144 - ``num_workers``: keep ≥ 2 for real datasets.
145 - ``shuffle``: when ``False``, sampler reads samples in dataset order.
146 """
147 type: str = "dummy"
148 train_path: Optional[Any] = None
149 subset: Optional[str] = None
150 max_seq_len: int = 2048
151 text_key: str = "text"
152 train_size: Optional[int] = None
153 # multimodal / VL (synthetic vl_dummy path)
154 template: str = "empty"
155 image_key: str = "image"
156 messages_key: str = "messages"
157 image_token_id: int = 151655
158 video_token_id: int = 151656
159 vl_video: bool = False
160 vl_grid_t: int = 2
161 vl_grid_h: int = 2
162 vl_grid_w: int = 2
163 # loader perf
164 streaming: bool = False
165 num_workers: int = 0
166 prefetch_factor: Optional[int] = None
167 pin_memory: bool = True
168 shuffle: bool = True
169 # Megatron .bin/.idx options (data.type='megatron'). ``train_path`` is a
170 # single path-prefix, or a blend spec ("w1 prefix1 w2 prefix2 ...") /
171 # list-of-pairs. ``megatron_seed`` defaults to ``train.seed`` when None.
172 megatron_seed: Optional[int] = None
173 pad_token_id: int = 0
174 eod_token_id: Optional[int] = None
175 eod_mask_loss: bool = False
176 mmap_bin_files: bool = True
178# ============================================================================
179# train.* — sub-configs
180# ============================================================================
183@dataclass
184class AcceleratorConfig:
185 """``train.accelerator.*`` — parallelism topology.
187 Two ways to express data parallelism:
189 - **Legacy single field** (back-compat): ``dp`` only — maps to
190 ``dp_shard`` for FSDP.
191 - ** split**: ``dp_replicate`` (HSDP outer) and
192 ``dp_shard`` (FSDP inner). Pass ``dp_shard=-1`` to auto-fill from
193 ``world_size / (dp_replicate * cp * tp * pp)``.
195 For MoE: ``etp`` controls expert TP. Must equal ``tp`` or ``1``.
196 ``moe_token_dispatcher_type`` selects the EP token exchange algorithm.
197 ``npu_nums_per_device`` is the inner-EP degree for the deredundency
198 dispatcher; ``oep`` is inferred as ``ep // npu_nums_per_device``.
199 """
200 dp: Optional[int] = None
201 dp_replicate: int = 1
202 dp_shard: Optional[int] = None
203 tp: int = 1
204 cp: int = 1
205 pp: int = 1
206 # Number of pipeline micro-batches per optimizer step when ``pp > 1``.
207 # The global batch is split into this many micro-batches along dim 0 and
208 # streamed through the pipeline stages; ``global_batch_size`` must be
209 # divisible by it. Defaults to 1 (no micro-batching).
210 pp_micro_batch_num: int = 1
211 # Pipeline schedule when ``pp > 1``: ``"gpipe"`` (all-forward then
212 # all-backward) or ``"1f1b"`` (one-forward-one-backward steady state).
213 # ``None`` keeps each model's default (dense → gpipe, MoE → 1f1b).
214 pp_schedule: Optional[str] = None
215 # Virtual-pipeline (VPP) degree: number of non-contiguous stage chunks each
216 # PP rank owns. ``1`` (default) is the plain single-stage-per-rank pipeline.
217 # ``>1`` builds ``pp * pp_vpp`` interleaved global stages (rank ``r`` owns
218 # stages ``r, r+pp, r+2*pp, ...``) and drives them with interleaved 1F1B.
219 pp_vpp: int = 1
220 # Per-global-stage decoder-layer counts (length ``pp * pp_vpp``, summing to
221 # ``num_hidden_layers``). ``None`` (default) keeps the even split with the
222 # remainder on the later stages. Lets users rebalance stages whose extra
223 # modules (embed / visual tower / lm_head) dominate memory or latency.
224 pp_layer_split: Optional[List[int]] = None
225 ep: int = 1
226 etp: int = 1
227 moe_token_dispatcher_type: str = "all_to_all"
228 npu_nums_per_device: int = 8
229 zero_stage: int = 0
230 reshard_after_forward: bool = True
231 async_cp: bool = False
232 ulysses_degree: Optional[int] = None
233 # Qwen3.5 linear-attention CP execution and local GDN implementation.
234 linear_attention_cp_mode: str = "ulysses"
235 linear_attention_gdn_backend: str = "eager"
236 # Bucketed reduce-scatter: single fused RS per FSDP unit, stable fp32
237 # reduction order across runs.
238 comm_fusion: bool = True
239 # Offload sharded params, grads, and optimizer states to CPU so large
240 # checkpoints can fit without device-resident master weights and Adam state.
241 cpu_offload: bool = False
244@dataclass
245class MixedPrecisionConfig:
246 """``train.mixed_precision.*`` — mixed-precision forward configuration.
248 FSDP2 ``MixedPrecisionPolicy``: ``param_dtype`` is the all-gather'd
249 forward dtype, ``reduce_dtype`` is the reduce-scatter dtype for grads,
250 and ``output_dtype`` controls the forward-output dtype at FSDP wrap
251 boundaries (leave ``None`` to inherit from ``param_dtype``).
252 """
253 enabled: bool = False
254 param_dtype: str = "bfloat16"
255 reduce_dtype: str = "float32"
256 output_dtype: Optional[str] = None
259@dataclass
260class GradientCheckpointingConfig:
261 """``train.gradient_checkpointing.*`` — activation recomputation.
263 . Modes: ``"off"``, ``"full"``, or ``"selective"``.
264 """
265 activation_checkpoint: str = "off"
268@dataclass
269class OptimizerConfig:
270 """``train.optimizer.*`` — optimizer + LR schedule + grad clip.
272 ``loss_aggregation``: how the per-micro-batch loss is scaled before
273 backward. ``"token_weighted"`` divides the summed loss by the global
274 valid-token count; ``"rank_average"`` averages per-rank micro-batch
275 means and is preferred when batches have variable valid-token counts
276 across ranks.
277 ``max_grad_norm``: values <= 0 disable gradient clipping.
278 """
279 type: str = "adamw"
280 lr: float = 1e-4
281 lr_min: float = 1e-5
282 lr_decay_style: str = "cosine"
283 lr_warmup_ratio: float = 0.1
284 loss_aggregation: str = "token_weighted"
285 weight_decay: float = 0.01
286 max_grad_norm: float = 1.0
287 bsz_warmup_ratio: float = 0.0
288 eps: float = 1e-8
289 betas: tuple = (0.9, 0.999)
290 # ``None`` lets torch pick the foreach kernel; set ``False`` in YAML when a
291 # run must force the deterministic per-parameter loop.
292 foreach: Optional[bool] = None
295@dataclass
296class CheckpointConfig:
297 """``train.checkpoint.*`` — DCP save / load + HF export."""
298 output_dir: str = "outputs"
299 save_steps: int = 500
300 save_hf_weights: bool = True
301 load_path: Optional[str] = None
302 save_async: bool = False
305@dataclass
306class LoggingConfig:
307 """``train.logging.*`` — console / metric output (consumed by LoggingCallback)."""
308 report_to: str = "none"
309 report_global_loss: bool = False
310 log_steps: int = 10
311 report_throughput: bool = True
312 model_flops_per_token: Optional[int] = None
313 peak_tflops: Optional[float] = None # e.g. 312.0 for A100 bf16
316@dataclass
317class TensorBoardConfig:
318 """``train.tensorboard.*`` — TB SummaryWriter on rank 0."""
319 enabled: bool = False
320 output_dir: str = "tb_traces"
321 log_steps: int = 1
324@dataclass
325class WandbConfig:
326 """``train.wandb.*`` — W&B run logging on rank 0."""
327 enabled: bool = False
328 project: str = "hyper-parallel"
329 run_name: Optional[str] = None
330 log_steps: int = 1
333@dataclass
334class ProfileConfig:
335 """``train.profile.*`` — torch.profiler schedule ().
337 Schedule semantics: wait → warmup → active."""
338 enabled: bool = False
339 output_dir: str = "profiler_traces"
340 wait_steps: int = 1
341 warmup_steps: int = 1
342 active_steps: int = 3
345@dataclass
346class MemoryMonitorConfig:
347 """``train.memory_monitor.*`` — periodic device-memory snapshot."""
348 enabled: bool = False
349 log_steps: int = 50
350 reset_peak_each_step: bool = False
353@dataclass
354class MonitorConfig:
355 """``train.monitor.*`` — training-state monitor for loss / gradient scalars."""
356 monitor_on: bool = False
357 dump_path: str = "./dump"
358 target: Optional[list] = None
359 invert: bool = False
360 step_interval: int = 1
361 local_loss_format: Optional[list] = None
362 device_local_loss_format: Optional[list] = None
363 local_norm_format: Optional[list] = None
364 device_local_norm_format: Optional[list] = None
367@dataclass
368class MoEMonitorConfig:
369 """``train.moe_monitor.*`` — MoE routing / load-balance monitor.
371 When enabled, :class:`~hyper_parallel.core.moe_utils.MoEMonitorCallback`
372 automatically syncs ``tokens_per_expert`` across distributed ranks and
373 updates ``expert_bias`` after each optimizer step. The mean ``aux_loss``
374 across MoE layers is exposed via ``last_mean_aux_loss`` so that
375 :class:`LoggingCallback` can print it alongside the main training loss.
377 DP/TP+SP/CP group information is automatically obtained from the trainer's
378 device mesh — no manual group configuration needed.
380 Args:
381 enabled: Whether to activate the MoE monitor callback.
382 lr: Step size for expert bias updates. Defaults to ``1e-3``.
383 num_recomputations: Number of forward executions per optimizer step.
384 Default ``1``. Set to ``2`` when activation checkpoint is enabled.
385 """
386 enabled: bool = False
387 lr: float = 1e-3
388 num_recomputations: int = 1
391@dataclass
392class EvalConfig:
393 """``train.eval.*`` — eval cadence + dataset."""
394 eval_steps: int = 0
395 eval_dataset: Optional[str] = None
398@dataclass
399class DebugConfig:
400 """``train.debug.*`` — reproducibility and numerical-stability knobs.
402 All flags here are production-safe; they tune determinism (CI / paper
403 reproducibility), guard against numerical blow-ups, and bound memory
404 growth in long runs.
405 """
406 deterministic: bool = False
407 deterministic_warn_only: bool = False
408 check_nan_inf: bool = False
409 gc_steps: int = 0
411# ============================================================================
412# train: (top of the train section, holds the sub-configs)
413# ============================================================================
416@dataclass
417class TrainConfig:
418 """``train.*`` — full training-section config.
420 Flat fields cover the basic loop knobs (steps, batch shape, init device,
421 seed, comm backend); nested sub-configs cover everything else.
422 """
423 # Loop shape
424 max_steps: int = 100
425 num_train_epochs: int = 1
426 global_batch_size: int = 8
427 micro_batch_size: int = 1
428 seed: int = 42
430 # Runtime / device
431 backend: str = "torch"
432 init_device: str = "meta"
433 comm_backend: Optional[str] = None
434 local_rank: int = 0 # set from LOCAL_RANK env by parser
436 # Sub-configs
437 accelerator: AcceleratorConfig = field(default_factory=AcceleratorConfig)
438 mixed_precision: MixedPrecisionConfig = field(default_factory=MixedPrecisionConfig)
439 gradient_checkpointing: GradientCheckpointingConfig = field(
440 default_factory=GradientCheckpointingConfig
441 )
442 optimizer: OptimizerConfig = field(default_factory=OptimizerConfig)
443 checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig)
444 logging: LoggingConfig = field(default_factory=LoggingConfig)
445 tensorboard: TensorBoardConfig = field(default_factory=TensorBoardConfig)
446 wandb: WandbConfig = field(default_factory=WandbConfig)
447 profile: ProfileConfig = field(default_factory=ProfileConfig)
448 memory_monitor: MemoryMonitorConfig = field(default_factory=MemoryMonitorConfig)
449 monitor: MonitorConfig = field(default_factory=MonitorConfig)
450 moe_monitor: MoEMonitorConfig = field(default_factory=MoEMonitorConfig)
451 eval: EvalConfig = field(default_factory=EvalConfig)
452 debug: DebugConfig = field(default_factory=DebugConfig)
454# ============================================================================
455# Top-level: model / data / train (and only these three)
456# ============================================================================
459@dataclass
460class HyperTrainerConfig:
461 """Top-level config — strict three-tier ().
463 Allowed top-level keys: ``model``, ``data``, ``train``. Anything else in
464 the YAML is rejected by the parser with a typo-suggestion message.
465 """
466 model: ModelConfig = field(default_factory=ModelConfig)
467 data: DataConfig = field(default_factory=DataConfig)
468 train: TrainConfig = field(default_factory=TrainConfig)
470 # Computed (no user input)
471 train_steps: int = 0
473 def __post_init__(self):
474 self.train_steps = self.train.max_steps
476# ==============================================================================
477# CLI / YAML parser
478# ==============================================================================
479# Configuration parser: YAML file + CLI dot-path overrides.
480#
481# Supports:
482# - Unknown YAML/CLI keys emit a warning with difflib closest-match suggestions.
483# - Bool fields accept string aliases: ``true/yes/y/on/1/t`` -> ``True``,
484# ``false/no/n/off/0/f`` -> ``False``. Only applied when the dataclass
485# field type resolves to ``bool`` or ``Optional[bool]`` to avoid ambiguity.
488def _string_to_bool(value: Any) -> bool:
489 """Convert common string representations of booleans to ``bool``.
491 Accepts: ``true/yes/y/on/1/t`` → ``True``,
492 ``false/no/n/off/0/f`` → ``False``.
494 Args:
495 value: A string or bool value.
497 Returns:
498 The corresponding ``bool``.
500 Raises:
501 ValueError: When the string cannot be mapped to a bool.
502 """
503 if isinstance(value, bool):
504 return value
505 if isinstance(value, str):
506 lower = value.lower()
507 if lower in _BOOL_TRUE_STRINGS:
508 return True
509 if lower in _BOOL_FALSE_STRINGS:
510 return False
511 raise ValueError(
512 f"Cannot convert {value!r} to bool. "
513 "Expected one of: true/false/yes/no/y/n/on/off/1/0/t/f"
514 )
517def _resolve_field_type(cls: Type, dot_path: str) -> Optional[Type]:
518 """Walk a dataclass hierarchy to find the resolved type of a dot-path field.
520 Args:
521 cls: Root dataclass class.
522 dot_path: Dot-separated field path, e.g. ``"debug.deterministic"``.
524 Returns:
525 The resolved Python type, or ``None`` if the path cannot be resolved.
526 """
527 parts = dot_path.split(".")
528 current_cls = cls
529 for part in parts:
530 if not is_dataclass(current_cls):
531 return None
532 found = None
533 for f in fields(current_cls):
534 if f.name == part:
535 found = f
536 break
537 if found is None:
538 return None
539 field_type = found.type
540 # Unwrap Optional[X] → X
541 origin = get_origin(field_type)
542 if origin is Union:
543 unwrapped = [
544 a for a in get_args(field_type)
545 if a is not type(None) # pylint: disable=unidiomatic-typecheck
546 ]
547 field_type = unwrapped[0] if len(unwrapped) == 1 else field_type
548 current_cls = field_type
549 return current_cls
552def _coerce_cli_value(raw: str, dot_path: str, root_class: Type) -> Any:
553 """Parse a CLI string value, coercing to the correct type for the field.
555 Bool fields accept an extended string set. For all other
556 fields the existing int → float → str heuristic is used.
558 Args:
559 raw: Raw string from the CLI.
560 dot_path: Dot-separated field path used for type lookup.
561 root_class: Root dataclass class for type resolution.
563 Returns:
564 Coerced value.
565 """
566 field_type = _resolve_field_type(root_class, dot_path)
567 if field_type is bool:
568 try:
569 return _string_to_bool(raw)
570 except ValueError:
571 pass # fall through to heuristic below
572 # Existing heuristic: int → float → bool-string → str
573 try:
574 return int(raw)
575 except ValueError:
576 pass
577 try:
578 return float(raw)
579 except ValueError:
580 pass
581 if raw.lower() in ("true", "false"):
582 return raw.lower() == "true"
583 return raw
586def _deep_update(source: Dict[str, Any], overrides: Dict[str, Any]) -> Dict[str, Any]:
587 """Recursively update source dict with overrides dict."""
588 for key, value in overrides.items():
589 if isinstance(value, dict) and isinstance(source.get(key), dict):
590 _deep_update(source[key], value)
591 else:
592 source[key] = value
593 return source
595_ALLOWED_TOP_LEVEL_KEYS = frozenset({"model", "data", "train"})
598def _validate_top_level(config: Dict[str, Any]) -> None:
599 """Reject any top-level key other than ``model`` / ``data`` / ``train``.
601 Strict three-tier YAML shape. Any flat-style legacy key
602 (``parallel``, ``optim``, ``mixed_precision``, ``runtime``, ``debug`` ...)
603 must be moved under ``train.*`` — see schema.py for the canonical layout.
605 Raises:
606 ValueError: With migration hints when forbidden top-level keys are
607 present in the YAML.
608 """
609 forbidden = sorted(set(config) - _ALLOWED_TOP_LEVEL_KEYS)
610 if not forbidden:
611 return
613 legacy_to_train_path = {
614 "parallel": "train.accelerator",
615 "optim": "train.optimizer",
616 "mixed_precision": "train.mixed_precision",
617 "memory": "train.gradient_checkpointing",
618 "checkpoint": "train.checkpoint",
619 "logging": "train.logging",
620 "tensorboard": "train.tensorboard",
621 "wandb": "train.wandb",
622 "profiler": "train.profile",
623 "memory_monitor": "train.memory_monitor",
624 "moe_monitor": "train.moe_monitor",
625 "eval": "train.eval",
626 "runtime": "train (flatten init_device / backend / comm_backend)",
627 "debug": "train.debug",
628 "seed": "train.seed",
629 }
630 hints = []
631 for key in forbidden:
632 new_path = legacy_to_train_path.get(key)
633 if new_path:
634 hints.append(f" - top-level '{key}:' → move under {new_path}")
635 else:
636 hints.append(f" - top-level '{key}:' is not allowed")
637 raise ValueError(
638 "Forbidden top-level YAML keys: %s. The schema is strict three-tier "
639 "(model / data / train) — see config/schema.py. Migrate as follows:\n%s"
640 % (forbidden, "\n".join(hints))
641 )
644def _instantiate_recursive(cls: Type[T], config_dict: Dict[str, Any]) -> T:
645 """Recursively convert a dict into nested dataclass instances.
647 Unknown keys in ``config_dict`` that have no corresponding field on
648 ``cls`` emit a ``logger.warning`` with a closest-match suggestion from
649 ``difflib``, helping users catch typos in YAML configs.
650 """
651 if not is_dataclass(cls):
652 return config_dict
654 known = {f.name for f in fields(cls)}
655 unknown = set(config_dict) - known
656 for name in sorted(unknown):
657 matches = difflib.get_close_matches(name, known, n=1)
658 suggestion = f" Did you mean '{matches[0]}'?" if matches else ""
659 logger.warning(
660 "Unknown config key '%s' for %s ignored.%s",
661 name, cls.__name__, suggestion,
662 )
664 field_values = {}
665 for field_info in fields(cls):
666 if field_info.name not in config_dict:
667 continue
668 raw_value = config_dict[field_info.name]
669 field_type = field_info.type
671 # Unwrap Optional[X] → X
672 origin = get_origin(field_type)
673 if origin is Union:
674 unwrapped = [
675 a for a in get_args(field_type)
676 if a is not type(None) # pylint: disable=unidiomatic-typecheck
677 ]
678 if len(unwrapped) == 1:
679 field_type = unwrapped[0]
681 if is_dataclass(field_type) and isinstance(raw_value, dict):
682 field_values[field_info.name] = _instantiate_recursive(field_type, raw_value)
683 elif field_type is bool and isinstance(raw_value, str):
684 field_values[field_info.name] = _string_to_bool(raw_value)
685 else:
686 field_values[field_info.name] = raw_value
688 return cls(**field_values)
691def parse_args(root_class: Type[T]) -> T:
692 """Parse training config from YAML file + CLI overrides.
694 Usage::
696 args = parse_args(HyperTrainerConfig)
698 The first positional argument is the YAML config file path.
699 CLI arguments use dot-path notation under the strict three-tier schema:
700 ``--train.accelerator.dp_shard=8 --train.optimizer.lr=3e-4``
702 Bool fields accept extended string aliases (``yes/no/on/off/y/n/t/f/1/0``).
703 Unknown YAML keys emit a warning with a closest-match suggestion.
705 Args:
706 root_class: The root config dataclass type.
708 Returns:
709 An instance of root_class populated from YAML + CLI.
710 """
711 parser = argparse.ArgumentParser(description="HyperParallel Trainer")
712 parser.add_argument("config_file", nargs="?", help="Path to YAML config file")
713 args, remaining = parser.parse_known_args()
715 # Load YAML
716 final_config: dict = {}
717 if args.config_file:
718 if not os.path.isfile(args.config_file):
719 logger.warning(
720 "Config file not found: %s (cwd=%s). Using all defaults.",
721 args.config_file, os.getcwd(),
722 )
723 else:
724 with open(args.config_file, encoding="utf-8") as f:
725 yaml_config = yaml.safe_load(f)
726 if yaml_config:
727 final_config = yaml_config
729 # Parse CLI dot-path overrides: --train.accelerator.dp=8 → nested dict
730 cli_config: dict = {}
731 for item in remaining:
732 if item.startswith("--") and "=" in item:
733 dot_key, raw_value = item[2:].split("=", 1)
734 coerced = _coerce_cli_value(raw_value, dot_key, root_class)
735 keys = dot_key.split(".")
736 current = cli_config
737 for k in keys[:-1]:
738 current = current.setdefault(k, {})
739 current[keys[-1]] = coerced
741 # CLI overrides YAML
742 final_config = _deep_update(final_config, cli_config)
744 # Strict three-tier validation — only model / data / train allowed.
745 _validate_top_level(final_config)
747 # local_rank from environment (torchrun sets it).
748 local_rank = int(os.environ.get("LOCAL_RANK", "0"))
749 final_config.setdefault("train", {})["local_rank"] = local_rank
751 return _instantiate_recursive(root_class, final_config)