Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / config_adapter / _config_loader.py: 95%
124 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +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 names.
35# After field-name alignment, most HP YAML keys already match the internal
36# model_spec names. Only ``seq_length`` (MF variant) needs remapping.
37_HP_TO_INTERNAL: Dict[str, str] = {
38 "seq_length": "max_position_embeddings",
39}
42def _normalize_model_spec(model_spec: Dict[str, Any]) -> Dict[str, Any]:
43 """Rename non-standard config overrides keys to canonical names.
45 Only maps a key if the target name is not already present, so
46 explicit canonical names in the YAML take precedence.
47 """
48 for hf_key, internal_key in _HP_TO_INTERNAL.items():
49 if hf_key in model_spec and internal_key not in model_spec:
50 model_spec[internal_key] = model_spec.pop(hf_key)
51 return model_spec
54# Mapping from short dimension names (used in search config YAML parallelism section)
55# to canonical NormalizedConfig search_space keys.
56_UNIFIED_DIM_MAP: Dict[str, str] = {
57 "dp": "data_parallel_replicate_degree",
58 "fsdp": "data_parallel_shard_degree",
59 "tp": "tensor_parallel_degree",
60 "pp": "pipeline_parallel_degree",
61 "cp": "context_parallel_degree",
62 "ep": "expert_parallel_degree",
63 "etp": "expert_tensor_parallel_degree",
64 "micro_batch_num": "micro_batch_num",
65}
67# Mapping from short dimension names to constraint fixed_*_degree keys.
68_FIXED_DIM_MAP: Dict[str, str] = {
69 "dp": "fixed_dp_degree",
70 "fsdp": "fixed_fsdp_degree",
71 "tp": "fixed_tp_degree",
72 "pp": "fixed_pp_degree",
73 "cp": "fixed_cp_degree",
74 "ep": "fixed_ep_degree",
75 "etp": "fixed_etp_degree",
76 "micro_batch_num": "fixed_micro_batch_num",
77}
80def _get_dict(raw: Dict[str, Any], key: str) -> Dict[str, Any]:
81 """Return the value of a key if it is a dict, otherwise an empty dict."""
82 val = raw.get(key, {})
83 return val if isinstance(val, dict) else {}
86def _load_yaml(path: str) -> Dict[str, Any]:
87 """Read and parse a YAML file, returning the raw dict."""
88 if yaml is None:
89 raise ImportError(
90 "PyYAML is required to read HyperParallel YAML configs. "
91 "Install it with: pip install pyyaml"
92 )
93 if not os.path.isfile(path):
94 raise FileNotFoundError(f"Config file not found: {path}")
96 ext = os.path.splitext(path)[1].lower()
97 if ext not in (".yaml", ".yml"):
98 raise ValueError(
99 f"Unsupported config file format: {ext!r}. "
100 "Supported formats: .yaml, .yml"
101 )
103 try:
104 with open(path, "r", encoding="utf-8") as fh:
105 raw = yaml.safe_load(fh)
106 except yaml.YAMLError as exc:
107 raise ValueError(f"Failed to parse YAML file {path}: {exc}") from exc
109 if raw is None:
110 raw = {}
111 if not isinstance(raw, dict):
112 raise ValueError(
113 f"Config file {path} must contain a YAML mapping at the top level, "
114 f"got {type(raw).__name__}"
115 )
116 return raw
119# ── Search Config YAML reader (primary) ──────────────────────────────
122def _parse_unified_parallelism(
123 para_raw: Dict[str, Any],
124) -> Tuple[Dict[str, List[int]], Dict[str, Any]]:
125 """Convert the unified parallelism declaration into search_space + constraint.
127 Rules:
128 * Scalar integer value → fixed dimension (both ``constraint.fixed_*``
129 and ``search_space`` with a single-element list).
130 * List value → search candidates (placed into ``search_space`` only).
131 * String ``"auto"`` → dimension is left to the searcher's ``bound_space``
132 (neither fixed nor explicitly enumerated).
134 Returns:
135 A ``(search_space, constraint)`` tuple.
136 """
137 search_space: Dict[str, List[int]] = {}
138 constraint: Dict[str, Any] = {}
140 for short_key, canonical_key in _UNIFIED_DIM_MAP.items():
141 if short_key not in para_raw:
142 continue
143 value = para_raw[short_key]
145 if isinstance(value, int):
146 constraint[_FIXED_DIM_MAP[short_key]] = value
147 search_space[canonical_key] = [value]
148 elif isinstance(value, list):
149 search_space[canonical_key] = [int(v) for v in value]
150 elif isinstance(value, str) and value.strip().lower() == "auto":
151 continue
153 return search_space, constraint
156def _build_config_from_search_yaml(raw: Dict[str, Any]) -> NormalizedConfig:
157 """Construct a NormalizedConfig from a parsed Search Config YAML dict.
159 Supports two modes:
161 * **Standalone** (no ``train_yaml``) — ``model`` section and all other
162 info must be present in the search config file.
163 * **With train_yaml** — loads the specified ``train.yaml`` for model
164 parameters and current parallelism values, then overlays the search
165 config's ``cluster``, ``parallelism``, and ``constraint`` sections.
167 Undeclared parallelism dimensions are inherited from ``train.yaml``
168 as fixed (single-element) entries.
169 """
170 train_yaml_path = raw.get("train_yaml")
171 base_config: Optional[NormalizedConfig] = None
173 if train_yaml_path:
174 if not isinstance(train_yaml_path, str):
175 raise ValueError("'train_yaml' must be a file path string")
176 base_raw = _load_yaml(train_yaml_path)
177 base_config = _build_config_from_hp_yaml(base_raw)
179 model_spec: Dict[str, Any]
180 if base_config:
181 model_spec = dict(base_config.model_spec)
182 else:
183 model_spec = {}
185 # Override or supply model section from search.yaml
186 search_model = _get_dict(raw, "model")
187 if search_model:
188 model_spec.update(search_model)
190 model_spec.setdefault("max_position_embeddings", 4096)
191 model_spec.setdefault("local_batch_size", 1)
193 model_spec = _normalize_model_spec(model_spec)
195 cluster_spec = _get_dict(raw, "cluster")
197 pp_raw = _get_dict(raw, "pp_config")
198 parallelism_raw = _get_dict(raw, "parallelism")
199 constraint_raw = _get_dict(raw, "constraint")
201 search_space, parallelism_constraint = _parse_unified_parallelism(parallelism_raw)
203 # Inherit undeclared dimensions from train.yaml as fixed values.
204 if base_config:
205 for space_key, candidates in base_config.search_space.items():
206 if space_key not in search_space:
207 search_space[space_key] = candidates
209 if constraint_raw.get("global_batch_size", 0) is None or constraint_raw.get("global_batch_size", 0) == 0:
210 constraint_raw.setdefault(
211 "global_batch_size", base_config.constraint.get("global_batch_size", 0)
212 )
214 pp_config: Dict[str, Any] = {
215 "pp_degree": pp_raw.get("pp_degree",
216 parallelism_raw.get("pp", 1)),
217 "stage_partition_mode": pp_raw.get("stage_partition_mode", "uniform"),
218 "stage_partition": pp_raw.get("stage_partition", []),
219 "layer_offset_range": tuple(pp_raw.get("layer_offset_range", [0, 0])),
220 "layer_recompute_layers": pp_raw.get("layer_recompute_layers", []),
221 "micro_batch_num": pp_raw.get("micro_batch_num", 1),
222 "pp_interleave_num": pp_raw.get("pp_interleave_num", 1),
223 "pipeline_parallel_schedule": pp_raw.get("pipeline_schedule", "1F1B"),
224 }
226 estimator: Dict[str, Any] = {
227 "type": "symbolic",
228 "recompute_strategy": str(raw.get("recompute", "none")),
229 "enable_profiling_calibration": False,
230 }
232 constraint: Dict[str, Any] = {
233 "global_batch_size": constraint_raw.get("global_batch_size", 0),
234 "memory_limit_gb": constraint_raw.get("memory_limit_gb", 0.0),
235 **parallelism_constraint,
236 }
238 return NormalizedConfig(
239 model_spec=model_spec,
240 cluster_spec=cluster_spec,
241 search_space=search_space,
242 constraint=constraint,
243 estimator=estimator,
244 pp_config=pp_config,
245 )
248def read_search_config(path: str) -> NormalizedConfig:
249 """Read a Search Config YAML file and return a :class:`NormalizedConfig`.
251 The Search Config YAML format uses a unified ``parallelism`` section
252 where each dimension is declared as::
254 parallelism:
255 tp: 4 # scalar → fixed input
256 dp: [1, 2, 4] # list → search candidate
257 pp: auto # string → let the searcher decide
259 To reuse model parameters from an existing ``train.yaml`` without
260 duplicating them, set the ``train_yaml`` key::
262 train_yaml: "./train.yaml" # load model params from here
263 cluster:
264 num_nodes: 4
265 cards_per_node: 8
266 parallelism:
267 dp: [1, 2, 4]
268 tp: [1, 2, 4, 8]
270 Dimensions absent from ``parallelism`` are inherited from
271 ``train.yaml`` as fixed values. A ``model`` section in the search
272 config overrides values read from ``train_yaml``.
274 See ``auto_parallel/examples/dense_llm_search.yaml`` for a complete
275 standalone example.
277 Args:
278 path: Path to the YAML config file (``.yaml`` or ``.yml``).
280 Returns:
281 A :class:`NormalizedConfig` instance.
283 Raises:
284 FileNotFoundError: If the file does not exist.
285 ValueError: If the file cannot be parsed, or if ``cluster``
286 is missing when ``train_yaml`` is not used.
287 ImportError: If PyYAML is not installed.
288 """
289 raw = _load_yaml(path)
290 return _build_config_from_search_yaml(raw)
293# ── HyperParallel training YAML reader (secondary) ──────────────────
295_ACCEL_TO_SEARCH = {
296 "dp_shard": "data_parallel_shard_degree",
297 "dp_replicate": "data_parallel_replicate_degree",
298 "tp_degree": "tensor_parallel_degree",
299 "pipeline_parallel_degree": "pipeline_parallel_degree",
300 "context_parallel_degree": "context_parallel_degree",
301 "expert_parallel_degree": "expert_parallel_degree",
302 "expert_tensor_parallel_degree": "expert_tensor_parallel_degree",
303}
306def _build_config_from_hp_yaml(raw: Dict[str, Any]) -> NormalizedConfig:
307 """Construct a NormalizedConfig from a parsed HyperParallel YAML dict.
309 Extracts model identifiers from ``model.name`` / ``model.config_overrides``,
310 parallelism from ``train.accelerator.*``, batch settings from ``train.*``,
311 sequence length from ``data.max_seq_len``, and recompute mode from
312 ``train.gradient_checkpointing.activation_checkpoint``.
314 Model hyperparameters are extracted from ``model.config_overrides``.
315 """
316 model_raw = _get_dict(raw, "model")
317 train_raw = _get_dict(raw, "train")
318 data_raw = _get_dict(raw, "data")
319 accel_raw = _get_dict(train_raw, "accelerator")
320 gc_raw = _get_dict(train_raw, "gradient_checkpointing")
322 # --- model_spec ---
323 model_spec: Dict[str, Any] = {}
324 model_spec["name"] = model_raw.get("name", "unknown")
325 overrides = model_raw.get("config_overrides", {})
326 if isinstance(overrides, dict):
327 model_spec.update(overrides)
328 model_spec["max_position_embeddings"] = data_raw.get("max_seq_len", 4096)
329 model_spec["local_batch_size"] = train_raw.get("micro_batch_size", 1)
331 # dtype from train.mixed_precision
332 mp_raw = _get_dict(train_raw, "mixed_precision")
333 if mp_raw.get("enabled", True):
334 model_spec["compute_dtype"] = mp_raw.get("param_dtype", "bfloat16")
336 # --- cluster_spec (users should set via search config or directly) ---
337 cluster_spec: Dict[str, Any] = {}
339 # --- search_space from train.accelerator ---
340 search_space: Dict[str, List[int]] = {}
341 for hkey, skey in _ACCEL_TO_SEARCH.items():
342 val = accel_raw.get(hkey)
343 if val is not None:
344 search_space[skey] = [int(val)]
346 # --- constraint from train ---
347 gbs = train_raw.get("global_batch_size", 0)
348 constraint: Dict[str, Any] = {
349 "global_batch_size": gbs or 0,
350 "memory_limit_gb": 0.0,
351 }
352 mb_num = int(gbs) // int(model_spec["local_batch_size"]) if gbs and model_spec.get("local_batch_size") else 1
354 # --- pp_config ---
355 pp_degree = accel_raw.get("pipeline_parallel_degree", 1)
356 pp_degree = max(1, int(pp_degree) if pp_degree else 1)
357 pp_config: Dict[str, Any] = {
358 "pp_degree": pp_degree,
359 "stage_partition_mode": "uniform",
360 "micro_batch_num": max(1, mb_num // pp_degree),
361 }
363 # --- estimator from gradient_checkpointing ---
364 ac_mode = str(gc_raw.get("activation_checkpoint", "none"))
365 recompute_map = {"none": "none", "full": "full", "selective": "selective"}
366 estimator: Dict[str, Any] = {
367 "type": "symbolic",
368 "recompute_strategy": recompute_map.get(ac_mode, "none"),
369 }
371 model_spec = _normalize_model_spec(model_spec)
373 return NormalizedConfig(
374 model_spec=model_spec,
375 cluster_spec=cluster_spec,
376 search_space=search_space,
377 constraint=constraint,
378 estimator=estimator,
379 pp_config=pp_config,
380 )
383def read_hp_yaml_config(path: str) -> NormalizedConfig:
384 """Read a HyperParallel YAML configuration file.
386 This is a convenience reader for the native HyperParallel ``train.yaml``
387 format. It extracts parallelism from ``train.accelerator`` and model
388 fields from ``model.config_overrides``.
390 .. note::
391 Cluster configuration is **not** present in ``train.yaml``.
392 To perform a full strategy search, use :func:`read_search_config`
393 which accepts cluster and search-space parameters.
395 See :func:`_build_config_from_hp_yaml` for the full list of recognised
396 YAML sections.
398 Args:
399 path: Path to the YAML config file (``.yaml`` or ``.yml``).
401 Returns:
402 A :class:`NormalizedConfig` instance.
404 Raises:
405 FileNotFoundError: If the file does not exist.
406 ValueError: If the file cannot be parsed.
407 ImportError: If PyYAML is not installed.
408 """
409 raw = _load_yaml(path)
410 return _build_config_from_hp_yaml(raw)