Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / config_adapter / _config_loader.py: 98%
261 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""Configuration loader for auto parallel strategy search.
17Reads Search Config (``search.yaml``) and HyperParallel training config
18(``train.yaml``) files, producing :class:`NormalizedConfig` instances.
19"""
21import copy
22import logging
23import os
24from typing import Any, Dict, List, Optional, Set, Tuple
26try:
27 import yaml # type: ignore[import-untyped] # pylint: disable=C0415
28except ImportError:
29 yaml = None # pragma: no cover
31from hyper_parallel.auto_parallel._hf_model_spec import (
32 exec_overrides,
33 is_auto_models_schema,
34 resolve_hf_model_spec,
35)
36from hyper_parallel.auto_parallel._layer_census import replacement_specs
37from hyper_parallel.auto_parallel.config_adapter._normalized_config import NormalizedConfig
38from hyper_parallel.auto_parallel.sapp_nd.nd.recompute_dimension import read_activation_checkpoint_mode
40logger = logging.getLogger(__name__)
43def _recompute_strategy(stated: Any, where: str) -> str:
44 """The estimator's recompute strategy for an activation checkpoint mode a yaml states, which says off as none."""
45 mode = read_activation_checkpoint_mode(stated, where)
46 return "none" if mode == "off" else mode
48# Mapping from HuggingFace-style config_overrides keys to internal names.
49# After field-name alignment, most HP YAML keys already match the internal
50# model_spec names. Only ``seq_length`` (MF variant) needs remapping.
51_HP_TO_INTERNAL: Dict[str, str] = {
52 "seq_length": "max_position_embeddings",
53}
56def _normalize_model_spec(model_spec: Dict[str, Any]) -> Dict[str, Any]:
57 """Rename non-standard config overrides keys to canonical names.
59 Only maps a key if the target name is not already present, so
60 explicit canonical names in the YAML take precedence.
61 """
62 for hf_key, internal_key in _HP_TO_INTERNAL.items():
63 if hf_key in model_spec and internal_key not in model_spec:
64 model_spec[internal_key] = model_spec.pop(hf_key)
65 # A resolved Transformers config states its expert count, never a
66 # moe_enabled flag; the constraint checker and the emitted strategy
67 # both read the flag, so derive it rather than reporting a MoE as dense.
68 if "moe_enabled" not in model_spec:
69 model_spec["moe_enabled"] = int(model_spec.get("num_experts") or 1) > 1
70 return model_spec
73# Mapping from short dimension names (used in search config YAML parallelism section)
74# to canonical NormalizedConfig search_space keys.
75_UNIFIED_DIM_MAP: Dict[str, str] = {
76 "dp": "data_parallel_replicate_degree",
77 "fsdp": "data_parallel_shard_degree",
78 "tp": "tensor_parallel_degree",
79 "pp": "pipeline_parallel_degree",
80 "cp": "context_parallel_degree",
81 "ep": "expert_parallel_degree",
82 "etp": "expert_tensor_parallel_degree",
83 "micro_batch_num": "micro_batch_num",
84}
86# Mapping from short dimension names to constraint fixed_*_degree keys.
87_FIXED_DIM_MAP: Dict[str, str] = {
88 "dp": "fixed_dp_degree",
89 "fsdp": "fixed_fsdp_degree",
90 "tp": "fixed_tp_degree",
91 "pp": "fixed_pp_degree",
92 "cp": "fixed_cp_degree",
93 "ep": "fixed_ep_degree",
94 "etp": "fixed_etp_degree",
95 "micro_batch_num": "fixed_micro_batch_num",
96}
99def _get_dict(raw: Dict[str, Any], key: str) -> Dict[str, Any]:
100 """Return the value of a key if it is a dict, otherwise an empty dict."""
101 val = raw.get(key, {})
102 return val if isinstance(val, dict) else {}
105# The model section's keys that state how the run builds, loads and computes
106# the model, not the model; the cost model reads them there.
107_MODEL_RUN_KEYS = ("_target_", "torch_dtype", "param_init_type", "compute_dtype", "softmax_compute_type")
109# The root keys of an AutoModels train.yaml that state the run. Its
110# plan_overrides state the modules it builds, which a census measures.
111_RUN_SECTIONS = ("model_init_dtype", "accelerator", "fsdp_config", "training", "optimizer", "plan_overrides")
113# The ``context`` keys of a train.yaml that state how the cost model prices
114# the run: a census of the layers, the census's spec of the model, a vision
115# tower's token count, a launcher that shards each strategy's experts over
116# its whole expert group, and what a MoE layer's token dispatch costs on this
117# cluster. The device count and the memory budget are a search's own.
118_RUN_CONTEXT_KEYS = ("census", "census_spec", "visual_seq_len", "expert_shard", "moe_dispatch")
120# The keys of a legacy train.yaml's ``train`` section that are not its
121# training settings: recompute, the search's, and precision, which the cost
122# model reads from the model section.
123_LEGACY_TRAIN_SECTIONS = ("accelerator", "optimizer", "gradient_checkpointing", "mixed_precision")
126def _model_run(model_raw: Dict[str, Any]) -> Dict[str, Any]:
127 """Return the dtypes a model section states for the run."""
128 return {key: model_raw[key] for key in _MODEL_RUN_KEYS if model_raw.get(key) is not None}
131def _stated_run(raw: Dict[str, Any]) -> Dict[str, Any]:
132 """Return the run a train.yaml states, in the AutoModels sections' names.
134 The model's dtypes, FSDP's precision and resharding, the optimizer,
135 gradient clipping and the accelerator's settings are the run's, and the
136 model spec carries none of them. The strategy the same sections state
137 stays in: the search runner writes the one it searches over it. A legacy
138 train.yaml states them under ``train``. The context's pricing options,
139 :data:`_RUN_CONTEXT_KEYS`, ride along.
140 """
141 run: Dict[str, Any] = {"model": _model_run(_get_dict(raw, "model"))}
142 context = {key: value for key, value in _get_dict(raw, "context").items() if key in _RUN_CONTEXT_KEYS}
143 if context:
144 run["context"] = copy.deepcopy(context)
145 if is_auto_models_schema(raw):
146 for key in _RUN_SECTIONS:
147 if raw.get(key) is not None:
148 run[key] = copy.deepcopy(raw[key])
149 return run
150 train_raw = _get_dict(raw, "train")
151 run["training"] = {
152 key: copy.deepcopy(value) for key, value in train_raw.items() if key not in _LEGACY_TRAIN_SECTIONS
153 }
154 for key in ("accelerator", "optimizer"):
155 if isinstance(train_raw.get(key), dict):
156 run[key] = copy.deepcopy(train_raw[key])
157 return run
160def _load_auto_models_model_spec(
161 model_raw: Dict[str, Any],
162 visual_seq_len: Optional[int] = None,
163 census_seq_len: int = 0,
164 census_spec: bool = False,
165 replacements: tuple = (),
166) -> Dict[str, Any]:
167 """Resolve model dimensions through the shared AutoModels path.
169 Delegates to :func:`resolve_hf_model_spec` so this reader and the
170 SAPP-ND parser cannot disagree about field names or fallbacks. The run
171 keys of ``config_overrides`` ride along, as the adapter hands them on.
172 """
173 spec = resolve_hf_model_spec(model_raw, visual_seq_len, census_seq_len, census_spec, replacements)
174 spec.update(exec_overrides(model_raw))
175 return _normalize_model_spec(spec)
178def _load_yaml(path: str) -> Dict[str, Any]:
179 """Read and parse a YAML file, returning the raw dict."""
180 if yaml is None:
181 raise ImportError(
182 "PyYAML is required to read HyperParallel YAML configs. "
183 "Install it with: pip install pyyaml"
184 )
185 if not os.path.isfile(path):
186 raise FileNotFoundError(f"Config file not found: {path}")
188 ext = os.path.splitext(path)[1].lower()
189 if ext not in (".yaml", ".yml"):
190 raise ValueError(
191 f"Unsupported config file format: {ext!r}. "
192 "Supported formats: .yaml, .yml"
193 )
195 try:
196 with open(path, "r", encoding="utf-8") as fh:
197 raw = yaml.safe_load(fh)
198 except yaml.YAMLError as exc:
199 raise ValueError(f"Failed to parse YAML file {path}: {exc}") from exc
201 if raw is None:
202 raw = {}
203 if not isinstance(raw, dict):
204 raise ValueError(
205 f"Config file {path} must contain a YAML mapping at the top level, "
206 f"got {type(raw).__name__}"
207 )
208 return raw
211# ── Search Config YAML reader (primary) ──────────────────────────────
214def _parse_unified_parallelism(
215 para_raw: Dict[str, Any],
216) -> Tuple[Dict[str, List[int]], Dict[str, Any], Set[str]]:
217 """Convert the unified parallelism declaration into search_space + constraint.
219 Rules:
220 * Scalar integer value → fixed dimension (both ``constraint.fixed_*``
221 and ``search_space`` with a single-element list).
222 * List value → search candidates (placed into ``search_space`` only).
223 * String ``"auto"`` → dimension is left to the searcher's ``bound_space``
224 (neither fixed nor explicitly enumerated).
226 Returns:
227 A ``(search_space, constraint, auto)`` tuple, where *auto* names the
228 dimensions the user handed to the searcher. They carry no candidates,
229 so without naming them they are indistinguishable from dimensions the
230 file never mentioned, and a ``train_yaml`` would pin them.
231 """
232 search_space: Dict[str, List[int]] = {}
233 constraint: Dict[str, Any] = {}
234 auto: Set[str] = set()
236 for short_key, canonical_key in _UNIFIED_DIM_MAP.items():
237 if short_key not in para_raw:
238 continue
239 value = para_raw[short_key]
241 if isinstance(value, int):
242 constraint[_FIXED_DIM_MAP[short_key]] = value
243 search_space[canonical_key] = [value]
244 elif isinstance(value, list):
245 search_space[canonical_key] = [int(v) for v in value]
246 elif isinstance(value, str) and value.strip().lower() == "auto":
247 auto.add(canonical_key)
249 return search_space, constraint, auto
252# Every unified dimension ND can actually search. ``etp`` is declarable but
253# has no ND dimension behind it, so it is always a fixed input.
254_NOT_SEARCHABLE = frozenset({"etp"})
257def _describe_parallelism(
258 search_space: Dict[str, List[int]],
259 auto_dims: Any,
260 declared: Dict[str, Any],
261 from_train_yaml: bool,
262) -> str:
263 """Report how every parallelism dimension was resolved.
265 A dimension left out of the file is not switched off: it is pinned to the
266 train.yaml value when there is one and searched freely when there is not.
267 Pin a dimension to 1 to take it out of the strategy.
268 """
269 parts = []
270 for short_key, canonical_key in _UNIFIED_DIM_MAP.items():
271 candidates = search_space.get(canonical_key)
272 if short_key in _NOT_SEARCHABLE:
273 fixed = candidates[0] if candidates else "default"
274 state = f"fixed {fixed} (no ND dimension)"
275 elif canonical_key in auto_dims:
276 state = "searched (auto)"
277 elif candidates is None:
278 state = "searched (not declared)"
279 elif len(candidates) == 1:
280 source = (
281 "declared" if short_key in declared
282 else "train.yaml" if from_train_yaml else "default"
283 )
284 state = f"fixed {candidates[0]} ({source})"
285 else:
286 state = f"searched over {candidates}"
287 parts.append(f"{short_key}={state}")
288 return "; ".join(parts)
291def _stated_recompute_modes(raw: Dict[str, Any]) -> Dict[str, Tuple[str, ...]]:
292 """Return the trainer modes ``recompute_modes`` lets a search choose among, as the estimator's entry.
294 Under ``recompute: auto`` or ``per_layer`` the search chooses among the
295 trainer's activation checkpoint modes: off and full, and its selective
296 policy where a census prices it. ``recompute_modes`` narrows them. YAML
297 reads an unquoted ``off`` as False.
299 Returns:
300 ``{"recompute_modes": modes}``, or an empty dict where the search
301 config states none.
303 Raises:
304 ValueError: For an empty list, or a mode the trainer does not run.
305 """
306 stated = raw.get("recompute_modes")
307 if stated is None:
308 return {}
309 listed = list(stated) if isinstance(stated, (list, tuple)) else [stated]
310 modes = tuple("off" if mode is False else str(mode) for mode in listed)
311 unknown = sorted(set(modes) - {"off", "selective", "full"})
312 if unknown or not modes:
313 raise ValueError(f"recompute_modes {listed}: expected some of off, selective and full")
314 return {"recompute_modes": modes}
317def _stated_recompute(raw: Dict[str, Any], parallelism_raw: Dict[str, Any]) -> Dict[str, Any]:
318 """Return the estimator's recompute entries, from ``parallelism.recompute`` where it is stated.
320 ``parallelism.recompute`` states recompute as a degree is stated: one
321 mode, a list of modes, or ``auto`` for every mode the trainer runs that
322 the search can price. The search then chooses among them for each
323 strategy, as ``recompute: auto`` with ``recompute_modes`` does, or for each
324 layer where the top-level ``recompute`` says ``per_layer``. A mode left
325 out is never chosen.
327 Returns:
328 ``recompute_strategy``, and ``recompute_modes`` where a list narrows
329 the modes.
331 Raises:
332 ValueError: For a mode the trainer does not run, or when recompute
333 is stated twice: the dimension beside ``recompute_modes``, or
334 beside a top-level ``recompute`` that fixes one mode.
335 """
336 stated = parallelism_raw.get("recompute")
337 strategy = str(raw.get("recompute", "none"))
338 if stated is None:
339 return {"recompute_strategy": strategy, **_stated_recompute_modes(raw)}
340 if "recompute_modes" in raw:
341 raise ValueError("recompute modes are stated twice, in parallelism.recompute and recompute_modes: "
342 "keep parallelism.recompute")
343 if strategy not in ("none", "auto", "per_layer"):
344 raise ValueError("recompute is stated twice: parallelism.recompute searches it as a dimension and the "
345 f"top-level recompute fixes {strategy}; keep parallelism.recompute, with recompute: "
346 "per_layer for a mode per layer")
347 chosen = "per_layer" if strategy == "per_layer" else "auto"
348 if not isinstance(stated, (list, tuple)) and str(stated).strip().lower() == "auto":
349 return {"recompute_strategy": chosen}
350 return {"recompute_strategy": chosen, **_stated_recompute_modes({"recompute_modes": stated})}
353def _describe_recompute(estimator: Dict[str, Any]) -> str:
354 """Report how a recompute dimension was resolved, as :func:`_describe_parallelism` reports a degree."""
355 modes = estimator.get("recompute_modes")
356 among = f"among {list(modes)}" if modes else "among every mode it can price"
357 per_layer = " per layer" if estimator["recompute_strategy"] == "per_layer" else ""
358 return f"recompute=chosen{per_layer} {among}"
361def _build_config_from_search_yaml(raw: Dict[str, Any]) -> NormalizedConfig:
362 """Construct a NormalizedConfig from a parsed Search Config YAML dict.
364 Supports two modes:
366 * **Standalone** (no ``train_yaml``) — ``model`` section and all other
367 info must be present in the search config file.
368 * **With train_yaml** — loads the specified ``train.yaml`` for model
369 parameters and current parallelism values, then overlays the search
370 config's ``cluster``, ``parallelism``, and ``constraint`` sections.
372 Undeclared parallelism dimensions are inherited from ``train.yaml``
373 as fixed (single-element) entries.
374 """
375 train_yaml_path = raw.get("train_yaml")
376 base_config: Optional[NormalizedConfig] = None
378 if train_yaml_path:
379 if not isinstance(train_yaml_path, str):
380 raise ValueError("'train_yaml' must be a file path string")
381 base_raw = _load_yaml(train_yaml_path)
382 base_config = _build_config_from_hp_yaml(base_raw)
384 model_spec: Dict[str, Any]
385 run: Dict[str, Any]
386 if base_config:
387 model_spec = dict(base_config.model_spec)
388 run = copy.deepcopy(base_config.run)
389 else:
390 model_spec = {}
391 run = {}
393 # Override or supply model section from search.yaml
394 search_model = _get_dict(raw, "model")
395 if search_model:
396 model_spec.update(search_model)
397 run.setdefault("model", {}).update(_model_run(search_model))
399 model_spec.setdefault("max_position_embeddings", 4096)
400 model_spec.setdefault("local_batch_size", 1)
402 model_spec = _normalize_model_spec(model_spec)
404 cluster_spec = _get_dict(raw, "cluster")
406 pp_raw = _get_dict(raw, "pp_config")
407 parallelism_raw = _get_dict(raw, "parallelism")
408 constraint_raw = _get_dict(raw, "constraint")
410 search_space, parallelism_constraint, auto_dims = _parse_unified_parallelism(
411 parallelism_raw
412 )
414 # Inherit undeclared dimensions from train.yaml as fixed values.
415 if base_config:
416 for space_key, candidates in base_config.search_space.items():
417 if space_key not in search_space and space_key not in auto_dims:
418 search_space[space_key] = candidates
420 if constraint_raw.get("global_batch_size", 0) is None or constraint_raw.get("global_batch_size", 0) == 0:
421 constraint_raw.setdefault(
422 "global_batch_size", base_config.constraint.get("global_batch_size", 0)
423 )
425 pp_config: Dict[str, Any] = {
426 "pp_degree": pp_raw.get("pp_degree",
427 parallelism_raw.get("pp", 1)),
428 "stage_partition_mode": pp_raw.get("stage_partition_mode", "uniform"),
429 "stage_partition": pp_raw.get("stage_partition", []),
430 "layer_offset_range": tuple(pp_raw.get("layer_offset_range", [0, 0])),
431 "layer_recompute_layers": pp_raw.get("layer_recompute_layers", []),
432 "micro_batch_num": pp_raw.get("micro_batch_num", 1),
433 "pp_interleave_num": pp_raw.get("pp_interleave_num", 1),
434 "pipeline_parallel_schedule": pp_raw.get("pipeline_schedule", "1F1B"),
435 }
437 estimator: Dict[str, Any] = {
438 "type": "symbolic",
439 **_stated_recompute(raw, parallelism_raw),
440 "enable_profiling_calibration": False,
441 }
443 constraint: Dict[str, Any] = {
444 "global_batch_size": constraint_raw.get("global_batch_size", 0),
445 "memory_limit_gb": constraint_raw.get("memory_limit_gb", 0.0),
446 **parallelism_constraint,
447 }
449 summary = _describe_parallelism(
450 search_space, auto_dims, parallelism_raw, bool(base_config)
451 )
452 if parallelism_raw.get("recompute") is not None:
453 summary += "; " + _describe_recompute(estimator)
454 logger.info("parallelism resolved: %s", summary)
456 return NormalizedConfig(
457 model_spec=model_spec,
458 cluster_spec=cluster_spec,
459 search_space=search_space,
460 constraint=constraint,
461 parallelism_summary=summary,
462 estimator=estimator,
463 pp_config=pp_config,
464 run=run,
465 )
468def read_search_config(path: str) -> NormalizedConfig:
469 """Read a Search Config YAML file and return a :class:`NormalizedConfig`.
471 The Search Config YAML format uses a unified ``parallelism`` section
472 where each dimension is declared as::
474 parallelism:
475 tp: 4 # scalar → fixed input
476 dp: [1, 2, 4] # list → search candidate
477 pp: auto # string → let the searcher decide
479 To reuse model parameters from an existing ``train.yaml`` without
480 duplicating them, set the ``train_yaml`` key::
482 train_yaml: "./train.yaml" # load model params from here
483 cluster:
484 num_nodes: 4
485 cards_per_node: 16 # optional: defaults to the -A device type
486 parallelism:
487 dp: [1, 2, 4]
488 tp: [1, 2, 4, 8]
490 Dimensions absent from ``parallelism`` are inherited from
491 ``train.yaml`` as fixed values. A ``model`` section in the search
492 config overrides values read from ``train_yaml``.
494 See ``auto_parallel/examples/dense_llm_search.yaml`` for a complete
495 standalone example.
497 Args:
498 path: Path to the YAML config file (``.yaml`` or ``.yml``).
500 Returns:
501 A :class:`NormalizedConfig` instance.
503 Raises:
504 FileNotFoundError: If the file does not exist.
505 ValueError: If the file cannot be parsed, or if ``cluster``
506 is missing when ``train_yaml`` is not used.
507 ImportError: If PyYAML is not installed.
508 """
509 raw = _load_yaml(path)
510 return _build_config_from_search_yaml(raw)
513# ── HyperParallel training YAML reader (secondary) ──────────────────
515_ACCEL_TO_SEARCH = {
516 "dp_shard": "data_parallel_shard_degree",
517 "dp_replicate": "data_parallel_replicate_degree",
518 "tp_degree": "tensor_parallel_degree",
519 "pipeline_parallel_degree": "pipeline_parallel_degree",
520 "context_parallel_degree": "context_parallel_degree",
521 "expert_parallel_degree": "expert_parallel_degree",
522 "expert_tensor_parallel_degree": "expert_tensor_parallel_degree",
523}
525_AUTO_MODELS_ACCEL_TO_SEARCH = {
526 "tp_size": "tensor_parallel_degree",
527 "pp_size": "pipeline_parallel_degree",
528 "cp_size": "context_parallel_degree",
529 "ep_size": "expert_parallel_degree",
530}
533def _census_replacements(raw: Dict[str, Any], census_seq_len: int) -> tuple:
534 """The module replacements a census runs its layers with: the ones the train.yaml's plan_overrides install."""
535 if not census_seq_len:
536 return ()
537 return replacement_specs(raw.get("plan_overrides") or ())
540def _training_seq_len(raw: Dict[str, Any]) -> Any:
541 """Return the training sequence length an AutoModels YAML states, or None.
543 Reads the three spellings the SAPP-ND parser accepts, in its order, so the
544 two halves of the cost model agree on where the length comes from: the
545 Online path's, an Indexed Dataset's, the legacy one.
546 """
547 dataset_raw = _get_dict(raw, "dataset")
548 return (_get_dict(dataset_raw, "data_transform").get("max_seq_len")
549 or _get_dict(dataset_raw, "data_config").get("seq_length")
550 or _get_dict(raw, "data").get("max_seq_len"))
553def _load_auto_models_model_spec_from_yaml(raw: Dict[str, Any]) -> Dict[str, Any]:
554 """Load model dimensions and training attributes from AutoModels YAML."""
555 model_raw = _get_dict(raw, "model")
556 training_raw = _get_dict(raw, "training")
557 context_raw = _get_dict(raw, "context")
558 seq_len = _training_seq_len(raw)
559 # A census builds its layers from the checkpoint's config, which only
560 # this reader resolves: the search hands ND the spec, its records in it.
561 census_seq_len = int(seq_len or 4096) if context_raw.get("census") else 0
562 model_spec = _load_auto_models_model_spec(
563 model_raw, context_raw.get("visual_seq_len"), census_seq_len, bool(context_raw.get("census_spec")),
564 _census_replacements(raw, census_seq_len),
565 )
566 if seq_len:
567 model_spec["max_position_embeddings"] = seq_len
568 else:
569 logger.warning(
570 "no dataset.data_transform.max_seq_len, dataset.data_config.seq_length "
571 "nor data.max_seq_len: costing the model's context limit of %s tokens, "
572 "which for a long-context model puts every candidate out of memory",
573 model_spec.get("max_position_embeddings", 4096),
574 )
575 model_spec.setdefault("max_position_embeddings", 4096)
576 if context_raw.get("device_num") is not None:
577 model_spec["device_num"] = int(context_raw["device_num"])
578 model_spec["local_batch_size"] = training_raw.get("micro_batch_size", 1)
579 model_spec["compute_dtype"] = model_raw.get("torch_dtype", "bfloat16")
580 return model_spec
583def _load_auto_models_parallelism(
584 raw: Dict[str, Any],
585) -> Tuple[Dict[str, List[int]], int, int]:
586 """Load fixed parallelism degrees from AutoModels YAML."""
587 accelerator_raw = _get_dict(raw, "accelerator")
588 fsdp_raw = _get_dict(raw, "fsdp_config")
589 search_space: Dict[str, List[int]] = {}
591 dp_shard_size = fsdp_raw.get("dp_shard_size")
592 if dp_shard_size is not None:
593 search_space["data_parallel_shard_degree"] = [int(dp_shard_size)]
594 for field_name, search_name in _AUTO_MODELS_ACCEL_TO_SEARCH.items():
595 value = accelerator_raw.get(field_name)
596 if value is not None:
597 search_space[search_name] = [int(value)]
599 data_parallel_size = int(dp_shard_size or 1)
600 pp_degree = max(1, int(accelerator_raw.get("pp_size", 1) or 1))
601 return search_space, data_parallel_size, pp_degree
604def _build_config_from_auto_models_yaml(raw: Dict[str, Any]) -> NormalizedConfig:
605 """Construct a normalized config from the current AutoModels schema."""
606 model_spec = _load_auto_models_model_spec_from_yaml(raw)
607 search_space, data_parallel_size, pp_degree = _load_auto_models_parallelism(raw)
608 training_raw = _get_dict(raw, "training")
609 activation_raw = _get_dict(raw, "activation_checkpoint")
611 global_batch_size = int(training_raw.get("global_batch_size", 0) or 0)
612 local_batch_size = int(model_spec["local_batch_size"] or 1)
613 micro_batch_num = (
614 global_batch_size // (local_batch_size * data_parallel_size)
615 if global_batch_size
616 and global_batch_size % (local_batch_size * data_parallel_size) == 0
617 else 1
618 )
620 recompute_strategy = _recompute_strategy(activation_raw.get("mode"), "activation_checkpoint.mode")
621 return NormalizedConfig(
622 model_spec=model_spec,
623 cluster_spec={},
624 search_space=search_space,
625 constraint={
626 "global_batch_size": global_batch_size,
627 "memory_limit_gb": 0.0,
628 },
629 estimator={
630 "type": "symbolic",
631 "recompute_strategy": recompute_strategy,
632 },
633 pp_config={
634 "pp_degree": pp_degree,
635 "stage_partition_mode": "uniform",
636 "micro_batch_num": max(1, micro_batch_num),
637 },
638 run=_stated_run(raw),
639 )
642def _build_config_from_hp_yaml(raw: Dict[str, Any]) -> NormalizedConfig:
643 """Construct a NormalizedConfig from a parsed HyperParallel YAML dict.
645 Extracts current AutoModels fields when root-level ``training`` and
646 ``accelerator`` sections are present. The legacy ``model.name`` /
647 ``model.config_overrides`` schema remains supported for compatibility.
649 Legacy parsing extracts model identifiers from ``model.config_overrides``,
650 parallelism from ``train.accelerator.*``, batch settings from ``train.*``,
651 sequence length from ``data.max_seq_len``, and recompute mode from
652 ``train.gradient_checkpointing.activation_checkpoint``.
654 Model hyperparameters are extracted from ``model.config_overrides``.
655 """
656 if is_auto_models_schema(raw):
657 return _build_config_from_auto_models_yaml(raw)
659 model_raw = _get_dict(raw, "model")
660 train_raw = _get_dict(raw, "train")
661 data_raw = _get_dict(raw, "data")
662 accel_raw = _get_dict(train_raw, "accelerator")
663 gc_raw = _get_dict(train_raw, "gradient_checkpointing")
665 # --- model_spec ---
666 model_spec: Dict[str, Any] = {}
667 model_spec["name"] = model_raw.get("name", "unknown")
668 overrides = model_raw.get("config_overrides", {})
669 if isinstance(overrides, dict):
670 model_spec.update(overrides)
671 model_spec["max_position_embeddings"] = data_raw.get("max_seq_len", 4096)
672 model_spec["local_batch_size"] = train_raw.get("micro_batch_size", 1)
674 # dtype from train.mixed_precision
675 mp_raw = _get_dict(train_raw, "mixed_precision")
676 if mp_raw.get("enabled", True):
677 model_spec["compute_dtype"] = mp_raw.get("param_dtype", "bfloat16")
679 # --- cluster_spec (users should set via search config or directly) ---
680 cluster_spec: Dict[str, Any] = {}
682 # --- search_space from train.accelerator ---
683 search_space: Dict[str, List[int]] = {}
684 for hkey, skey in _ACCEL_TO_SEARCH.items():
685 val = accel_raw.get(hkey)
686 if val is not None:
687 search_space[skey] = [int(val)]
689 # --- constraint from train ---
690 gbs = train_raw.get("global_batch_size", 0)
691 constraint: Dict[str, Any] = {
692 "global_batch_size": gbs or 0,
693 "memory_limit_gb": 0.0,
694 }
695 mb_num = int(gbs) // int(model_spec["local_batch_size"]) if gbs and model_spec.get("local_batch_size") else 1
697 # --- pp_config ---
698 pp_degree = accel_raw.get("pipeline_parallel_degree", 1)
699 pp_degree = max(1, int(pp_degree) if pp_degree else 1)
700 pp_config: Dict[str, Any] = {
701 "pp_degree": pp_degree,
702 "stage_partition_mode": "uniform",
703 "micro_batch_num": max(1, mb_num // pp_degree),
704 }
706 # --- estimator from gradient_checkpointing ---
707 estimator: Dict[str, Any] = {
708 "type": "symbolic",
709 "recompute_strategy": _recompute_strategy(
710 gc_raw.get("activation_checkpoint"), "train.gradient_checkpointing.activation_checkpoint"
711 ),
712 }
714 model_spec = _normalize_model_spec(model_spec)
716 return NormalizedConfig(
717 model_spec=model_spec,
718 cluster_spec=cluster_spec,
719 search_space=search_space,
720 constraint=constraint,
721 estimator=estimator,
722 pp_config=pp_config,
723 run=_stated_run(raw),
724 )
727def read_hp_yaml_config(path: str) -> NormalizedConfig:
728 """Read a HyperParallel YAML configuration file.
730 For the current AutoModels Trainer schema, model dimensions are resolved
731 from ``model.pretrained_model_name_or_path`` through the shared
732 Transformers config path. Parallelism is read from root-level
733 ``accelerator`` / ``fsdp_config`` sections. The legacy
734 ``model.config_overrides`` / ``train.accelerator`` schema remains
735 supported for standalone cost-model configurations.
737 .. note::
738 Cluster configuration is **not** present in ``train.yaml``.
739 To perform a full strategy search, use :func:`read_search_config`
740 which accepts cluster and search-space parameters.
742 See :func:`_build_config_from_hp_yaml` for the full list of recognised
743 YAML sections.
745 Args:
746 path: Path to the YAML config file (``.yaml`` or ``.yml``).
748 Returns:
749 A :class:`NormalizedConfig` instance.
751 Raises:
752 FileNotFoundError: If the file does not exist.
753 ValueError: If the file cannot be parsed.
754 ImportError: If PyYAML is not installed.
755 """
756 raw = _load_yaml(path)
757 return _build_config_from_hp_yaml(raw)