Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / config_adapter / _search_runner.py: 96%
246 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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.
17Hands a :class:`NormalizedConfig` to the ND search as the mapping
18``Parallelize`` parses directly, post-filters by user candidate lists and
19memory budget, and returns the optimal strategy.
20"""
22import copy
23import logging
24from typing import Any, Dict, FrozenSet, List, Mapping, Optional, Set, Tuple, TYPE_CHECKING
26from hyper_parallel.auto_parallel._hf_model_spec import EXEC_OVERRIDE_KEYS
27from hyper_parallel.auto_parallel._model_spec import ModelSpec, model_fields
28from 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__)
38# The activation_checkpoint modes HyperParallel's trainer runs every layer
39# with, that recompute "auto" chooses among. Its selective mode recomputes
40# every other matmul, which the cost model's formulas do not price; a census
41# of the layers measures it (IR phase 5), so a run that states one offers it
42# too, where the census priced it.
43TRAINER_RECOMPUTE_MODES = ("off", "full")
44TRAINER_CENSUS_MODES = ("off", "selective", "full")
45# The search config's recompute values that choose among those modes: "auto"
46# one for every layer, "per_layer" one for each layer, which the trainer runs
47# as activation_checkpoint.layers.
48CHOSEN_RECOMPUTE = ("auto", "per_layer")
51def _get_dim_module():
52 """Lazy-import the sapp_nd dimensions module."""
53 import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as dim_mod # pylint: disable=C0415
54 return dim_mod
57def _get_machine_mod():
58 """Lazy-import the sapp_nd hardware module."""
59 import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as hw_mod # pylint: disable=C0415
60 return hw_mod
63def _search_dim_map():
64 """Return the mapping of NormalizedConfig keys to sapp_nd Dimension objects.
66 Lazily loaded to avoid importing sapp_nd at module-import time.
67 """
68 dim_mod = _get_dim_module()
69 return {
70 "data_parallel_replicate_degree": dim_mod.DP,
71 "tensor_parallel_degree": dim_mod.TP,
72 "pipeline_parallel_degree": dim_mod.PP,
73 "context_parallel_degree": dim_mod.CP,
74 "expert_parallel_degree": dim_mod.EP,
75 "micro_batch_num": dim_mod.MBN,
76 # OP carries the FSDP/optimizer shard degree (ccfg.os_max_shard),
77 # so mapping it here is what lets ND search that dimension.
78 "data_parallel_shard_degree": dim_mod.OP,
79 }
82def _validate_before_search(config: NormalizedConfig) -> None:
83 """Check required model fields are populated (>0) before search.
85 Raises:
86 ValueError: If any required field is missing or zero.
87 """
88 model = config.model_spec
89 required = {
90 "model_spec.num_hidden_layers": model.get("num_hidden_layers", 0),
91 "model_spec.hidden_size": model.get("hidden_size", 0),
92 "model_spec.num_attention_heads": model.get("num_attention_heads", 0),
93 "model_spec.vocab_size": model.get("vocab_size", 0),
94 "cluster_spec": config.cluster_spec,
95 }
96 missing = []
97 for name, value in required.items():
98 if name == "cluster_spec":
99 if not isinstance(value, dict) or not value:
100 missing.append(name)
101 elif value <= 0:
102 missing.append(name)
103 if missing:
104 raise ValueError(
105 "Required fields missing or zero before ND search: "
106 f"{', '.join(missing)}"
107 )
110def _build_model_dict(model: Dict[str, Any]) -> Dict[str, Any]:
111 """Build the ``model`` section of the HP YAML from *model* spec.
113 The model's fields go through the model spec, whose names are the HP
114 YAML's, and the run keys the parser reads from ``config_overrides``
115 pass through as they are. The adapter's own keys, such as the
116 micro-batch size, reach the HP YAML through their own sections.
118 Args:
119 model: The ``model_spec`` dict from :class:`NormalizedConfig`.
121 Returns:
122 A dict suitable for the ``model`` key of a HP ``train.yaml``.
123 """
124 spec = ModelSpec.from_dict(model_fields(model))
125 overrides = spec.to_dict()
126 overrides.pop("name", None)
127 overrides.update({key: model[key] for key in EXEC_OVERRIDE_KEYS if key in model})
128 return {"name": spec.name, "config_overrides": overrides}
131# The strategy, under every name the cost model reads it by, in the sections
132# a train.yaml states its run in: the search sets it, from its candidates or
133# its own config, so the train.yaml's values give way. The rest of each
134# section is the run's.
135_STRATEGY_KEYS: Dict[str, FrozenSet[str]] = {
136 "accelerator": frozenset({
137 "dp_replicate", "dp_shard", "tp_size", "tp_degree", "pp_size", "pipeline_parallel_degree",
138 "cp_size", "context_parallel_degree", "ep_size", "expert_parallel_degree",
139 "expert_tensor_parallel_degree", "micro_batch_num", "optimizer_weight_shard_size",
140 "pipeline_scheduler", "pp_interleave_num",
141 }),
142 "fsdp_config": frozenset({"dp_shard_size"}),
143 "training": frozenset({"global_batch_size", "micro_batch_size", "micro_batch_num"}),
144}
147def _stated_run(config: NormalizedConfig) -> Dict[str, Any]:
148 """Return the run the train.yaml states, less the strategy the search sets.
150 Args:
151 config: The normalized config whose ``run`` the train.yaml filled.
153 Returns:
154 A copy of ``config.run``, its sections without the strategy's keys.
155 """
156 run = copy.deepcopy(config.run)
157 for section, strategy in _STRATEGY_KEYS.items():
158 stated = run.get(section)
159 if isinstance(stated, dict):
160 run[section] = {key: value for key, value in stated.items() if key not in strategy}
161 return run
164def _memory_budget_gb(config: NormalizedConfig) -> float:
165 """Return the per-device memory budget in GB, or ``0.0`` if unconstrained.
167 ``constraint.memory_limit_gb`` is the budget the user asked for and
168 ``cluster_spec.device_memory_gb`` the hardware ceiling, so the search
169 must respect whichever of the two is tighter.
171 Args:
172 config: The normalized config carrying constraint and cluster spec.
174 Returns:
175 The tighter positive budget in GB, or ``0.0`` when neither is set.
176 """
177 limit = float(config.constraint.get("memory_limit_gb", 0.0) or 0.0)
178 device = float(config.cluster_spec.get("device_memory_gb", 0.0) or 0.0)
179 budgets = [value for value in (limit, device) if value > 0]
180 return min(budgets) if budgets else 0.0
183def _pinned_or_first(config: NormalizedConfig, constraint_key: str, space_key: str,
184 default: List[Any]) -> Any:
185 """Return the degree *constraint_key* pins, else the first candidate.
187 Args:
188 config: The normalized config carrying constraint and search space.
189 constraint_key: The ``constraint`` entry that pins the degree.
190 space_key: The ``search_space`` entry listing its candidates.
191 default: The candidates when the search space lists none.
193 Returns:
194 The pinned degree when it is positive, else the first candidate,
195 a placeholder the search replaces.
196 """
197 pinned = config.constraint.get(constraint_key)
198 if pinned is not None and pinned > 0:
199 return pinned
200 return config.search_space.get(space_key, default)[0]
203def _build_strategy_dicts(
204 config: NormalizedConfig, run: Dict[str, Any],
205) -> Tuple[Dict[str, Any], Dict[str, Any]]:
206 """Build the ``accelerator`` and ``fsdp_config`` sections of the HP YAML.
208 The strategy goes over the sections the train.yaml states, which this
209 takes out of *run*.
210 """
211 accel: Dict[str, Any] = dict(run.pop("accelerator", None) or {})
212 fsdp: Dict[str, Any] = dict(run.pop("fsdp_config", None) or {})
214 # Fixed dimensions -- write actual value.
215 fixed_map = {
216 "fixed_dp_degree": ("dp_replicate", "data_parallel_replicate_degree", [1]),
217 "fixed_tp_degree": ("tp_size", "tensor_parallel_degree", [1]),
218 "fixed_pp_degree": ("pp_size", "pipeline_parallel_degree", [1]),
219 "fixed_cp_degree": ("cp_size", "context_parallel_degree", [1]),
220 "fixed_ep_degree": ("ep_size", "expert_parallel_degree", [1]),
221 "fixed_etp_degree": ("expert_tensor_parallel_degree", "expert_tensor_parallel_degree", [0]),
222 }
223 for constraint_key, (accel_key, space_key, default) in fixed_map.items():
224 accel[accel_key] = _pinned_or_first(config, constraint_key, space_key, default)
226 # micro_batch_num has no accelerator entry of its own, so a fixed value
227 # would be dropped and the global-batch-size check would then reject the
228 # only config the user asked for.
229 accel["micro_batch_num"] = int(
230 _pinned_or_first(config, "fixed_micro_batch_num", "micro_batch_num", [1])
231 )
232 fsdp["dp_shard_size"] = int(
233 _pinned_or_first(config, "fixed_fsdp_degree", "data_parallel_shard_degree", [1])
234 )
236 # CP algorithm: propagate to yaml so CostModelParserHyperV2 can read it.
237 cp_algo = config.estimator.get("cp_algo")
238 if cp_algo:
239 accel["context_parallel_algo"] = cp_algo
241 # Optional accelerator fields that affect memory estimation.
242 owss = config.model_spec.get("optimizer_weight_shard_size")
243 if owss and owss > 0:
244 accel["optimizer_weight_shard_size"] = owss
246 use_sp = config.model_spec.get("use_seq_parallel", True)
247 accel.setdefault("sequence_parallel", bool(use_sp))
248 return accel, fsdp
251def _build_hp_yaml_dict(config: NormalizedConfig) -> dict:
252 """Build an AutoModels-shaped cost-model YAML dict from *config*.
254 Fixed dimensions (``constraint.fixed_*_degree``) are written directly
255 into the strategy sections. Dimensions with search-space candidates
256 use the first candidate as a placeholder -- the actual search is driven
257 by the ``dimensions`` parameter passed to :class:`Parallelize`. The
258 strategy goes over the run the train.yaml states, which stays as stated.
259 """
260 model = config.model_spec
261 constraint = config.constraint
262 run = _stated_run(config)
263 accel, fsdp = _build_strategy_dicts(config, run)
265 recompute = config.estimator.get("recompute_strategy", "none")
267 cluster = config.cluster_spec
268 # The train.yaml's pricing options, beneath the search's own device count
269 # and memory budget.
270 context: Dict[str, Any] = dict(run.pop("context", None) or {})
271 # ND prunes the space against ccfg.device_capacity, which the parser reads
272 # from this field, so the user's memory_limit_gb has to reach it here.
273 budget_gb = _memory_budget_gb(config)
274 if budget_gb > 0:
275 context["max_device_memory"] = f"{budget_gb}GB"
276 device_num = cluster.get("num_nodes", 0) * cluster.get("cards_per_node", 0)
277 if device_num > 0:
278 context["device_num"] = int(device_num)
279 visual_seq_len = model.get("visual_seq_len")
280 if visual_seq_len:
281 context["visual_seq_len"] = int(visual_seq_len)
283 # Choosing the recompute, the search keeps what fits fully recomputed:
284 # describe it so.
285 stated_mode = "full" if recompute in CHOSEN_RECOMPUTE else {"none": "off"}.get(recompute, recompute)
286 gc_dict: Dict[str, Any] = {"mode": stated_mode}
287 recompute_slice = model.get("recompute_slice_activation")
288 if recompute_slice is not None:
289 gc_dict["recompute_slice_activation"] = bool(recompute_slice)
291 model_dict = _build_model_dict(model)
293 hp_yaml: dict = {
294 **run,
295 "model": {**run.get("model", {}), **model_dict},
296 "training": {
297 **run.get("training", {}),
298 "global_batch_size": constraint.get("global_batch_size", 0),
299 "micro_batch_size": model.get("local_batch_size", 1),
300 "micro_batch_num": accel.pop("micro_batch_num", 1),
301 },
302 "accelerator": accel,
303 "fsdp_config": fsdp,
304 "activation_checkpoint": gc_dict,
305 "dataset": {
306 "data_transform": {
307 "max_seq_len": model.get("max_position_embeddings", 4096),
308 },
309 },
310 }
312 if context:
313 hp_yaml["context"] = context
315 return hp_yaml
319# The names a search config may give its devices, and the sapp_nd device
320# codes they mean: the Ascend 910B is the Atlas A2 series' chip, and the
321# 910C, CANN's ascend910_93, the A3's. A generic "ascend" stays A2.
322_DEVICE_CODES: Dict[str, str] = {
323 "a2": "A2", "a3": "A3", "v100": "V100",
324 "ascend": "A2", "ascend910": "A2", "ascend910b": "A2",
325 "ascend910c": "A3", "ascend910_93": "A3",
326}
329def _build_machine(config: NormalizedConfig) -> Any:
330 """Build a ``Hard.Machine`` from cluster_spec."""
331 hw_mod = _get_machine_mod()
332 cluster = config.cluster_spec
333 nodes = max(1, cluster.get("num_nodes", 1))
334 cards_per_node = max(1, cluster.get("cards_per_node", 8))
335 total_devices = nodes * cards_per_node
336 device_type = str(cluster.get("device_type", "A2"))
337 return hw_mod.Machine(total_devices, _DEVICE_CODES.get(device_type.lower(), device_type))
340def _resolve_search_dimensions(config: NormalizedConfig) -> Tuple[List[Any], Set[Any]]:
341 """Return search dimensions and the set of dimensions with user candidates.
343 List-valued entries in ``config.search_space`` are treated as
344 **output** (search) dimensions. Entries absent from
345 ``config.search_space`` (``"auto"`` in YAML) are also included -- they
346 will be determined by ND's ``bound_space()``.
348 Returns:
349 A tuple ``(dims, candidate_dims)`` where *dims* is the list of
350 ``Dim`` objects to pass to ND and *candidate_dims* is the set of
351 Dim objects for which the user supplied an explicit candidate list
352 (used by :func:`_post_filter`).
353 """
354 dims: List[Any] = []
355 candidate_dims: Set[Any] = set()
356 space = config.search_space
357 for space_key, dim_obj in _search_dim_map().items():
358 candidates = space.get(space_key)
359 if candidates is not None and len(candidates) > 1:
360 dims.append(dim_obj)
361 candidate_dims.add(dim_obj)
362 elif space_key not in space:
363 dims.append(dim_obj)
364 return dims, candidate_dims
367def _filter_by_memory(scored_space: list, budget_gb: float) -> list:
368 """Drop entries whose peak-memory estimate exceeds *budget_gb*.
370 ND already prunes against the device capacity while it generates the
371 space; this is the gate that holds when a strategy is scored under a
372 capacity looser than the budget the caller asked for.
374 Args:
375 scored_space: Scored entries ``(strategy, memory_mb, score, ...)``.
376 budget_gb: Per-device budget in GB; ``0`` disables the gate.
378 Returns:
379 The entries that fit. May be empty, which means the budget rules
380 out every strategy the search found.
381 """
382 if budget_gb <= 0:
383 return list(scored_space)
384 budget_mb = budget_gb * 1024.0
385 kept = [entry for entry in scored_space if float(entry[1]) <= budget_mb]
386 if len(kept) < len(scored_space):
387 logger.info(
388 "Memory filter dropped %d of %d strategies above %.1f GB.",
389 len(scored_space) - len(kept),
390 len(scored_space),
391 budget_gb,
392 )
393 return kept
396def _post_filter(
397 scored_space: list,
398 config: NormalizedConfig,
399 candidate_dims: Optional[Set[Any]] = None,
400) -> list:
401 """Keep only entries whose dimension values are in the user's candidate lists.
403 Args:
404 scored_space: The scored strategy list from ND engine.
405 config: The normalized config containing ``search_space``.
406 candidate_dims: The set of Dim objects that have user-supplied
407 candidate lists with more than one value. If *None*, the
408 set is derived from *config* (backward-compatible).
410 Returns:
411 The entries that fit the memory budget and sit in every list. Empty
412 when none does: a strategy outside the lists is never returned in
413 their place, and the caller refuses the search.
414 """
415 space = config.search_space
416 if candidate_dims is None:
417 candidate_dims = set()
418 for space_key, dim_obj in _search_dim_map().items():
419 candidates = space.get(space_key)
420 if candidates is not None and len(candidates) > 1:
421 candidate_dims.add(dim_obj)
423 candidate_map: Dict[Any, List[int]] = {}
424 for space_key, dim_obj in _search_dim_map().items():
425 if dim_obj not in candidate_dims:
426 continue
427 candidates = space.get(space_key)
428 if candidates is not None:
429 candidate_map[dim_obj] = candidates
431 feasible = _filter_by_memory(scored_space, _memory_budget_gb(config))
433 filtered = []
434 for entry in feasible:
435 dims_val = entry[0].dims_val # type: ignore[index]
436 keep = True
437 for dim_obj, allowed in candidate_map.items():
438 actual = dims_val.get(dim_obj)
439 if actual is not None and actual not in allowed:
440 keep = False
441 break
442 if keep:
443 filtered.append(entry)
444 return filtered
447def _no_candidate_match(feasible: list, config: NormalizedConfig, candidate_dims: Set[Any]) -> str:
448 """Say why no strategy that fits the budget is in the candidate lists.
450 ND searches each dimension within its own bounds and the lists only
451 filter what it found, so a list can ask for degrees the search never
452 reached. Naming the degrees it did reach shows which.
454 Args:
455 feasible: The scored entries that fit the memory budget.
456 config: The normalized config whose ``search_space`` holds the lists.
457 candidate_dims: The dimensions that have a list of their own.
459 Returns:
460 The refusal, naming each list no fitting strategy meets.
461 """
462 unmet = []
463 for space_key, dim_obj in _search_dim_map().items():
464 if dim_obj not in candidate_dims:
465 continue
466 asked = list(config.search_space.get(space_key, []))
467 found = sorted({entry[0].dims_val.get(dim_obj) for entry in feasible} - {None})
468 if not set(found) & set(asked):
469 unmet.append(f"{space_key} asks {asked}, the strategies that fit take {found}")
470 why = "; ".join(unmet) if unmet else "each list is met by some of them, but never all at once"
471 return (
472 f"None of the {len(feasible)} strategies that fit the memory budget is in the candidate "
473 f"lists: {why}. ND searches each dimension within its own bounds and the lists only "
474 "filter what it finds: widen the lists, or the budget."
475 )
478def _pinned_degree(config: NormalizedConfig, constraint_key: str, space_key: str) -> int:
479 """Return the degree of a dimension the search did not vary.
481 Args:
482 config: The normalized config carrying constraint and search space.
483 constraint_key: The ``constraint`` entry that pins the degree.
484 space_key: The ``search_space`` entry listing its candidates.
486 Returns:
487 The pinned degree, else the only candidate, else 1.
488 """
489 pinned = config.constraint.get(constraint_key)
490 if not pinned:
491 candidates = config.search_space.get(space_key) or []
492 pinned = candidates[0] if len(candidates) == 1 else 1
493 return int(pinned)
496def _format_result(best_entry: tuple, config: NormalizedConfig) -> Dict[str, Any]:
497 """Convert the best ND result entry into a flat result dict."""
498 dim_mod = _get_dim_module()
499 dims_val = best_entry[0].dims_val # type: ignore[index]
500 dim_to_key = {
501 dim_mod.DP: "dp",
502 dim_mod.TP: "tp",
503 dim_mod.PP: "pp",
504 dim_mod.CP: "cp",
505 dim_mod.EP: "ep",
506 dim_mod.MBN: "micro_batch_num",
507 dim_mod.OP: "dp_shard",
508 }
509 result: Dict[str, Any] = {
510 "memory_estimate_mb": float(best_entry[1]),
511 "score": float(best_entry[2]),
512 }
513 for dim_obj, key in dim_to_key.items():
514 if dim_obj in dims_val:
515 result[key] = int(dims_val[dim_obj])
516 # A pinned dimension is absent from dims_val, so fill it from what pinned
517 # it: the searcher never varied it, but every consumer still reads it.
518 fixed_from = {
519 "dp": ("fixed_dp_degree", "data_parallel_replicate_degree"),
520 "tp": ("fixed_tp_degree", "tensor_parallel_degree"),
521 "pp": ("fixed_pp_degree", "pipeline_parallel_degree"),
522 "cp": ("fixed_cp_degree", "context_parallel_degree"),
523 "ep": ("fixed_ep_degree", "expert_parallel_degree"),
524 "micro_batch_num": ("fixed_micro_batch_num", "micro_batch_num"),
525 }
526 for key, (constraint_key, space_key) in fixed_from.items():
527 if key not in result:
528 result[key] = _pinned_degree(config, constraint_key, space_key)
529 total_dp = result.get("dp", 1)
530 if "dp_shard" not in result:
531 # OP absent from the searched dimensions: fall back to the declared
532 # degree, which is what the parser used for the whole run.
533 fsdp_candidates = config.search_space.get("data_parallel_shard_degree", [1])
534 configured_fsdp = config.constraint.get("fixed_fsdp_degree")
535 result["dp_shard"] = int(configured_fsdp or fsdp_candidates[0])
536 # The yaml the search prices is AutoModels-shaped, whose FSDP domain is
537 # DP times CP: the shard may span it, and the trainer replicates the rest
538 # (MeshContext.build_meshs). Bounded by DP, a shard the search priced
539 # under CP was cut back and its replicas undercounted.
540 domain = total_dp * max(1, int(result.get("cp") or 1))
541 dp_shard = max(1, min(int(result["dp_shard"]), domain))
542 result["dp_shard"] = dp_shard
543 result["dp_replicate"] = max(1, domain // dp_shard)
544 return result
547def _recompute_result(nd_runner: Any, best_entry: tuple) -> Dict[str, Any]:
548 """The recompute the search chose for the best configuration, for the result.
550 Args:
551 nd_runner: The search, run with recompute "auto" or "per_layer".
552 best_entry: The best scored entry.
554 Returns:
555 ``activation_checkpoint``, the mode the trainer runs every layer
556 with: the chosen one, or ``full``, the policy the configuration was
557 scored with when there was nothing to choose. Under "per_layer",
558 ``activation_checkpoint_layers`` gives the layers that run another
559 mode, as the trainer's ``activation_checkpoint.layers`` states them,
560 where there are any: the resolved yaml writes them where the
561 installed trainer declares that field, and else runs every layer the
562 mode and logs them as advice (``_strategy_output``). With a choice
563 per layer possible,
564 ``recompute_per_layer`` also gives each layer's own fastest option of
565 its kind's front and the score and memory it would reach: what the
566 trainer would gain from running each layer its own way, switch by
567 switch. ``offloaded_layers`` names the layers a choice offloads, as
568 ``"first"`` or ``"first-last"``, where it offloads any: the plan
569 states them off, the trainer running no offload yet.
570 ``offloaded_gib`` gives, by the same names, the GiB each of those
571 layers moves to the host per micro-batch, and ``offloaded_share``
572 the share of what it keeps that is.
573 """
574 # pylint: disable=C0415
575 from hyper_parallel.auto_parallel.sapp_nd.recompute.candidate import offloaded_share, to_records, trainer_plan
576 choice = nd_runner.recompute_choices.get(best_entry[0])
577 mode, layers = trainer_plan(choice) if choice is not None else ("full", {})
578 result: Dict[str, Any] = {"activation_checkpoint": mode}
579 if layers:
580 result["activation_checkpoint_layers"] = layers
581 # Only a choice per layer offloads; one mode for every layer states no ranges of its own.
582 offloaded = [item for item in getattr(choice, "ranges", ()) if item.option.link_bandwidth]
583 if offloaded:
584 names = [str(item.first) if item.count == 1 else f"{item.first}-{item.first + item.count - 1}"
585 for item in offloaded]
586 result["offloaded_layers"] = names
587 result["offloaded_gib"] = {name: item.option.link_bandwidth / 2 ** 30 for name, item in zip(names, offloaded)}
588 result["offloaded_share"] = {name: offloaded_share(item.option) for name, item in zip(names, offloaded)}
589 per_layer, score = nd_runner.recompute_per_layer(best_entry[0])
590 if per_layer is not None:
591 result["recompute_per_layer"] = {
592 "score": float(score),
593 "memory_estimate_mb": float(per_layer.memory),
594 "ranges": to_records(per_layer),
595 }
596 return result
599def search_strategies(config: NormalizedConfig, offload: Optional[Mapping[str, Any]] = None) -> Dict[str, Any]:
600 """Run the ND strategy search and return the optimal strategy.
602 This is the main entry point for end-to-end strategy search:
604 1. Validates required model fields.
605 2. Shapes the ``NormalizedConfig`` into the cost-model config ND
606 parses, as a mapping rather than a file.
607 3. Launches the ND search engine (:class:`Parallelize`).
608 4. Post-filters results against the user's candidate lists.
609 5. Returns the best strategy as a flat dictionary.
611 Args:
612 config: A fully populated ``NormalizedConfig`` from
613 :func:`read_search_config` or :func:`read_hp_yaml_config`.
614 ``config.estimator["ratios_file"]``, when set, names the ratios
615 file ``run_nd --write_ratios`` fits, which ND scores with
616 (``run_nd -s ... -c`` sets it); without it ND scores raw.
617 offload: With a choice of a recompute mode per layer, let each
618 stage's first layers run off and offload what they keep, over the
619 device's host link with these of its figures replaced
620 (:meth:`HostLink.of`), ``None`` for one left as the device states
621 it. ``None`` offloads nothing.
623 Returns:
624 A dict with keys ``dp``, ``tp``, ``pp``, ``cp``, ``ep``,
625 ``micro_batch_num``, ``memory_estimate_mb``, and ``score``.
627 Raises:
628 ValueError: If required fields are missing or no strategy is found,
629 or for *offload* without a choice of a mode per layer.
630 """
631 _validate_before_search(config)
633 hp_config = _build_hp_yaml_dict(config)
634 machine = _build_machine(config)
635 dims, candidate_dims = _resolve_search_dimensions(config)
637 import hyper_parallel.auto_parallel.sapp_nd.nd.parallelize as _Par # pylint: disable=C0415
638 from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import ( # pylint: disable=C0415
639 HYPER_SELECTIVE_REC_OP,
640 )
641 strategy = config.estimator.get("recompute_strategy")
642 auto = strategy in CHOSEN_RECOMPUTE
643 census = bool((hp_config.get("context") or {}).get("census"))
644 trainer = {
645 "auto_recompute": True,
646 # The search config may narrow the modes, as to leave out a mode the
647 # estimate prices faster than the device runs it.
648 "recompute_modes": tuple(config.estimator.get("recompute_modes")
649 or (TRAINER_CENSUS_MODES if census else TRAINER_RECOMPUTE_MODES)),
650 # The trainer's selective mode runs its policy, whatever switches
651 # the mode the model is described with sets.
652 "recompute_selective": dict(HYPER_SELECTIVE_REC_OP),
653 "recompute_mode_per_layer": strategy == "per_layer",
654 }
655 if offload is not None:
656 if strategy != "per_layer":
657 raise ValueError("offload is priced in a choice of a recompute mode per layer: give the search config "
658 "recompute: per_layer")
659 from hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware import HostLink # pylint: disable=C0415
660 trainer.update(auto_offload=True, host_link=HostLink.of(machine.device.host_link, offload))
661 nd_runner = _Par.Parallelize(
662 "hyper_v2",
663 hp_config,
664 machine,
665 global_batch_size=config.constraint.get("global_batch_size", 0),
666 dimensions=dims,
667 **(trainer if auto else {}),
668 )
669 scored_space = nd_runner.run_generation_to_ordering(
670 yaml_folder=None,
671 threads_num=None,
672 top_num=None,
673 cache_file=config.estimator.get("ratios_file"),
674 )
676 if not scored_space:
677 raise ValueError("ND search returned no valid strategies.")
679 filtered = _post_filter(scored_space, config, candidate_dims)
680 if not filtered:
681 feasible = _filter_by_memory(scored_space, _memory_budget_gb(config))
682 if feasible:
683 raise ValueError(_no_candidate_match(feasible, config, candidate_dims))
684 lightest_gb = min(float(entry[1]) for entry in scored_space) / 1024.0
685 raise ValueError(
686 "No strategy fits the memory budget of "
687 f"{_memory_budget_gb(config):.1f} GB; the lightest strategy "
688 f"found needs {lightest_gb:.1f} GB. Raise "
689 "constraint.memory_limit_gb or widen the search space."
690 )
691 best = filtered[0]
692 result = _format_result(best, config)
693 if auto:
694 result.update(_recompute_result(nd_runner, best))
695 else:
696 # Without "auto" the search prices every candidate fully recomputed,
697 # whatever the search yaml says, so the trainer runs what was priced.
698 result["activation_checkpoint"] = "full"
700 logger.info(
701 "Optimal strategy found: dp=%(dp)s tp=%(tp)s pp=%(pp)s "
702 "cp=%(cp)s ep=%(ep)s mb_num=%(micro_batch_num)s "
703 "recompute=%(activation_checkpoint)s "
704 "mem=%(memory_estimate_mb).0f MB score=%(score).2e",
705 result,
706 )
707 return result