Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / config_adapter / _strategy_output.py: 92%
104 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"""Strategy output writer for auto parallel strategy search.
17Serializes :class:`NormalizedConfig` instances to JSON or YAML files,
18generates human-readable summaries, and provides a PPB (PR524)
19configuration stub.
20"""
22import json
23import logging
24import os
25from typing import Any, Dict
27import yaml # type: ignore[import-untyped]
29from hyper_parallel.auto_parallel.config_adapter._normalized_config import NormalizedConfig
32logger = logging.getLogger(__name__)
35def _resolve_pp_degree(pp_cfg: Dict) -> int:
36 """Resolve pp_degree to a scalar from pp_config."""
37 pp_degree_val = pp_cfg.get("pp_degree", 1)
38 if isinstance(pp_degree_val, list):
39 return pp_degree_val[0] if pp_degree_val else 1
40 return pp_degree_val
43def _write_json(data: Dict[str, Any], output_path: str) -> None:
44 """Write a dict as a pretty-printed JSON file."""
45 with open(output_path, "w", encoding="utf-8") as fh:
46 json.dump(data, fh, indent=2, default=str)
47 fh.write("\n")
50# Mapping from resolved strategy keys to HP YAML accelerator field names.
51_YAML_KEY_MAP: Dict[str, str] = {
52 "dp_shard": "dp_shard",
53 "dp_replicate": "dp_replicate",
54 "data_parallel_shard_degree": "dp_shard",
55 "data_parallel_replicate_degree": "dp_replicate",
56 "data_parallel_degree": "dp_replicate",
57 "dp": "dp_replicate",
58 "tensor_parallel_degree": "tp_degree",
59 "tp_degree": "tp_degree",
60 "tp": "tp_degree",
61 "pipeline_parallel_degree": "pipeline_parallel_degree",
62 "context_parallel_degree": "context_parallel_degree",
63 "expert_parallel_degree": "expert_parallel_degree",
64 "expert_tensor_parallel_degree": "expert_tensor_parallel_degree",
65 "pp_degree": "pipeline_parallel_degree",
66 "pp": "pipeline_parallel_degree",
67 "cp_degree": "context_parallel_degree",
68 "cp": "context_parallel_degree",
69 "ep_degree": "expert_parallel_degree",
70 "ep": "expert_parallel_degree",
71 "etp_degree": "expert_tensor_parallel_degree",
72 "etp": "expert_tensor_parallel_degree",
73}
76def _validate_strategy_and_yaml(
77 config: NormalizedConfig,
78 original_yaml_path: str,
79) -> None:
80 """Validate that resolved_strategy is set and the original YAML exists."""
81 if config.resolved_strategy is None:
82 raise ValueError(
83 "config.resolved_strategy is None — no strategy to write. "
84 "Set config.resolved_strategy first."
85 )
86 if not os.path.isfile(original_yaml_path):
87 raise FileNotFoundError(
88 f"Original YAML not found: {original_yaml_path}"
89 )
92def _load_yaml_to_inject(original_yaml_path: str) -> Dict[str, Any]:
93 """Load and validate the original YAML, ensuring train/accelerator exist."""
94 with open(original_yaml_path, "r", encoding="utf-8") as fh:
95 data = yaml.safe_load(fh)
96 if data is None or not isinstance(data, dict):
97 raise ValueError(
98 f"Original YAML {original_yaml_path} must contain a top-level mapping."
99 )
100 if "train" not in data or not isinstance(data["train"], dict):
101 data["train"] = {}
102 if "accelerator" not in data["train"] or not isinstance(data["train"]["accelerator"], dict):
103 data["train"]["accelerator"] = {}
104 return data
107def _inject_resolved_strategy(data: Dict[str, Any], resolved: Dict[str, Any]) -> None:
108 """Inject resolved strategy values into the YAML data dict."""
109 train = data["train"]
110 accel = train["accelerator"]
112 for src_key, dst_key in _YAML_KEY_MAP.items():
113 if src_key in resolved:
114 accel[dst_key] = int(resolved[src_key])
116 if "global_batch_size" in resolved:
117 train["global_batch_size"] = int(resolved["global_batch_size"])
119 if "micro_batch_num" in resolved:
120 train["micro_batch_num"] = int(resolved["micro_batch_num"])
123def _write_output_yaml(
124 data: Dict[str, Any],
125 output_path: str,
126 overwrite: bool,
127 original_yaml_path: str,
128) -> None:
129 """Write the data dict to the output YAML file."""
130 write_path = original_yaml_path if overwrite else output_path
132 parent_dir = os.path.dirname(os.path.abspath(write_path))
133 if parent_dir and not os.path.isdir(parent_dir):
134 os.makedirs(parent_dir, exist_ok=True)
136 with open(write_path, "w", encoding="utf-8") as fh:
137 yaml.dump(data, fh, default_flow_style=False, sort_keys=False)
139 logger.info(
140 "Resolved YAML written to %s (overwrite=%s)",
141 write_path, overwrite,
142 )
145def write_strategy_config(
146 config: NormalizedConfig,
147 output_path: str,
148 fmt: str = "json",
149) -> None:
150 """Write a normalized configuration to a file.
152 Args:
153 config: The normalized configuration to serialize.
154 output_path: Destination file path.
155 fmt: Output format, ``"json"`` (default) only.
157 Raises:
158 ValueError: If the output format is unsupported.
159 OSError: If the file cannot be written.
160 """
161 data = config.to_dict()
163 parent_dir = os.path.dirname(os.path.abspath(output_path))
164 if parent_dir and not os.path.isdir(parent_dir):
165 os.makedirs(parent_dir, exist_ok=True)
167 try:
168 if fmt == "json":
169 _write_json(data, output_path)
170 else:
171 raise ValueError(
172 f"Unsupported output format: {fmt!r}. "
173 "Supported format: 'json'"
174 )
175 except OSError as exc:
176 raise OSError(f"Failed to write config to {output_path}: {exc}") from exc
178 logger.info("Strategy config written to %s", output_path)
181def write_resolved_strategy(
182 config: NormalizedConfig,
183 output_path: str,
184 fmt: str = "json",
185) -> None:
186 """Write only the resolved strategy to a file.
188 This produces a compact config snippet suitable for consumption
189 by training scripts.
191 Args:
192 config: The normalized configuration (must have ``resolved_strategy`` set).
193 output_path: Destination file path.
194 fmt: Output format, ``"json"`` (default) only.
196 Raises:
197 ValueError: If ``resolved_strategy`` is ``None``.
198 """
199 if config.resolved_strategy is None:
200 raise ValueError("config.resolved_strategy is None — no strategy to write")
202 parent_dir = os.path.dirname(os.path.abspath(output_path))
203 if parent_dir and not os.path.isdir(parent_dir):
204 os.makedirs(parent_dir, exist_ok=True)
206 data = {"resolved_strategy": dict(config.resolved_strategy)}
208 try:
209 if fmt == "json":
210 _write_json(data, output_path)
211 else:
212 raise ValueError(
213 f"Unsupported output format: {fmt!r}. "
214 "Supported format: 'json'"
215 )
216 except OSError as exc:
217 raise OSError(f"Failed to write resolved strategy to {output_path}: {exc}") from exc
219 logger.info("Resolved strategy written to %s", output_path)
222def write_ppb_config(
223 config: NormalizedConfig,
224 output_path: str,
225) -> None:
226 """Write a PPB (PR524 pipeline balancer) input JSON file.
228 Writes the PP-relevant fields from ``NormalizedConfig``
229 in a structure compatible with ``args_for_pipeline_parallel.json``.
231 Args:
232 config: The normalized configuration.
233 output_path: Destination JSON file path.
235 Raises:
236 OSError: If the file cannot be written.
237 """
238 model = config.model_spec
239 pp_cfg = config.pp_config
240 constraint = config.constraint
242 ppb_data = {
243 "llm_class": "0",
244 "train_yaml": "",
245 "mindformers_dir": "",
246 "layer_ratio": 0.33,
247 "backward_ratio": 2.0,
248 "head_loss": 1.5,
249 "recompute_ratio": 1,
250 "time_limit": 2147483647,
251 "dryrun": True,
252 "check": True,
253 "fit_level": 0,
254 "extract": False,
255 "env_json": "./config/env_config.json",
256 "dryrun_lim": 16,
257 "_hyper_model": {
258 "n_layers": model.get("n_layers", 0),
259 "dim": model.get("dim", 0),
260 "n_heads": model.get("n_heads", 0),
261 "vocab_size": model.get("vocab_size", 0),
262 "seq_len": model.get("seq_len", 4096),
263 "moe_enabled": model.get("moe_enabled", False),
264 "num_experts": model.get("num_experts", 1),
265 },
266 "_hyper_pp": {
267 "pp_degree": _resolve_pp_degree(pp_cfg),
268 "stage_partition_mode": pp_cfg.get("stage_partition_mode", "uniform"),
269 "micro_batch_num": pp_cfg.get("micro_batch_num", 1),
270 "layer_offset_range": list(pp_cfg.get("layer_offset_range", (0, 0))),
271 "layer_recompute_layers": pp_cfg.get("layer_recompute_layers", []),
272 },
273 "_hyper_constraint": {
274 "global_batch_size": constraint.get("global_batch_size", 0),
275 "memory_limit_gb": constraint.get("memory_limit_gb", 0.0),
276 },
277 }
279 parent_dir = os.path.dirname(os.path.abspath(output_path))
280 if parent_dir and not os.path.isdir(parent_dir):
281 os.makedirs(parent_dir, exist_ok=True)
283 try:
284 with open(output_path, "w", encoding="utf-8") as fh:
285 json.dump(ppb_data, fh, indent=4)
286 fh.write("\n")
287 except OSError as exc:
288 raise OSError(f"Failed to write PPB config to {output_path}: {exc}") from exc
290 logger.info("PPB config stub written to %s", output_path)
293def write_resolved_yaml(
294 config: NormalizedConfig,
295 original_yaml_path: str,
296 output_path: str,
297 overwrite: bool = False,
298) -> None:
299 """Write a resolved strategy into a HyperParallel YAML file.
301 Copies the original ``train.yaml`` and replaces the parallel
302 dimension fields with the resolved strategy values. This produces
303 a complete, immediately launchable training configuration.
305 The resolved strategy is read from ``config.resolved_strategy``.
306 Supported keys: ``dp_shard``, ``dp_replicate``, ``tp_degree``,
307 ``pipeline_parallel_degree``, ``context_parallel_degree``,
308 ``expert_parallel_degree``, ``global_batch_size``.
310 Args:
311 config: NormalizedConfig with ``resolved_strategy`` set.
312 original_yaml_path: Path to the original ``train.yaml``.
313 output_path: Destination file path for the resolved YAML.
314 overwrite: If ``True``, write to ``original_yaml_path`` instead
315 of ``output_path`` (default ``False``).
317 Raises:
318 ValueError: If ``resolved_strategy`` is ``None``.
319 FileNotFoundError: If ``original_yaml_path`` does not exist.
320 """
321 _validate_strategy_and_yaml(config, original_yaml_path)
322 data = _load_yaml_to_inject(original_yaml_path)
323 _inject_resolved_strategy(data, config.resolved_strategy)
324 _write_output_yaml(data, output_path, overwrite, original_yaml_path)
327def normalized_to_summary(config: NormalizedConfig) -> Dict[str, Any]:
328 """Generate a human-readable summary of the configuration.
330 Args:
331 config: The normalized configuration to summarize.
333 Returns:
334 A dictionary with summary fields suitable for logging or display.
335 """
336 model = config.model_spec
337 cluster = config.cluster_spec
338 search = config.search_space
339 constraint = config.constraint
340 estimator = config.estimator
341 pp_cfg = config.pp_config
343 return {
344 "model": {
345 "name": model.get("name", "unknown"),
346 "n_layers": model.get("n_layers", 0),
347 "dim": model.get("dim", 0),
348 "inter_dim": model.get("inter_dim", 0),
349 "n_heads": model.get("n_heads", 0),
350 "n_kv_heads": model.get("n_kv_heads", 0),
351 "vocab_size": model.get("vocab_size", 0),
352 "seq_len": model.get("seq_len", 0),
353 "moe_enabled": model.get("moe_enabled", False),
354 "num_experts": model.get("num_experts", 0),
355 },
356 "cluster": {
357 "num_nodes": cluster.get("num_nodes", 0),
358 "cards_per_node": cluster.get("cards_per_node", 0),
359 "total_cards": (
360 cluster.get("num_nodes", 0) * cluster.get("cards_per_node", 0)
361 ),
362 "device_memory_gb": cluster.get("device_memory_gb", 0),
363 "device_type": cluster.get("device_type", "unknown"),
364 },
365 "search_space": dict(sorted(search.items())),
366 "constraints": {
367 "global_batch_size": constraint.get("global_batch_size", 0),
368 "memory_limit_gb": constraint.get("memory_limit_gb", 0.0),
369 "fixed_dimensions": {
370 k: constraint.get(f"fixed_{k}_degree")
371 for k in ("dp", "tp", "pp", "cp", "ep")
372 if constraint.get(f"fixed_{k}_degree") is not None
373 },
374 },
375 "estimator": {
376 "type": estimator.get("type", "symbolic"),
377 "recompute_strategy": estimator.get("recompute_strategy", "none"),
378 },
379 "pipeline": {
380 "pp_degree": _resolve_pp_degree(pp_cfg),
381 "stage_partition_mode": pp_cfg.get("stage_partition_mode", "uniform"),
382 "micro_batch_num": pp_cfg.get("micro_batch_num", 1),
383 },
384 "resolved_strategy": config.resolved_strategy,
385 }