Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / common / apply_exec.py: 99%
68 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"""How an execution spec reaches a cost-model config.
17:func:`apply_exec` is the one way a strategy reaches a config: it writes the
18fields an :class:`~hyper_parallel.auto_parallel._exec_spec.ExecSpec` states
19and derives the rest. :func:`exec_of` reads a config's ExecSpec back, for a
20parser that fills the config itself, and :func:`strategy_exec` is the
21ExecSpec a search's keyword strategy states. A config a search owns refuses
22any other write of a degree, but :func:`apply_layer_strategy`, which gives
23the layer about to be priced the degrees its kind runs with.
24"""
25from typing import Any, List, Mapping, Optional, Tuple
27from hyper_parallel.auto_parallel._exec_spec import RECOMPUTE_OPS, ExecSpec, RecomputeRange
28from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import derive
29from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_order import get_model_order, stated_recompute
30from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
32# Each ExecSpec field, and the config field that holds it.
33CONFIG_FIELDS = {
34 "dp": "d",
35 "tp": "t",
36 "pp": "p",
37 "vpp": "vp",
38 "cp": "cp",
39 "ep": "ep",
40 "etp": "etp",
41 "sequence_parallel": "sequence_parallel",
42 "shard_activations": "shard_activations",
43 "loss_parallel": "loss_parallel",
44 "micro_batch_size": "b",
45 "micro_batch_num": "m",
46 "global_batch_size": "gbs",
47 "optimizer_parallel": "has_op",
48 "optimizer_shard": "os_max_shard",
49 "expert_shard": "expert_shard",
50 "expert_shard_group": "expert_shard_group",
51 "grad_shard": "has_grad_shard",
52 "grad_shard_as_params": "grad_shard_as_params",
53 "shard_spans_cp": "shard_spans_cp",
54 "grad_accumulation": "grad_accumulation",
55 "reshard_params": "reshard_params",
56 "deferred_grad_accumulation": "deferred_grad_accumulation",
57 "overlapped_grad_reduce": "overlapped_grad_reduce",
58 "param_casts": "param_casts",
59 "pp_schedule": "pp_sched",
60 "offset": "offset",
61 "seq_split": "n_s_split",
62 "pp_partition": "pp_partition",
63 "mtp_in_offset": "is_mtp_in_offset",
64 "emb_out_in_offset": "emb_out_in_offset",
65 "full_recompute": "full_rec",
66 "selective_recompute": "sel_rec",
67 "selective_comm_recompute": "sel_comm_rec",
68 "selective_rule": "sel_rec_rule",
69 "recompute_slice_activation": "recompute_slice_activation",
70 "recompute": "recompute_ranges",
71 "param_bytes": "bytes_p",
72 "compute_bytes": "bytes_compute",
73 "softmax_bytes": "bytes_softmax",
74 "grad_bytes": "grad_bytes",
75 "optimizer_state_bytes": "optimizer_state_bytes",
76 "optimizer_states": "optimizer_states",
77 "main_param_bytes": "main_param_bytes",
78 "norm_bytes": "norm_bytes",
79 "dropout_bytes": "dropout_bytes",
80 "flash_attention": "has_fa",
81 "grad_clip": "has_clip",
82 "grouped_gemm": "gmm",
83 "vocab_emb_dp": "vocab_emb_dp",
84 "emb_dp_sharded": "emb_dp_sharded",
85 "tie_embeddings": "tie_emb_out",
86 "cp_algo": "cp_algo",
87 "capacity_factor": "cap_fact",
88 "shard_mtp_param": "is_shard_mtp_param",
89 "frozen": "freeze",
90 "optimizer": "optimizer",
91 "seq_length": "s",
92 "device_memory": "device_capacity",
93}
95# The keyword strategy of ``CostModelConfig.set_strategy``: each key that
96# states an integer, and the ExecSpec field it states.
97STRATEGY_KEYS = (
98 ("dp", "dp"),
99 ("mp", "tp"),
100 ("ep", "ep"),
101 ("etp", "etp"),
102 ("cp", "cp"),
103 ("vpp", "vpp"),
104 ("pp", "pp"),
105 ("mb", "micro_batch_num"),
106 ("mbs", "micro_batch_size"),
107)
110# The recompute option of each layer type the partition generator places.
111_OPTIONS = {LayerType.FULL_REC_LAYER: "full", LayerType.SEL_REC_LAYER: "selective"}
114def _ranges(options: List[str], ops: Mapping[str, str]) -> Tuple[RecomputeRange, ...]:
115 """Return the runs of *options*, one layer each in model order, as ranges.
117 A layer that does not recompute is left out, since a layer no range
118 covers is not recomputed.
119 """
120 out = []
121 first = 0
122 for index in range(1, len(options) + 1):
123 if index < len(options) and options[index] == options[first]:
124 continue
125 if options[first] != "none":
126 out.append(RecomputeRange(
127 first=first, count=index - first, option=options[first],
128 ops=dict(ops) if options[first] == "selective" else None,
129 ))
130 first = index
131 return tuple(out)
134def recompute_of(ccfg: Any) -> Optional[Tuple[RecomputeRange, ...]]:
135 """Return how a config's layers recompute, as ranges in model order.
137 The ranges it states, else the option its partition gives each layer
138 from ``full_rec`` and ``sel_rec``, with the switches of its selective
139 recompute as each op's state. ``None`` for a config the partition
140 generator does not lay out alone, such as a multimodal parent.
141 """
142 stated = stated_recompute(ccfg)
143 if stated is not None:
144 return stated
145 layout = getattr(ccfg, "generate_partitions_vpp_unimodal", None)
146 if not callable(layout) or getattr(ccfg, "multimodal", False) or not ccfg.n_lay:
147 return None
148 stages = layout()
149 options = [
150 _OPTIONS.get(stages[stage_id][chunk_id][lay_id], "none")
151 for stage_id, chunk_id, lay_id in get_model_order(ccfg, stages)
152 ]
153 # From vars: a Config answers 0, recompute, for a switch it lacks.
154 stated = vars(ccfg.rec_op) if ccfg.rec_op is not None else {}
155 ops = {op: "keep" if stated.get(op, 1) else "recompute" for op in RECOMPUTE_OPS}
156 return _ranges(options, ops)
159def exec_of(ccfg: Any) -> ExecSpec:
160 """Read back the ExecSpec a config holds, every field it states.
162 Its recompute comes back as ranges too (:func:`recompute_of`), whichever
163 form the config states it in.
164 """
165 stated = {name: getattr(ccfg, field) for name, field in CONFIG_FIELDS.items()}
166 stated["recompute"] = recompute_of(ccfg)
167 return ExecSpec(**stated)
170def apply_exec(ccfg: Any, exec_spec: ExecSpec, strict: bool = True) -> None:
171 """Write the fields *exec_spec* states on *ccfg*, then derive the rest.
173 A field the spec does not state keeps the value the config holds.
175 Args:
176 ccfg: The config to change.
177 exec_spec: What to change it to.
178 strict: As in :func:`~hyper_parallel.auto_parallel.sapp_nd.nd.common.derive.derive`.
179 """
180 for name, field in CONFIG_FIELDS.items():
181 value = getattr(exec_spec, name)
182 if value is None:
183 continue
184 if name == "device_memory" and isinstance(value, str):
185 value = _memory(value)
186 # The one sanctioned write of a strategy field, past the guard.
187 object.__setattr__(ccfg, field, value)
188 if exec_spec.recompute is None and (
189 exec_spec.full_recompute is not None or exec_spec.selective_recompute is not None
190 ):
191 # A per-stage recompute replaces the ranges the config states.
192 object.__setattr__(ccfg, "recompute_ranges", None)
193 derive(ccfg, strict)
196def apply_layer_strategy(ccfg: Any, degrees: Mapping[str, Any]) -> None:
197 """Give the layer about to be priced the degrees its kind runs with.
199 A dense kind in a MoE model runs with no expert parallelism. Only the
200 degrees are written: the fields derived from the model's degrees stay as
201 the estimators read them.
203 Args:
204 ccfg: The config the layer is priced on.
205 degrees: The strategy fields the kind states.
206 """
207 for name, value in degrees.items():
208 object.__setattr__(ccfg, name, value)
211def strategy_exec(ccfg: Any, strategy: Mapping[str, Any]) -> ExecSpec:
212 """Return the ExecSpec a keyword strategy states, over the config it changes.
214 A degree or a micro-batching is stated when it is given as an integer,
215 and ``op`` states the optimizer sharding. Sequence parallelism is the
216 candidate's where it states ``sp``, else the run's, and on at every TP
217 degree for a run that states none; the global batch follows the
218 micro-batching the strategy leaves. The recompute layers and the offset
219 are stated when given.
221 Args:
222 ccfg: The config the strategy changes, which holds what it leaves.
223 strategy: ``set_strategy``'s keywords.
224 """
225 stated = {
226 name: strategy[key] for key, name in STRATEGY_KEYS if isinstance(strategy.get(key), int)
227 }
228 op = strategy.get("op")
229 if isinstance(op, int):
230 stated["optimizer_shard"] = op
231 # op <= 1 means no optimizer sharding
232 stated["optimizer_parallel"] = op > 1
233 # Every candidate ran sequence parallelism, so a run without it had the
234 # activations TP does not shard divided by TP all the same, and a search
235 # over SP priced its twins alike.
236 sp = strategy.get("sp")
237 run_sp = getattr(ccfg, "sequence_parallel", None)
238 stated["sequence_parallel"] = bool(sp) if sp is not None else run_sp is None or bool(run_sp)
239 stated["global_batch_size"] = (
240 stated.get("micro_batch_size", ccfg.b) * stated.get("dp", ccfg.d)
241 * stated.get("micro_batch_num", ccfg.m)
242 )
243 if strategy.get("full_rec") is not None:
244 stated["full_recompute"] = strategy["full_rec"]
245 if strategy.get("sel_rec") is not None:
246 stated["selective_recompute"] = strategy["sel_rec"]
247 if isinstance(strategy.get("offset"), (int, list)):
248 stated["offset"] = strategy["offset"]
249 return ExecSpec(**stated)
252def _memory(text: str) -> Any:
253 """The cost model's memory size for a string such as ``"64GB"``."""
254 # The memory estimation package imports the config module, which imports
255 # this one: importing it here keeps this module importable first.
256 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory # pylint: disable=import-outside-toplevel
257 return Memory.from_string(text)