Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / config_adapter / _config_loader.py: 96%
124 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"""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 logging
22import os
23from typing import Any, Dict, List, Tuple, Optional
25try:
26 import yaml # type: ignore[import-untyped] # pylint: disable=C0415
27except ImportError:
28 yaml = None # pragma: no cover
30from hyper_parallel.auto_parallel.config_adapter._normalized_config import NormalizedConfig
32logger = logging.getLogger(__name__)
34# Mapping from HuggingFace-style config_overrides keys to internal short names.
35_HP_TO_INTERNAL: Dict[str, str] = {
36 "num_hidden_layers": "n_layers",
37 "hidden_size": "dim",
38 "num_attention_heads": "n_heads",
39 "intermediate_size": "inter_dim",
40 "num_key_value_heads": "n_kv_heads",
41 "max_position_embeddings": "seq_len",
42 "seq_length": "seq_len",
43}
46def _normalize_model_spec(model_spec: Dict[str, Any]) -> Dict[str, Any]:
47 """Rename HuggingFace-style config overrides keys to internal short names.
49 Only maps a key if the target short name is not already present, so
50 explicit short names in the YAML take precedence.
51 """
52 for hf_key, internal_key in _HP_TO_INTERNAL.items():
53 if hf_key in model_spec and internal_key not in model_spec:
54 model_spec[internal_key] = model_spec.pop(hf_key)
55 return model_spec
58# Mapping from short dimension names (used in search config YAML parallelism section)
59# to canonical NormalizedConfig search_space keys.
60_UNIFIED_DIM_MAP: Dict[str, str] = {
61 "dp": "data_parallel_replicate_degree",
62 "fsdp": "data_parallel_shard_degree",
63 "tp": "tensor_parallel_degree",
64 "pp": "pipeline_parallel_degree",
65 "cp": "context_parallel_degree",
66 "ep": "expert_parallel_degree",
67 "etp": "expert_tensor_parallel_degree",
68 "micro_batch_num": "micro_batch_num",
69}
71# Mapping from short dimension names to constraint fixed_*_degree keys.
72_FIXED_DIM_MAP: Dict[str, str] = {
73 "dp": "fixed_dp_degree",
74 "fsdp": "fixed_fsdp_degree",
75 "tp": "fixed_tp_degree",
76 "pp": "fixed_pp_degree",
77 "cp": "fixed_cp_degree",
78 "ep": "fixed_ep_degree",
79 "micro_batch_num": "fixed_micro_batch_num",
80}
83def _get_dict(raw: Dict[str, Any], key: str) -> Dict[str, Any]:
84 """Return the value of a key if it is a dict, otherwise an empty dict."""
85 val = raw.get(key, {})
86 return val if isinstance(val, dict) else {}
89def _load_yaml(path: str) -> Dict[str, Any]:
90 """Read and parse a YAML file, returning the raw dict."""
91 if yaml is None:
92 raise ImportError(
93 "PyYAML is required to read HyperParallel YAML configs. "
94 "Install it with: pip install pyyaml"
95 )
96 if not os.path.isfile(path):
97 raise FileNotFoundError(f"Config file not found: {path}")
99 ext = os.path.splitext(path)[1].lower()
100 if ext not in (".yaml", ".yml"):
101 raise ValueError(
102 f"Unsupported config file format: {ext!r}. "
103 "Supported formats: .yaml, .yml"
104 )
106 try:
107 with open(path, "r", encoding="utf-8") as fh:
108 raw = yaml.safe_load(fh)
109 except yaml.YAMLError as exc:
110 raise ValueError(f"Failed to parse YAML file {path}: {exc}") from exc
112 if raw is None:
113 raw = {}
114 if not isinstance(raw, dict):
115 raise ValueError(
116 f"Config file {path} must contain a YAML mapping at the top level, "
117 f"got {type(raw).__name__}"
118 )
119 return raw
122# ── Search Config YAML reader (primary) ──────────────────────────────
125def _parse_unified_parallelism(
126 para_raw: Dict[str, Any],
127) -> Tuple[Dict[str, List[int]], Dict[str, Any]]:
128 """Convert the unified parallelism declaration into search_space + constraint.
130 Rules:
131 * Scalar integer value → fixed dimension (both ``constraint.fixed_*``
132 and ``search_space`` with a single-element list).
133 * List value → search candidates (placed into ``search_space`` only).
134 * String ``"auto"`` → dimension is left to the searcher's ``bound_space``
135 (neither fixed nor explicitly enumerated).
137 Returns:
138 A ``(search_space, constraint)`` tuple.
139 """
140 search_space: Dict[str, List[int]] = {}
141 constraint: Dict[str, Any] = {}
143 for short_key, canonical_key in _UNIFIED_DIM_MAP.items():
144 if short_key not in para_raw:
145 continue
146 value = para_raw[short_key]
148 if isinstance(value, int):
149 constraint[_FIXED_DIM_MAP[short_key]] = value
150 search_space[canonical_key] = [value]
151 elif isinstance(value, list):
152 search_space[canonical_key] = [int(v) for v in value]
153 elif isinstance(value, str) and value.strip().lower() == "auto":
154 continue
156 return search_space, constraint
159def _build_config_from_search_yaml(raw: Dict[str, Any]) -> NormalizedConfig:
160 """Construct a NormalizedConfig from a parsed Search Config YAML dict.
162 Supports two modes:
164 * **Standalone** (no ``train_yaml``) — ``model`` section and all other
165 info must be present in the search config file.
166 * **With train_yaml** — loads the specified ``train.yaml`` for model
167 parameters and current parallelism values, then overlays the search
168 config's ``cluster``, ``parallelism``, and ``constraint`` sections.
170 Undeclared parallelism dimensions are inherited from ``train.yaml``
171 as fixed (single-element) entries.
172 """
173 train_yaml_path = raw.get("train_yaml")
174 base_config: Optional[NormalizedConfig] = None
176 if train_yaml_path:
177 if not isinstance(train_yaml_path, str):
178 raise ValueError("'train_yaml' must be a file path string")
179 base_raw = _load_yaml(train_yaml_path)
180 base_config = _build_config_from_hp_yaml(base_raw)
182 model_spec: Dict[str, Any]
183 if base_config:
184 model_spec = dict(base_config.model_spec)
185 else:
186 model_spec = {}
188 # Override or supply model section from search.yaml
189 search_model = _get_dict(raw, "model")
190 if search_model:
191 model_spec.update(search_model)
193 model_spec.setdefault("seq_len", 4096)
194 model_spec.setdefault("local_batch_size", 1)
196 model_spec = _normalize_model_spec(model_spec)
198 cluster_spec = _get_dict(raw, "cluster")
200 pp_raw = _get_dict(raw, "pp_config")
201 parallelism_raw = _get_dict(raw, "parallelism")
202 constraint_raw = _get_dict(raw, "constraint")
204 search_space, parallelism_constraint = _parse_unified_parallelism(parallelism_raw)
206 # Inherit undeclared dimensions from train.yaml as fixed values.
207 if base_config:
208 for space_key, candidates in base_config.search_space.items():
209 if space_key not in search_space:
210 search_space[space_key] = candidates
212 if constraint_raw.get("global_batch_size", 0) is None or constraint_raw.get("global_batch_size", 0) == 0:
213 constraint_raw.setdefault(
214 "global_batch_size", base_config.constraint.get("global_batch_size", 0)
215 )
217 pp_config: Dict[str, Any] = {
218 "pp_degree": pp_raw.get("pp_degree",
219 parallelism_raw.get("pp", 1)),
220 "stage_partition_mode": pp_raw.get("stage_partition_mode", "uniform"),
221 "stage_partition": pp_raw.get("stage_partition", []),
222 "layer_offset_range": tuple(pp_raw.get("layer_offset_range", [0, 0])),
223 "layer_recompute_layers": pp_raw.get("layer_recompute_layers", []),
224 "micro_batch_num": pp_raw.get("micro_batch_num", 1),
225 "pp_interleave_num": pp_raw.get("pp_interleave_num", 1),
226 "pipeline_parallel_schedule": pp_raw.get("pipeline_schedule", "1F1B"),
227 }
229 estimator: Dict[str, Any] = {
230 "type": "symbolic",
231 "recompute_strategy": str(raw.get("recompute", "none")),
232 "enable_profiling_calibration": False,
233 }
235 constraint: Dict[str, Any] = {
236 "global_batch_size": constraint_raw.get("global_batch_size", 0),
237 "memory_limit_gb": constraint_raw.get("memory_limit_gb", 0.0),
238 **parallelism_constraint,
239 }
241 return NormalizedConfig(
242 model_spec=model_spec,
243 cluster_spec=cluster_spec,
244 search_space=search_space,
245 constraint=constraint,
246 estimator=estimator,
247 pp_config=pp_config,
248 )
251def read_search_config(path: str) -> NormalizedConfig:
252 """Read a Search Config YAML file and return a :class:`NormalizedConfig`.
254 The Search Config YAML format uses a unified ``parallelism`` section
255 where each dimension is declared as::
257 parallelism:
258 tp: 4 # scalar → fixed input
259 dp: [1, 2, 4] # list → search candidate
260 pp: auto # string → let the searcher decide
262 To reuse model parameters from an existing ``train.yaml`` without
263 duplicating them, set the ``train_yaml`` key::
265 train_yaml: "./train.yaml" # load model params from here
266 cluster:
267 num_nodes: 4
268 cards_per_node: 8
269 parallelism:
270 dp: [1, 2, 4]
271 tp: [1, 2, 4, 8]
273 Dimensions absent from ``parallelism`` are inherited from
274 ``train.yaml`` as fixed values. A ``model`` section in the search
275 config overrides values read from ``train_yaml``.
277 See ``auto_parallel/examples/dense_llm_search.yaml`` for a complete
278 standalone example.
280 Args:
281 path: Path to the YAML config file (``.yaml`` or ``.yml``).
283 Returns:
284 A :class:`NormalizedConfig` instance.
286 Raises:
287 FileNotFoundError: If the file does not exist.
288 ValueError: If the file cannot be parsed, or if ``cluster``
289 is missing when ``train_yaml`` is not used.
290 ImportError: If PyYAML is not installed.
291 """
292 raw = _load_yaml(path)
293 return _build_config_from_search_yaml(raw)
296# ── HyperParallel training YAML reader (secondary) ──────────────────
298_ACCEL_TO_SEARCH = {
299 "dp_shard": "data_parallel_shard_degree",
300 "dp_replicate": "data_parallel_replicate_degree",
301 "tp_degree": "tensor_parallel_degree",
302 "pipeline_parallel_degree": "pipeline_parallel_degree",
303 "context_parallel_degree": "context_parallel_degree",
304 "expert_parallel_degree": "expert_parallel_degree",
305 "expert_tensor_parallel_degree": "expert_tensor_parallel_degree",
306}
309def _build_config_from_hp_yaml(raw: Dict[str, Any]) -> NormalizedConfig:
310 """Construct a NormalizedConfig from a parsed HyperParallel YAML dict.
312 Extracts model identifiers from ``model.name`` / ``model.config_overrides``,
313 parallelism from ``train.accelerator.*``, batch settings from ``train.*``,
314 sequence length from ``data.max_seq_len``, and recompute mode from
315 ``train.gradient_checkpointing.activation_checkpoint``.
317 Model hyperparameters are extracted from ``model.config_overrides``.
318 """
319 model_raw = _get_dict(raw, "model")
320 train_raw = _get_dict(raw, "train")
321 data_raw = _get_dict(raw, "data")
322 accel_raw = _get_dict(train_raw, "accelerator")
323 gc_raw = _get_dict(train_raw, "gradient_checkpointing")
325 # --- model_spec ---
326 model_spec: Dict[str, Any] = {}
327 model_spec["name"] = model_raw.get("name", "unknown")
328 overrides = model_raw.get("config_overrides", {})
329 if isinstance(overrides, dict):
330 model_spec.update(overrides)
331 model_spec["seq_len"] = data_raw.get("max_seq_len", 4096)
332 model_spec["local_batch_size"] = train_raw.get("micro_batch_size", 1)
334 # dtype from train.mixed_precision
335 mp_raw = _get_dict(train_raw, "mixed_precision")
336 if mp_raw.get("enabled", True):
337 model_spec["compute_dtype"] = mp_raw.get("param_dtype", "bfloat16")
339 # --- cluster_spec (users should set via search config or directly) ---
340 cluster_spec: Dict[str, Any] = {}
342 # --- search_space from train.accelerator ---
343 search_space: Dict[str, List[int]] = {}
344 for hkey, skey in _ACCEL_TO_SEARCH.items():
345 val = accel_raw.get(hkey)
346 if val is not None:
347 search_space[skey] = [int(val)]
349 # --- constraint from train ---
350 gbs = train_raw.get("global_batch_size", 0)
351 constraint: Dict[str, Any] = {
352 "global_batch_size": gbs or 0,
353 "memory_limit_gb": 0.0,
354 }
355 mb_num = int(gbs) // int(model_spec["local_batch_size"]) if gbs and model_spec.get("local_batch_size") else 1
357 # --- pp_config ---
358 pp_degree = accel_raw.get("pipeline_parallel_degree", 1)
359 pp_degree = max(1, int(pp_degree) if pp_degree else 1)
360 pp_config: Dict[str, Any] = {
361 "pp_degree": pp_degree,
362 "stage_partition_mode": "uniform",
363 "micro_batch_num": max(1, mb_num // pp_degree),
364 }
366 # --- estimator from gradient_checkpointing ---
367 ac_mode = str(gc_raw.get("activation_checkpoint", "none"))
368 recompute_map = {"none": "none", "full": "full", "selective": "selective"}
369 estimator: Dict[str, Any] = {
370 "type": "symbolic",
371 "recompute_strategy": recompute_map.get(ac_mode, "none"),
372 }
374 model_spec = _normalize_model_spec(model_spec)
376 return NormalizedConfig(
377 model_spec=model_spec,
378 cluster_spec=cluster_spec,
379 search_space=search_space,
380 constraint=constraint,
381 estimator=estimator,
382 pp_config=pp_config,
383 )
386def read_hp_yaml_config(path: str) -> NormalizedConfig:
387 """Read a HyperParallel YAML configuration file.
389 This is a convenience reader for the native HyperParallel ``train.yaml``
390 format. It extracts parallelism from ``train.accelerator`` and model
391 fields from ``model.config_overrides``.
393 .. note::
394 Cluster configuration is **not** present in ``train.yaml``.
395 To perform a full strategy search, use :func:`read_search_config`
396 which accepts cluster and search-space parameters.
398 See :func:`_build_config_from_hp_yaml` for the full list of recognised
399 YAML sections.
401 Args:
402 path: Path to the YAML config file (``.yaml`` or ``.yml``).
404 Returns:
405 A :class:`NormalizedConfig` instance.
407 Raises:
408 FileNotFoundError: If the file does not exist.
409 ValueError: If the file cannot be parsed.
410 ImportError: If PyYAML is not installed.
411 """
412 raw = _load_yaml(path)
413 return _build_config_from_hp_yaml(raw)