Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / config_adapter / _search_runner.py: 95%
165 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"""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, Optional, Set, Tuple, TYPE_CHECKING
27import yaml # type: ignore[import-untyped]
29from hyper_parallel.auto_parallel.config_adapter._normalized_config import NormalizedConfig
32CONFIG_OVERRIDE_FIELDS = [
33 "hidden_size", "num_hidden_layers", "num_attention_heads", "vocab_size",
34 "intermediate_size", "num_key_value_heads", "max_position_embeddings",
35 "num_experts", "num_experts_per_tok", "num_shared_experts",
36 "moe_intermediate_size", "first_k_dense_replace", "mtp_depth",
37 "multiple_of", "ffn_dim_multiplier", "kv_lora_rank", "q_lora_rank",
38 "qk_rope_head_dim", "v_head_dim", "capacity_factor", "offset",
39 "param_init_type", "compute_dtype", "softmax_compute_type",
40]
42if TYPE_CHECKING:
43 import hyper_parallel.auto_parallel.sapp_nd.nd.parallelize as Par
44 import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim
45 import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as Hard
47logger = logging.getLogger(__name__)
49def _get_dim_module():
50 """Lazy-import the sapp_nd dimensions module."""
51 import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as dim_mod # pylint: disable=C0415
52 return dim_mod
55def _get_machine_mod():
56 """Lazy-import the sapp_nd hardware module."""
57 import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as hw_mod # pylint: disable=C0415
58 return hw_mod
61def _search_dim_map():
62 """Return the mapping of NormalizedConfig keys to sapp_nd Dimension objects.
64 Lazily loaded to avoid importing sapp_nd at module-import time.
65 """
66 dim_mod = _get_dim_module()
67 return {
68 "data_parallel_replicate_degree": dim_mod.DP,
69 "tensor_parallel_degree": dim_mod.TP,
70 "pipeline_parallel_degree": dim_mod.PP,
71 "context_parallel_degree": dim_mod.CP,
72 "expert_parallel_degree": dim_mod.EP,
73 "micro_batch_num": dim_mod.MBN,
74 }
76# Accelerator field names for fixed dimensions (written into the temp YAML).
77_ACCEL_FIELD_MAP: Dict[str, str] = {
78 "data_parallel_shard_degree": "dp_shard",
79 "data_parallel_replicate_degree": "dp_replicate",
80 "tensor_parallel_degree": "tp_degree",
81 "pipeline_parallel_degree": "pipeline_parallel_degree",
82 "context_parallel_degree": "context_parallel_degree",
83 "expert_parallel_degree": "expert_parallel_degree",
84 "micro_batch_num": "micro_batch_num",
85}
88def _validate_before_search(config: NormalizedConfig) -> None:
89 """Check required model fields are populated (>0) before search.
91 Raises:
92 ValueError: If any required field is missing or zero.
93 """
94 model = config.model_spec
95 required = {
96 "model_spec.num_hidden_layers": model.get("num_hidden_layers", 0),
97 "model_spec.hidden_size": model.get("hidden_size", 0),
98 "model_spec.num_attention_heads": model.get("num_attention_heads", 0),
99 "model_spec.vocab_size": model.get("vocab_size", 0),
100 "cluster_spec": config.cluster_spec,
101 }
102 missing = []
103 for name, value in required.items():
104 if name == "cluster_spec":
105 if not isinstance(value, dict) or not value:
106 missing.append(name)
107 elif value <= 0:
108 missing.append(name)
109 if missing:
110 raise ValueError(
111 "Required fields missing or zero before ND search: "
112 f"{', '.join(missing)}"
113 )
116def _build_model_dict(model: Dict[str, Any]) -> Dict[str, Any]:
117 """Build the ``model`` section of the HP YAML from *model* spec.
119 All ``config_overrides`` field names in *model* already match the
120 HP YAML convention, so they are passed through directly without
121 any name mapping.
123 Args:
124 model: The ``model_spec`` dict from :class:`NormalizedConfig`.
126 Returns:
127 A dict suitable for the ``model`` key of a HP ``train.yaml``.
128 """
129 model_dict: Dict[str, Any] = {
130 "name": model.get("name", "custom"),
131 "config_overrides": {},
132 }
133 overrides = model_dict["config_overrides"]
134 for key in CONFIG_OVERRIDE_FIELDS:
135 val = model.get(key)
136 if val is not None:
137 overrides[key] = val
139 return model_dict
142def _build_hp_yaml_dict(config: NormalizedConfig) -> dict:
143 """Build a HyperParallel ``train.yaml`` dict from *config*.
145 Fixed dimensions (``constraint.fixed_*_degree``) are written directly
146 into ``train.accelerator``. Dimensions with search-space candidates
147 use the first candidate as a placeholder -- the actual search is driven
148 by the ``dimensions`` parameter passed to :class:`Parallelize`.
149 """
150 model = config.model_spec
151 constraint = config.constraint
152 space = config.search_space
154 accel: Dict[str, Any] = {}
156 # Fixed dimensions -- write actual value.
157 fixed_map = {
158 "fixed_dp_degree": ("dp_replicate", "data_parallel_replicate_degree", [1]),
159 "fixed_fsdp_degree": ("dp_shard", "data_parallel_shard_degree", [1]),
160 "fixed_tp_degree": ("tp_degree", "tensor_parallel_degree", [1]),
161 "fixed_pp_degree": ("pipeline_parallel_degree", "pipeline_parallel_degree", [1]),
162 "fixed_cp_degree": ("context_parallel_degree", "context_parallel_degree", [1]),
163 "fixed_ep_degree": ("expert_parallel_degree", "expert_parallel_degree", [1]),
164 "fixed_etp_degree": ("expert_tensor_parallel_degree", "expert_tensor_parallel_degree", [0]),
165 }
166 for constraint_key, (accel_key, space_key, default) in fixed_map.items():
167 fixed_val = constraint.get(constraint_key)
168 if fixed_val is not None and fixed_val > 0:
169 accel[accel_key] = fixed_val
170 else:
171 candidates = space.get(space_key, default)
172 accel[accel_key] = candidates[0]
174 # Enable parallel optimizer by default.
175 accel.setdefault("enable_parallel_optimizer", True)
177 # CP algorithm: propagate to yaml so CostModelParserHyperV2 can read it.
178 cp_algo = config.estimator.get("cp_algo")
179 if cp_algo:
180 accel["context_parallel_algo"] = cp_algo
182 # Optional accelerator fields that affect memory estimation.
183 owss = model.get("optimizer_weight_shard_size")
184 if owss and owss > 0:
185 accel["optimizer_weight_shard_size"] = owss
187 use_sp = model.get("use_seq_parallel", True)
188 accel.setdefault("use_seq_parallel", bool(use_sp))
190 recompute = config.estimator.get("recompute_strategy", "none")
192 cluster = config.cluster_spec
193 device_mem_gb = cluster.get("device_memory_gb", 0)
194 context: Dict[str, Any] = {}
195 if device_mem_gb > 0:
196 context["max_device_memory"] = f"{device_mem_gb}GB"
198 gc_dict: Dict[str, Any] = {"activation_checkpoint": recompute}
199 recompute_slice = model.get("recompute_slice_activation")
200 if recompute_slice is not None:
201 gc_dict["recompute_slice_activation"] = bool(recompute_slice)
203 model_dict = _build_model_dict(model)
205 hp_yaml: dict = {
206 "model": model_dict,
207 "train": {
208 "global_batch_size": constraint.get("global_batch_size", 0),
209 "micro_batch_size": model.get("local_batch_size", 1),
210 "micro_batch_num": accel.pop("micro_batch_num", 1),
211 "accelerator": accel,
212 "gradient_checkpointing": gc_dict,
213 "mixed_precision": {
214 "enabled": True,
215 "param_dtype": model.get("compute_dtype", "bfloat16"),
216 },
217 },
218 "data": {
219 "max_seq_len": model.get("max_position_embeddings", 4096),
220 },
221 }
223 if context:
224 hp_yaml["context"] = context
226 return hp_yaml
229def _write_temp_hp_yaml(config: NormalizedConfig) -> str:
230 """Write a temp ``train.yaml`` and return its absolute path."""
231 data = _build_hp_yaml_dict(config)
232 fd, path = tempfile.mkstemp(suffix=".yaml", prefix="hp_search_")
233 os.close(fd)
234 with open(path, "w", encoding="utf-8") as fh:
235 yaml.dump(data, fh, default_flow_style=False, sort_keys=False)
236 logger.debug("Temp HP YAML written to %s", path)
237 return path
240def _build_machine(config: NormalizedConfig) -> Any:
241 """Build a ``Hard.Machine`` from cluster_spec."""
242 hw_mod = _get_machine_mod()
243 cluster = config.cluster_spec
244 nodes = max(1, cluster.get("num_nodes", 1))
245 cards_per_node = max(1, cluster.get("cards_per_node", 8))
246 total_devices = nodes * cards_per_node
247 device_type = cluster.get("device_type", "A2")
248 # Map generic names to sapp_nd device codes.
249 device_code_map = {"ascend": "A2", "ascend910": "A2", "ascend910b": "A3"}
250 device_type = device_code_map.get(str(device_type).lower(), device_type)
251 return hw_mod.Machine(total_devices, device_type)
254def _resolve_search_dimensions(config: NormalizedConfig) -> Tuple[List[Any], Set[Any]]:
255 """Return search dimensions and the set of dimensions with user candidates.
257 List-valued entries in ``config.search_space`` are treated as
258 **output** (search) dimensions. Entries absent from
259 ``config.search_space`` (``"auto"`` in YAML) are also included -- they
260 will be determined by ND's ``bound_space()``.
262 Returns:
263 A tuple ``(dims, candidate_dims)`` where *dims* is the list of
264 ``Dim`` objects to pass to ND and *candidate_dims* is the set of
265 Dim objects for which the user supplied an explicit candidate list
266 (used by :func:`_post_filter`).
267 """
268 dims: List[Any] = []
269 candidate_dims: Set[Any] = set()
270 space = config.search_space
271 for space_key, dim_obj in _search_dim_map().items():
272 candidates = space.get(space_key)
273 if candidates is not None and len(candidates) > 1:
274 dims.append(dim_obj)
275 candidate_dims.add(dim_obj)
276 elif space_key not in space:
277 dims.append(dim_obj)
278 return dims, candidate_dims
281def _post_filter(
282 scored_space: list,
283 config: NormalizedConfig,
284 candidate_dims: Optional[Set[Any]] = None,
285) -> list:
286 """Keep only entries whose dimension values are in the user's candidate lists.
288 Args:
289 scored_space: The scored strategy list from ND engine.
290 config: The normalized config containing ``search_space``.
291 candidate_dims: The set of Dim objects that have user-supplied
292 candidate lists with more than one value. If *None*, the
293 set is derived from *config* (backward-compatible).
295 Returns:
296 A filtered list. May be empty if no entry satisfies all
297 candidate constraints -- the caller decides how to handle this.
298 """
299 space = config.search_space
300 if candidate_dims is None:
301 candidate_dims = set()
302 for space_key, dim_obj in _search_dim_map().items():
303 candidates = space.get(space_key)
304 if candidates is not None and len(candidates) > 1:
305 candidate_dims.add(dim_obj)
307 candidate_map: Dict[Any, List[int]] = {}
308 for space_key, dim_obj in _search_dim_map().items():
309 if dim_obj not in candidate_dims:
310 continue
311 candidates = space.get(space_key)
312 if candidates is not None:
313 candidate_map[dim_obj] = candidates
315 filtered = []
316 for entry in scored_space:
317 dims_val = entry[0].dims_val # type: ignore[index]
318 keep = True
319 for dim_obj, allowed in candidate_map.items():
320 actual = dims_val.get(dim_obj)
321 if actual is not None and actual not in allowed:
322 keep = False
323 break
324 if keep:
325 filtered.append(entry)
327 if not filtered and scored_space:
328 logger.warning(
329 "Post-filter removed ALL %d candidates; "
330 "no strategy matches the user's candidate constraints.",
331 len(scored_space),
332 )
333 return scored_space[:1]
334 return filtered
337def _format_result(best_entry: tuple) -> Dict[str, Any]:
338 """Convert the best ND result entry into a flat result dict."""
339 dim_mod = _get_dim_module()
340 dims_val = best_entry[0].dims_val # type: ignore[index]
341 dim_to_key = {
342 dim_mod.DP: "dp",
343 dim_mod.TP: "tp",
344 dim_mod.PP: "pp",
345 dim_mod.CP: "cp",
346 dim_mod.EP: "ep",
347 dim_mod.MBN: "micro_batch_num",
348 }
349 result: Dict[str, Any] = {
350 "memory_estimate_mb": float(best_entry[1]),
351 "score": float(best_entry[2]),
352 }
353 for dim_obj, key in dim_to_key.items():
354 if dim_obj in dims_val:
355 result[key] = int(dims_val[dim_obj])
356 result.setdefault("cp", 1)
357 result.setdefault("ep", 1)
358 return result
361def search_strategies(config: NormalizedConfig) -> Dict[str, Any]:
362 """Run the ND strategy search and return the optimal strategy.
364 This is the main entry point for end-to-end strategy search:
366 1. Validates required model fields.
367 2. Converts the ``NormalizedConfig`` to a temporary HyperParallel
368 ``train.yaml`` and writes it to disk.
369 3. Launches the ND search engine (:class:`Parallelize`).
370 4. Post-filters results against the user's candidate lists.
371 5. Returns the best strategy as a flat dictionary.
373 Args:
374 config: A fully populated ``NormalizedConfig`` from
375 :func:`read_search_config` or :func:`read_hp_yaml_config`.
377 Returns:
378 A dict with keys ``dp``, ``tp``, ``pp``, ``cp``, ``ep``,
379 ``micro_batch_num``, ``memory_estimate_mb``, and ``score``.
381 Raises:
382 ValueError: If required fields are missing or no strategy is found.
383 ImportError: If PyYAML is not installed.
384 """
385 _validate_before_search(config)
387 yaml_path = _write_temp_hp_yaml(config)
388 machine = _build_machine(config)
389 dims, candidate_dims = _resolve_search_dimensions(config)
391 import hyper_parallel.auto_parallel.sapp_nd.nd.parallelize as _Par # pylint: disable=C0415
392 try:
393 nd_runner = _Par.Parallelize(
394 "hyper_v2",
395 yaml_path,
396 machine,
397 global_batch_size=config.constraint.get("global_batch_size", 0),
398 dimensions=dims,
399 )
400 scored_space = nd_runner.run_generation_to_ordering(
401 yaml_folder=None,
402 threads_num=None,
403 top_num=None,
404 )
405 finally:
406 try:
407 os.remove(yaml_path)
408 except OSError:
409 pass
411 if not scored_space:
412 raise ValueError("ND search returned no valid strategies.")
414 filtered = _post_filter(scored_space, config, candidate_dims)
415 best = filtered[0]
416 result = _format_result(best)
418 logger.info(
419 "Optimal strategy found: dp=%(dp)s tp=%(tp)s pp=%(pp)s "
420 "cp=%(cp)s ep=%(ep)s mb_num=%(micro_batch_num)s "
421 "mem=%(memory_estimate_mb).0f MB score=%(score).2e",
422 result,
423 )
424 return result