Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / config.py: 98%
292 statements
« 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"""Training configuration schema — strict three-tier (model/data/train).
17Top-level keys are exactly ``model``, ``data``, ``train`` and nothing else.
18arguments schema.
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, 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 # Universal transformer architecture overrides.
86 num_hidden_layers: Optional[int] = None
87 hidden_size: Optional[int] = None
88 intermediate_size: Optional[int] = None
89 num_attention_heads: Optional[int] = None
90 num_key_value_heads: Optional[int] = None
91 vocab_size: Optional[int] = None
92 max_position_embeddings: Optional[int] = None
93 # Free-form per-model overrides handed to ``build_model_fn``.
94 config_overrides: Optional[dict] = None
96# ============================================================================
97# data:
98# ============================================================================
101@dataclass
102class DataConfig:
103 """``data.*`` — dataset, tokenizer/processor, sampler, batch shape.
105 - ``streaming``: only ``False`` supported today.
106 - ``num_workers``: keep ≥ 2 for real datasets.
107 - ``shuffle``: when ``False``, sampler reads samples in dataset order.
108 """
109 type: str = "dummy"
110 train_path: Optional[str] = None
111 subset: Optional[str] = None
112 max_seq_len: int = 2048
113 text_key: str = "text"
114 train_size: Optional[int] = None
115 # multimodal / VL (synthetic vl_dummy path)
116 template: str = "empty"
117 image_key: str = "image"
118 messages_key: str = "messages"
119 image_token_id: int = 151655
120 vl_grid_t: int = 2
121 vl_grid_h: int = 2
122 vl_grid_w: int = 2
123 # loader perf
124 streaming: bool = False
125 num_workers: int = 0
126 prefetch_factor: Optional[int] = None
127 pin_memory: bool = True
128 shuffle: bool = True
130# ============================================================================
131# train.* — sub-configs
132# ============================================================================
135@dataclass
136class AcceleratorConfig:
137 """``train.accelerator.*`` — parallelism topology.
139 Two ways to express data parallelism:
141 - **Legacy single field** (back-compat): ``dp`` only — maps to
142 ``dp_shard`` for FSDP.
143 - ** split**: ``dp_replicate`` (HSDP outer) and
144 ``dp_shard`` (FSDP inner). Pass ``dp_shard=-1`` to auto-fill from
145 ``world_size / (dp_replicate * cp * tp * pp)``.
147 For MoE: ``etp`` controls expert TP. Must equal ``tp`` or ``1``.
148 ``moe_token_dispatcher_type`` selects the EP token exchange algorithm.
149 ``npu_nums_per_device`` is the inner-EP degree for the deredundency
150 dispatcher; ``oep`` is inferred as ``ep // npu_nums_per_device``.
151 """
152 dp: Optional[int] = None
153 dp_replicate: int = 1
154 dp_shard: Optional[int] = None
155 tp: int = 1
156 cp: int = 1
157 pp: int = 1
158 ep: int = 1
159 etp: int = 1
160 moe_token_dispatcher_type: str = "all_to_all"
161 npu_nums_per_device: int = 8
162 zero_stage: int = 0
163 reshard_after_forward: bool = True
164 async_cp: bool = False
165 ulysses_degree: Optional[int] = None
166 # Bucketed reduce-scatter: single fused RS per FSDP unit, stable fp32
167 # reduction order across runs.
168 comm_fusion: bool = True
171@dataclass
172class MixedPrecisionConfig:
173 """``train.mixed_precision.*`` — FSDP2 mp_policy fields.
175 ``output_dtype`` controls forward-output dtype at FSDP wrap boundaries.
176 Set this to ``"bfloat16"`` to keep the cross-FSDP-unit forward output
177 in bf16 (matches the typical "uniform low-precision forward" mp_policy
178 setup); leave ``None`` to inherit from ``param_dtype``.
179 """
180 enabled: bool = False
181 param_dtype: str = "bfloat16"
182 reduce_dtype: str = "float32"
183 output_dtype: Optional[str] = None
186@dataclass
187class GradientCheckpointingConfig:
188 """``train.gradient_checkpointing.*`` — activation recomputation.
190 . Modes: ``"off"``, ``"full"``, or ``"selective"``.
191 """
192 activation_checkpoint: str = "off"
195@dataclass
196class OptimizerConfig:
197 """``train.optimizer.*`` — optimizer + LR schedule + grad clip.
199 ``loss_aggregation``: how the per-micro-batch loss is scaled before
200 backward. ``"token_weighted"`` divides the summed loss by the global
201 valid-token count; ``"rank_average"`` averages per-rank micro-batch
202 means and is preferred when batches have variable valid-token counts
203 across ranks.
204 """
205 type: str = "adamw"
206 lr: float = 1e-4
207 lr_min: float = 1e-5
208 lr_decay_style: str = "cosine"
209 lr_warmup_ratio: float = 0.1
210 loss_aggregation: str = "token_weighted"
211 weight_decay: float = 0.01
212 max_grad_norm: float = 1.0
213 bsz_warmup_ratio: float = 0.0
214 eps: float = 1e-8
215 betas: tuple = (0.9, 0.999)
218@dataclass
219class CheckpointConfig:
220 """``train.checkpoint.*`` — DCP save / load + HF export."""
221 output_dir: str = "outputs"
222 save_steps: int = 500
223 save_hf_weights: bool = True
224 load_path: Optional[str] = None
225 save_async: bool = False
228@dataclass
229class LoggingConfig:
230 """``train.logging.*`` — console / metric output (consumed by LoggingCallback)."""
231 report_to: str = "none"
232 report_global_loss: bool = False
233 log_steps: int = 10
234 report_throughput: bool = True
235 model_flops_per_token: Optional[int] = None
236 peak_tflops: Optional[float] = None # e.g. 312.0 for A100 bf16
239@dataclass
240class TensorBoardConfig:
241 """``train.tensorboard.*`` — TB SummaryWriter on rank 0."""
242 enabled: bool = False
243 output_dir: str = "tb_traces"
244 log_steps: int = 1
247@dataclass
248class WandbConfig:
249 """``train.wandb.*`` — W&B run logging on rank 0."""
250 enabled: bool = False
251 project: str = "hyper-parallel"
252 run_name: Optional[str] = None
253 log_steps: int = 1
256@dataclass
257class ProfileConfig:
258 """``train.profile.*`` — torch.profiler schedule ().
260 Schedule semantics: wait → warmup → active.
261 """
262 enabled: bool = False
263 output_dir: str = "profiler_traces"
264 wait_steps: int = 1
265 warmup_steps: int = 1
266 active_steps: int = 3
269@dataclass
270class MemoryMonitorConfig:
271 """``train.memory_monitor.*`` — periodic device-memory snapshot."""
272 enabled: bool = False
273 log_steps: int = 50
274 reset_peak_each_step: bool = False
277@dataclass
278class MoEMonitorConfig:
279 """``train.moe_monitor.*`` — MoE routing / load-balance monitor.
281 When enabled, :class:`~hyper_parallel.core.moe_utils.MoEMonitorCallback`
282 automatically syncs ``tokens_per_expert`` across distributed ranks and
283 updates ``expert_bias`` after each optimizer step. The mean ``aux_loss``
284 across MoE layers is exposed via ``last_mean_aux_loss`` so that
285 :class:`LoggingCallback` can print it alongside the main training loss.
287 DP/TP+SP/CP group information is automatically obtained from the trainer's
288 device mesh — no manual group configuration needed.
290 Args:
291 enabled: Whether to activate the MoE monitor callback.
292 lr: Step size for expert bias updates. Defaults to ``1e-3``.
293 num_recomputations: Number of forward executions per optimizer step.
294 Default ``1``. Set to ``2`` when activation checkpoint is enabled.
295 """
296 enabled: bool = False
297 lr: float = 1e-3
298 num_recomputations: int = 1
301@dataclass
302class EvalConfig:
303 """``train.eval.*`` — eval cadence + dataset."""
304 eval_steps: int = 0
305 eval_dataset: Optional[str] = None
308@dataclass
309class DebugConfig:
310 """``train.debug.*`` — reproducibility and numerical-stability knobs.
312 All flags here are production-safe; they tune determinism (CI / paper
313 reproducibility), guard against numerical blow-ups, and bound memory
314 growth in long runs.
315 """
316 deterministic: bool = False
317 deterministic_warn_only: bool = False
318 check_nan_inf: bool = False
319 gc_steps: int = 0
321# ============================================================================
322# train: (top of the train section, holds the sub-configs)
323# ============================================================================
326@dataclass
327class TrainConfig:
328 """``train.*`` — full training-section config.
330 Flat fields cover the basic loop knobs (steps, batch shape, init device,
331 seed, comm backend); nested sub-configs cover everything else.
332 """
333 # Loop shape
334 max_steps: int = 100
335 num_train_epochs: int = 1
336 global_batch_size: int = 8
337 micro_batch_size: int = 1
338 seed: int = 42
340 # Runtime / device
341 backend: str = "torch"
342 init_device: str = "meta"
343 comm_backend: Optional[str] = None
344 local_rank: int = 0 # set from LOCAL_RANK env by parser
346 # Sub-configs
347 accelerator: AcceleratorConfig = field(default_factory=AcceleratorConfig)
348 mixed_precision: MixedPrecisionConfig = field(default_factory=MixedPrecisionConfig)
349 gradient_checkpointing: GradientCheckpointingConfig = field(
350 default_factory=GradientCheckpointingConfig
351 )
352 optimizer: OptimizerConfig = field(default_factory=OptimizerConfig)
353 checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig)
354 logging: LoggingConfig = field(default_factory=LoggingConfig)
355 tensorboard: TensorBoardConfig = field(default_factory=TensorBoardConfig)
356 wandb: WandbConfig = field(default_factory=WandbConfig)
357 profile: ProfileConfig = field(default_factory=ProfileConfig)
358 memory_monitor: MemoryMonitorConfig = field(default_factory=MemoryMonitorConfig)
359 moe_monitor: MoEMonitorConfig = field(default_factory=MoEMonitorConfig)
360 eval: EvalConfig = field(default_factory=EvalConfig)
361 debug: DebugConfig = field(default_factory=DebugConfig)
363# ============================================================================
364# Top-level: model / data / train (and only these three)
365# ============================================================================
368@dataclass
369class HyperTrainerConfig:
370 """Top-level config — strict three-tier ().
372 Allowed top-level keys: ``model``, ``data``, ``train``. Anything else in
373 the YAML is rejected by the parser with a typo-suggestion message.
374 """
375 model: ModelConfig = field(default_factory=ModelConfig)
376 data: DataConfig = field(default_factory=DataConfig)
377 train: TrainConfig = field(default_factory=TrainConfig)
379 # Computed (no user input)
380 train_steps: int = 0
382 def __post_init__(self):
383 self.train_steps = self.train.max_steps
385# ==============================================================================
386# CLI / YAML parser
387# ==============================================================================
388# Configuration parser: YAML file + CLI dot-path overrides.
389#
390# Supports:
391# - Unknown YAML/CLI keys emit a warning with difflib closest-match suggestions.
392# - Bool fields accept string aliases: ``true/yes/y/on/1/t`` -> ``True``,
393# ``false/no/n/off/0/f`` -> ``False``. Only applied when the dataclass
394# field type resolves to ``bool`` or ``Optional[bool]`` to avoid ambiguity.
397def _string_to_bool(value: Any) -> bool:
398 """Convert common string representations of booleans to ``bool``.
400 Accepts: ``true/yes/y/on/1/t`` → ``True``,
401 ``false/no/n/off/0/f`` → ``False``.
403 Args:
404 value: A string or bool value.
406 Returns:
407 The corresponding ``bool``.
409 Raises:
410 ValueError: When the string cannot be mapped to a bool.
411 """
412 if isinstance(value, bool):
413 return value
414 if isinstance(value, str):
415 lower = value.lower()
416 if lower in _BOOL_TRUE_STRINGS:
417 return True
418 if lower in _BOOL_FALSE_STRINGS:
419 return False
420 raise ValueError(
421 f"Cannot convert {value!r} to bool. "
422 "Expected one of: true/false/yes/no/y/n/on/off/1/0/t/f"
423 )
426def _resolve_field_type(cls: Type, dot_path: str) -> Optional[Type]:
427 """Walk a dataclass hierarchy to find the resolved type of a dot-path field.
429 Args:
430 cls: Root dataclass class.
431 dot_path: Dot-separated field path, e.g. ``"debug.deterministic"``.
433 Returns:
434 The resolved Python type, or ``None`` if the path cannot be resolved.
435 """
436 parts = dot_path.split(".")
437 current_cls = cls
438 for part in parts:
439 if not is_dataclass(current_cls):
440 return None
441 found = None
442 for f in fields(current_cls):
443 if f.name == part:
444 found = f
445 break
446 if found is None:
447 return None
448 field_type = found.type
449 # Unwrap Optional[X] → X
450 origin = get_origin(field_type)
451 if origin is Union:
452 unwrapped = [a for a in get_args(field_type) if a is not type(None)]
453 field_type = unwrapped[0] if len(unwrapped) == 1 else field_type
454 current_cls = field_type
455 return current_cls
458def _coerce_cli_value(raw: str, dot_path: str, root_class: Type) -> Any:
459 """Parse a CLI string value, coercing to the correct type for the field.
461 Bool fields accept an extended string set. For all other
462 fields the existing int → float → str heuristic is used.
464 Args:
465 raw: Raw string from the CLI.
466 dot_path: Dot-separated field path used for type lookup.
467 root_class: Root dataclass class for type resolution.
469 Returns:
470 Coerced value.
471 """
472 field_type = _resolve_field_type(root_class, dot_path)
473 if field_type is bool:
474 try:
475 return _string_to_bool(raw)
476 except ValueError:
477 pass # fall through to heuristic below
478 # Existing heuristic: int → float → bool-string → str
479 try:
480 return int(raw)
481 except ValueError:
482 pass
483 try:
484 return float(raw)
485 except ValueError:
486 pass
487 if raw.lower() in ("true", "false"):
488 return raw.lower() == "true"
489 return raw
492def _deep_update(source: Dict[str, Any], overrides: Dict[str, Any]) -> Dict[str, Any]:
493 """Recursively update source dict with overrides dict."""
494 for key, value in overrides.items():
495 if isinstance(value, dict) and isinstance(source.get(key), dict):
496 source[key] = _deep_update(source[key], value)
497 else:
498 source[key] = value
499 return source
501_ALLOWED_TOP_LEVEL_KEYS = frozenset({"model", "data", "train"})
504def _validate_top_level(config: Dict[str, Any]) -> None:
505 """Reject any top-level key other than ``model`` / ``data`` / ``train``.
507 Strict three-tier YAML shape. Any flat-style legacy key
508 (``parallel``, ``optim``, ``mixed_precision``, ``runtime``, ``debug`` ...)
509 must be moved under ``train.*`` — see schema.py for the canonical layout.
511 Raises:
512 ValueError: With migration hints when forbidden top-level keys are
513 present in the YAML.
514 """
515 forbidden = sorted(set(config) - _ALLOWED_TOP_LEVEL_KEYS)
516 if not forbidden:
517 return
519 legacy_to_train_path = {
520 "parallel": "train.accelerator",
521 "optim": "train.optimizer",
522 "mixed_precision": "train.mixed_precision",
523 "memory": "train.gradient_checkpointing",
524 "checkpoint": "train.checkpoint",
525 "logging": "train.logging",
526 "tensorboard": "train.tensorboard",
527 "wandb": "train.wandb",
528 "profiler": "train.profile",
529 "memory_monitor": "train.memory_monitor",
530 "moe_monitor": "train.moe_monitor",
531 "eval": "train.eval",
532 "runtime": "train (flatten init_device / backend / comm_backend)",
533 "debug": "train.debug",
534 "seed": "train.seed",
535 }
536 hints = []
537 for key in forbidden:
538 new_path = legacy_to_train_path.get(key)
539 if new_path:
540 hints.append(f" - top-level '{key}:' → move under {new_path}")
541 else:
542 hints.append(f" - top-level '{key}:' is not allowed")
543 raise ValueError(
544 "Forbidden top-level YAML keys: %s. The schema is strict three-tier "
545 "(model / data / train) — see config/schema.py. Migrate as follows:\n%s"
546 % (forbidden, "\n".join(hints))
547 )
550def _instantiate_recursive(cls: Type[T], config_dict: Dict[str, Any]) -> T:
551 """Recursively convert a dict into nested dataclass instances.
553 Unknown keys in ``config_dict`` that have no corresponding field on
554 ``cls`` emit a ``logger.warning`` with a closest-match suggestion from
555 ``difflib``, helping users catch typos in YAML configs.
556 """
557 if not is_dataclass(cls):
558 return config_dict
560 known = {f.name for f in fields(cls)}
561 unknown = set(config_dict) - known
562 for name in sorted(unknown):
563 matches = difflib.get_close_matches(name, known, n=1)
564 suggestion = f" Did you mean '{matches[0]}'?" if matches else ""
565 logger.warning(
566 "Unknown config key '%s' for %s ignored.%s",
567 name, cls.__name__, suggestion,
568 )
570 field_values = {}
571 for field_info in fields(cls):
572 if field_info.name not in config_dict:
573 continue
574 raw_value = config_dict[field_info.name]
575 field_type = field_info.type
577 # Unwrap Optional[X] → X
578 origin = get_origin(field_type)
579 if origin is Union:
580 unwrapped = [a for a in get_args(field_type) if a is not type(None)]
581 if len(unwrapped) == 1:
582 field_type = unwrapped[0]
584 if is_dataclass(field_type) and isinstance(raw_value, dict):
585 field_values[field_info.name] = _instantiate_recursive(field_type, raw_value)
586 elif field_type is bool and isinstance(raw_value, str):
587 field_values[field_info.name] = _string_to_bool(raw_value)
588 else:
589 field_values[field_info.name] = raw_value
591 return cls(**field_values)
594def parse_args(root_class: Type[T]) -> T:
595 """Parse training config from YAML file + CLI overrides.
597 Usage::
599 args = parse_args(HyperTrainerConfig)
601 The first positional argument is the YAML config file path.
602 CLI arguments use dot-path notation under the strict three-tier schema:
603 ``--train.accelerator.dp_shard=8 --train.optimizer.lr=3e-4``
605 Bool fields accept extended string aliases (``yes/no/on/off/y/n/t/f/1/0``).
606 Unknown YAML keys emit a warning with a closest-match suggestion.
608 Args:
609 root_class: The root config dataclass type.
611 Returns:
612 An instance of root_class populated from YAML + CLI.
613 """
614 parser = argparse.ArgumentParser(description="HyperParallel Trainer")
615 parser.add_argument("config_file", nargs="?", help="Path to YAML config file")
616 args, remaining = parser.parse_known_args()
618 # Load YAML
619 final_config: dict = {}
620 if args.config_file:
621 if not os.path.isfile(args.config_file):
622 logger.warning(
623 "Config file not found: %s (cwd=%s). Using all defaults.",
624 args.config_file, os.getcwd(),
625 )
626 else:
627 with open(args.config_file, encoding="utf-8") as f:
628 yaml_config = yaml.safe_load(f)
629 if yaml_config:
630 final_config = yaml_config
632 # Parse CLI dot-path overrides: --train.accelerator.dp=8 → nested dict
633 cli_config: dict = {}
634 for item in remaining:
635 if item.startswith("--") and "=" in item:
636 dot_key, raw_value = item[2:].split("=", 1)
637 coerced = _coerce_cli_value(raw_value, dot_key, root_class)
638 keys = dot_key.split(".")
639 current = cli_config
640 for k in keys[:-1]:
641 current = current.setdefault(k, {})
642 current[keys[-1]] = coerced
644 # CLI overrides YAML
645 final_config = _deep_update(final_config, cli_config)
647 # Strict three-tier validation — only model / data / train allowed.
648 _validate_top_level(final_config)
650 # local_rank from environment (torchrun sets it).
651 local_rank = int(os.environ.get("LOCAL_RANK", "0"))
652 final_config.setdefault("train", {})["local_rank"] = local_rank
654 return _instantiate_recursive(root_class, final_config)