Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / config_adapter / _search_runner.py: 95%
129 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"""Search runner — bridges NormalizedConfig to the ND search engine.
17Converts a :class:`NormalizedConfig` into a temporary HyperParallel
18``train.yaml``, runs the ND search via :class:`Parallelize`,
19post-filters by user candidate lists, and returns the optimal strategy.
20"""
22import logging
23import os
24import tempfile
25from typing import Any, Dict, List, TYPE_CHECKING
27import yaml # type: ignore[import-untyped]
29from hyper_parallel.auto_parallel.config_adapter._normalized_config import NormalizedConfig
31if TYPE_CHECKING:
32 import hyper_parallel.auto_parallel.sapp_nd.nd.parallelize as Par
33 import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim
34 import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as Hard
36logger = logging.getLogger(__name__)
38def _get_dim_module():
39 """Lazy-import the sapp_nd dimensions module."""
40 import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as dim_mod # pylint: disable=C0415
41 return dim_mod
44def _get_machine_mod():
45 """Lazy-import the sapp_nd hardware module."""
46 import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as hw_mod # pylint: disable=C0415
47 return hw_mod
50def _search_dim_map():
51 """Return the mapping of NormalizedConfig keys to sapp_nd Dimension objects.
53 Lazily loaded to avoid importing sapp_nd at module-import time.
54 """
55 dim_mod = _get_dim_module()
56 return {
57 "data_parallel_replicate_degree": dim_mod.DP,
58 "tensor_parallel_degree": dim_mod.TP,
59 "pipeline_parallel_degree": dim_mod.PP,
60 "context_parallel_degree": dim_mod.CP,
61 "expert_parallel_degree": dim_mod.EP,
62 "micro_batch_num": dim_mod.MBN,
63 }
65# Accelerator field names for fixed dimensions (written into the temp YAML).
66_ACCEL_FIELD_MAP: Dict[str, str] = {
67 "data_parallel_shard_degree": "dp_shard",
68 "data_parallel_replicate_degree": "dp_replicate",
69 "tensor_parallel_degree": "tp_degree",
70 "pipeline_parallel_degree": "pipeline_parallel_degree",
71 "context_parallel_degree": "context_parallel_degree",
72 "expert_parallel_degree": "expert_parallel_degree",
73 "micro_batch_num": "micro_batch_num",
74}
77def _validate_before_search(config: NormalizedConfig) -> None:
78 """Check required model fields are populated (>0) before search.
80 Raises:
81 ValueError: If any required field is missing or zero.
82 """
83 model = config.model_spec
84 required = {
85 "model_spec.n_layers": model.get("n_layers", 0),
86 "model_spec.dim": model.get("dim", 0),
87 "model_spec.n_heads": model.get("n_heads", 0),
88 "model_spec.vocab_size": model.get("vocab_size", 0),
89 "cluster_spec": config.cluster_spec,
90 }
91 missing = []
92 for name, value in required.items():
93 if name == "cluster_spec":
94 if not isinstance(value, dict) or not value:
95 missing.append(name)
96 elif value <= 0:
97 missing.append(name)
98 if missing:
99 raise ValueError(
100 "Required fields missing or zero before ND search: "
101 f"{', '.join(missing)}"
102 )
105def _build_hp_yaml_dict(config: NormalizedConfig) -> dict:
106 """Build a HyperParallel ``train.yaml`` dict from *config*.
108 Fixed dimensions (``constraint.fixed_*_degree``) are written directly
109 into ``train.accelerator``. Dimensions with search-space candidates
110 use the first candidate as a placeholder — the actual search is driven
111 by the ``dimensions`` parameter passed to :class:`Parallelize`.
112 """
113 model = config.model_spec
114 constraint = config.constraint
115 space = config.search_space
117 accel: Dict[str, Any] = {}
119 # Fixed dimensions → write actual value.
120 fixed_map = {
121 "fixed_dp_degree": ("dp_replicate", "data_parallel_replicate_degree", [1]),
122 "fixed_fsdp_degree": ("dp_shard", "data_parallel_shard_degree", [1]),
123 "fixed_tp_degree": ("tp_degree", "tensor_parallel_degree", [1]),
124 "fixed_pp_degree": ("pipeline_parallel_degree", "pipeline_parallel_degree", [1]),
125 "fixed_cp_degree": ("context_parallel_degree", "context_parallel_degree", [1]),
126 "fixed_ep_degree": ("expert_parallel_degree", "expert_parallel_degree", [1]),
127 }
128 for constraint_key, (accel_key, space_key, default) in fixed_map.items():
129 fixed_val = constraint.get(constraint_key)
130 if fixed_val is not None and fixed_val > 0:
131 accel[accel_key] = fixed_val
132 else:
133 candidates = space.get(space_key, default)
134 accel[accel_key] = candidates[0]
136 # Enable parallel optimizer by default.
137 accel.setdefault("enable_parallel_optimizer", True)
139 recompute = config.estimator.get("recompute_strategy", "none")
141 hp_yaml: dict = {
142 "model": {
143 "name": model.get("name", "custom"),
144 "config_overrides": {
145 "hidden_size": model.get("dim", 4096),
146 "num_hidden_layers": model.get("n_layers", 32),
147 "num_attention_heads": model.get("n_heads", 32),
148 "vocab_size": model.get("vocab_size", 128256),
149 },
150 },
151 "train": {
152 "global_batch_size": constraint.get("global_batch_size", 0) or 1,
153 "micro_batch_size": model.get("local_batch_size", 1),
154 "micro_batch_num": accel.pop("micro_batch_num", 1),
155 "accelerator": accel,
156 "gradient_checkpointing": {
157 "activation_checkpoint": recompute,
158 },
159 "mixed_precision": {
160 "enabled": True,
161 "param_dtype": model.get("compute_dtype", "bfloat16"),
162 },
163 },
164 "data": {
165 "max_seq_len": model.get("seq_len", 4096),
166 },
167 }
169 # Optional model fields.
170 overrides = hp_yaml["model"]["config_overrides"]
171 if model.get("inter_dim"):
172 overrides["intermediate_size"] = model["inter_dim"]
173 if model.get("n_kv_heads"):
174 overrides["num_key_value_heads"] = model["n_kv_heads"]
176 return hp_yaml
179def _write_temp_hp_yaml(config: NormalizedConfig) -> str:
180 """Write a temp ``train.yaml`` and return its absolute path."""
181 data = _build_hp_yaml_dict(config)
182 fd, path = tempfile.mkstemp(suffix=".yaml", prefix="hp_search_")
183 os.close(fd)
184 with open(path, "w", encoding="utf-8") as fh:
185 yaml.dump(data, fh, default_flow_style=False, sort_keys=False)
186 logger.debug("Temp HP YAML written to %s", path)
187 return path
190def _build_machine(config: NormalizedConfig) -> Any:
191 """Build a ``Hard.Machine`` from cluster_spec."""
192 hw_mod = _get_machine_mod()
193 cluster = config.cluster_spec
194 nodes = max(1, cluster.get("num_nodes", 1))
195 cards_per_node = max(1, cluster.get("cards_per_node", 8))
196 total_devices = nodes * cards_per_node
197 device_type = cluster.get("device_type", "A2")
198 # Map generic names to sapp_nd device codes.
199 device_code_map = {"ascend": "A2", "ascend910": "A2", "ascend910b": "A3"}
200 device_type = device_code_map.get(str(device_type).lower(), device_type)
201 return hw_mod.Machine(total_devices, device_type)
204def _resolve_search_dimensions(config: NormalizedConfig) -> List[Any]:
205 """Return a list of ``Dim`` objects whose candidates contain >1 value.
207 List-valued entries in ``config.search_space`` are treated as
208 **output** (search) dimensions. Entries absent from
209 ``search_space`` (``"auto"`` in YAML) are also included — they
210 will be determined by ND's ``bound_space()``.
211 """
212 dims: List[Any] = []
213 space = config.search_space
214 for space_key, dim_obj in _search_dim_map().items():
215 candidates = space.get(space_key)
216 if candidates is not None and len(candidates) > 1:
217 dims.append(dim_obj)
218 elif space_key not in space:
219 dims.append(dim_obj)
220 return dims
223def _post_filter(
224 scored_space: list,
225 config: NormalizedConfig,
226) -> list:
227 """Keep only entries whose dimension values are in the user's candidate lists."""
228 space = config.search_space
229 candidate_map: Dict[Any, List[int]] = {}
230 for space_key, dim_obj in _search_dim_map().items():
231 candidates = space.get(space_key)
232 if candidates is not None and len(candidates) > 1:
233 candidate_map[dim_obj] = candidates
235 filtered = []
236 for entry in scored_space:
237 dims_val = entry[0].dims_val # type: ignore[index]
238 keep = True
239 for dim_obj, allowed in candidate_map.items():
240 actual = dims_val.get(dim_obj)
241 if actual is not None and actual not in allowed:
242 keep = False
243 break
244 if keep:
245 filtered.append(entry)
247 if not filtered and scored_space:
248 logger.warning(
249 "Post-filter removed ALL %d candidates. "
250 "Returning unfiltered best entry.",
251 len(scored_space),
252 )
253 return scored_space[:1]
254 return filtered
257def _format_result(best_entry: tuple) -> Dict[str, Any]:
258 """Convert the best ND result entry into a flat result dict."""
259 dim_mod = _get_dim_module()
260 dims_val = best_entry[0].dims_val # type: ignore[index]
261 dim_to_key = {
262 dim_mod.DP: "dp",
263 dim_mod.TP: "tp",
264 dim_mod.PP: "pp",
265 dim_mod.CP: "cp",
266 dim_mod.EP: "ep",
267 dim_mod.MBN: "micro_batch_num",
268 }
269 result: Dict[str, Any] = {
270 "memory_estimate_mb": float(best_entry[1]),
271 "score": float(best_entry[2]),
272 }
273 for dim_obj, key in dim_to_key.items():
274 if dim_obj in dims_val:
275 result[key] = int(dims_val[dim_obj])
276 return result
279def search_strategies(config: NormalizedConfig) -> Dict[str, Any]:
280 """Run the ND strategy search and return the optimal strategy.
282 This is the main entry point for end-to-end strategy search:
284 1. Validates required model fields.
285 2. Converts the ``NormalizedConfig`` to a temporary HyperParallel
286 ``train.yaml`` and writes it to disk.
287 3. Launches the ND search engine (:class:`Parallelize`).
288 4. Post-filters results against the user's candidate lists.
289 5. Returns the best strategy as a flat dictionary.
291 Args:
292 config: A fully populated ``NormalizedConfig`` from
293 :func:`read_search_config` or :func:`read_hp_yaml_config`.
295 Returns:
296 A dict with keys ``dp``, ``tp``, ``pp``, ``cp``, ``ep``,
297 ``micro_batch_num``, ``memory_estimate_mb``, and ``score``.
299 Raises:
300 ValueError: If required fields are missing or no strategy is found.
301 ImportError: If PyYAML is not installed.
302 """
303 _validate_before_search(config)
305 yaml_path = _write_temp_hp_yaml(config)
306 machine = _build_machine(config)
307 dims = _resolve_search_dimensions(config)
309 import hyper_parallel.auto_parallel.sapp_nd.nd.parallelize as _Par # pylint: disable=C0415
310 try:
311 nd_runner = _Par.Parallelize(
312 "hyper_v2",
313 yaml_path,
314 machine,
315 global_batch_size=config.constraint.get("global_batch_size", 0),
316 dimensions=dims,
317 )
318 scored_space = nd_runner.run_generation_to_ordering(
319 yaml_folder=None,
320 threads_num=None,
321 top_num=None,
322 )
323 finally:
324 try:
325 os.remove(yaml_path)
326 except OSError:
327 pass
329 if not scored_space:
330 raise ValueError("ND search returned no valid strategies.")
332 filtered = _post_filter(scored_space, config)
333 best = filtered[0]
334 result = _format_result(best)
336 logger.info(
337 "Optimal strategy found: dp=%(dp)s tp=%(tp)s pp=%(pp)s "
338 "cp=%(cp)s ep=%(ep)s mb_num=%(micro_batch_num)s "
339 "mem=%(memory_estimate_mb).0f MB score=%(score).2e",
340 result,
341 )
342 return result