Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / _exec_spec.py: 95%
195 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"""The execution IR: how a model is trained, beside what the model is.
17:class:`~hyper_parallel.auto_parallel._model_spec.ModelSpec` states the facts
18about an architecture that hold whatever strategy it runs under.
19:class:`ExecSpec` states the rest: what a search varies (the parallel
20degrees, the micro-batching, the sharding and the recompute policy) and the
21fixed facts of the run (the sequence length, the precisions, the kernels and
22the device's memory). The cost model derives everything else from the two.
24It follows the model spec's two rules. **Absent is not zero**: every field
25defaults to ``None``, which means "not stated", so an ExecSpec can state a
26whole run or only what a search candidate changes, and a consumer never
27mistakes an unstated degree for a degree of zero. **Refuse rather than
28default**: :meth:`ExecSpec.validate` rejects a stated value that cannot be
29right, and :meth:`ExecSpec.from_dict` refuses a key it does not know, by name.
30"""
31from dataclasses import dataclass, fields
32from typing import Any, Dict, Mapping, Optional, Tuple, Union
35class ExecSpecError(ValueError):
36 """An execution spec states something it cannot, or something unknown."""
39# The rules by which selective recompute picks the ops it recomputes.
40SELECTIVE_RULES = ("hyperparallel", "mindformers")
42# How a range of layers recomputes.
43RECOMPUTE_OPTIONS = ("none", "full", "selective")
44# The ops a selective range decides for, the cost model's recompute switches,
45# and the state it gives each: the op's output is kept for backward, or
46# recomputed. Offload will add a third state. attUp is an MLA layer's
47# up-projections, which build its query, key and value heads from their
48# latents.
49RECOMPUTE_OPS = ("attBMM", "headCast", "dropout", "softmax", "normOp", "gather", "ffAct", "attUp")
50OP_STATES = ("keep", "recompute")
53def _whole(value: Any, least: int) -> bool:
54 """Whether *value* is a whole number of at least *least*."""
55 return isinstance(value, int) and not isinstance(value, bool) and value >= least
58@dataclass(frozen=True)
59class RecomputeRange:
60 """How consecutive layers recompute, every layer of the range alike.
62 Layers are counted in model order from 0, the body layers first and then
63 the MTP layers, whatever stage each is placed on.
65 Attributes:
66 first: The range's first layer.
67 count: How many layers it covers, or ``None`` for every layer from
68 *first* on, so that ``RecomputeRange(option="full")`` recomputes
69 the whole model.
70 option: One of :data:`RECOMPUTE_OPTIONS`.
71 ops: For a selective range, the state of each op it names, one of
72 :data:`OP_STATES`; an op it does not name is kept. ``None``
73 takes the selective recompute of the run's rule
74 (``selective_rule``).
75 """
77 first: int = 0
78 count: Optional[int] = None
79 option: str = "none"
80 ops: Optional[Mapping[str, str]] = None
82 def validate(self) -> "RecomputeRange":
83 """Raise :class:`ExecSpecError` unless the range can be right; return it."""
84 if not _whole(self.first, 0):
85 raise ExecSpecError(f"a recompute range's first layer must be a whole number from 0, got {self.first!r}")
86 if self.count is not None and not _whole(self.count, 1):
87 raise ExecSpecError(f"a recompute range's count must be at least 1, got {self.count!r}")
88 if self.option not in RECOMPUTE_OPTIONS:
89 raise ExecSpecError(f"a recompute option must be one of {list(RECOMPUTE_OPTIONS)}, got {self.option!r}")
90 if self.ops is not None:
91 if self.option != "selective":
92 raise ExecSpecError(f"only a selective range states ops, not a {self.option!r} one")
93 unknown = sorted(set(self.ops) - set(RECOMPUTE_OPS))
94 if unknown:
95 raise ExecSpecError(f"unknown recompute ops {unknown}; the ops are {list(RECOMPUTE_OPS)}")
96 states = sorted({str(state) for state in self.ops.values()} - set(OP_STATES))
97 if states:
98 raise ExecSpecError(f"an op's state must be one of {list(OP_STATES)}, got {states}")
99 return self
101 def switches(self) -> Optional[Dict[str, int]]:
102 """Return every op's switch, 1 to keep and 0 to recompute, or ``None`` for the rule's."""
103 if self.ops is None:
104 return None
105 return {op: int(self.ops.get(op, "keep") == "keep") for op in RECOMPUTE_OPS}
107 def to_dict(self) -> Dict[str, Any]:
108 """Return the range as a plain mapping, ready for YAML."""
109 out: Dict[str, Any] = {"first": self.first}
110 if self.count is not None:
111 out["count"] = self.count
112 out["option"] = self.option
113 if self.ops is not None:
114 out["ops"] = dict(self.ops)
115 return out
117 @classmethod
118 def from_dict(cls, data: Any) -> "RecomputeRange":
119 """Build a validated range from a mapping, refusing an unknown key."""
120 if not isinstance(data, Mapping):
121 raise ExecSpecError(f"a recompute range must be a mapping, got {data!r}")
122 unknown = sorted(set(data) - {"first", "count", "option", "ops"})
123 if unknown:
124 raise ExecSpecError(f"unknown recompute range keys {unknown}; a range has first, count, option and ops")
125 ops = data.get("ops")
126 if ops is not None and not isinstance(ops, Mapping):
127 raise ExecSpecError(f"a recompute range's ops must map ops to states, got {ops!r}")
128 return cls(
129 first=data.get("first", 0),
130 count=data.get("count"),
131 option=str(data.get("option", "none")),
132 ops=None if ops is None else {str(op): str(state) for op, state in ops.items()},
133 ).validate()
136def _check_ranges(ranges: Any) -> None:
137 """Raise :class:`ExecSpecError` unless *ranges* run in model order without overlap."""
138 if not isinstance(ranges, tuple):
139 raise ExecSpecError(f"recompute must be a tuple of recompute ranges, got {ranges!r}")
140 end = 0
141 for index, item in enumerate(ranges):
142 if not isinstance(item, RecomputeRange):
143 raise ExecSpecError(f"recompute must hold recompute ranges, got {item!r}")
144 item.validate()
145 if item.first < end:
146 raise ExecSpecError(
147 f"recompute ranges must run in model order without overlap: layer {item.first} "
148 f"is covered already"
149 )
150 if item.count is None and index != len(ranges) - 1:
151 raise ExecSpecError("only the last recompute range may run to the last layer")
152 end = item.first + (item.count or 0)
155@dataclass(frozen=True)
156class ExecSpec:
157 """How a model is trained: the strategy a search varies, and the run.
159 Every field is optional; ``None`` means the spec does not state it. A
160 byte width, ``grad_accumulation`` or ``shard_activations`` that no spec
161 states takes the default of the model's family, from its op profile.
163 Attributes:
164 dp, tp, pp, vpp, cp, ep: The data, tensor, pipeline, virtual
165 pipeline, context and expert parallel degrees.
166 etp: The expert tensor-parallel degree; at most 1 means an expert
167 layer runs with the dense layers' tensor parallelism.
168 sequence_parallel: Whether activations are split along the sequence
169 over the tensor-parallel group.
170 shard_activations: Whether tensor parallelism shards the activations
171 between layers, which a fully recomputed layer keeps as its
172 input, and the output layer's.
173 loss_parallel: Whether the loss runs on logits sharded over the
174 vocabulary across the tensor-parallel group, each rank's share,
175 rather than on logits every rank gathers whole, as
176 HyperParallel's trainer does unless its ``loss_parallel``.
177 micro_batch_size, micro_batch_num, global_batch_size: The batch.
178 optimizer_parallel: Whether optimizer states are sharded.
179 optimizer_shard: How many data-parallel ranks optimizer sharding
180 splits each parameter over, on top of TP's split, as MindSpore's
181 ``optimizer_weight_shard_size`` and HyperParallel's ``dp_shard``
182 count them.
183 expert_shard: How many ranks of its expert data-parallel group FSDP
184 shards a routed expert over under expert parallelism, as
185 HyperParallel's ``edp_shard_size`` counts them; without expert
186 parallelism the experts shard with the other parameters. Stated
187 by no one, the optimizer shards them over the whole group.
188 expert_shard_group: Whether a launcher shards each strategy's routed
189 experts over its whole expert data-parallel group, whatever
190 ``expert_shard`` states, as ``context.expert_shard: group`` says.
191 grad_shard: Whether gradients are sharded too.
192 grad_shard_as_params: Whether each gradient is sharded as its
193 parameter is, as FSDP holds it; ``grad_shard`` does not apply.
194 shard_spans_cp: Whether the optimizer's shard spans context
195 parallelism's ranks too, as HyperParallel's ``dp_shard_size``
196 does over its ``dp * cp`` FSDP domain, so that CP shards no
197 parameter state beyond it and the tables it gathers are whole;
198 unstated, CP shards a parameter's state again.
199 grad_accumulation: Whether gradients take memory without pipeline
200 parallelism too; under it they always do.
201 deferred_grad_accumulation: Whether FSDP holds each layer's
202 reduce-scatter output until the micro-batch's backward ends, and
203 only then adds it to the accumulated gradient, as HyperParallel's
204 does; PyTorch's FSDP2 adds it as soon as it is reduced.
205 overlapped_grad_reduce: Whether FSDP reduce-scatters a layer's
206 gradients while the next layer's backward runs, holding them
207 whole until that one ends, and the root's, the embedding and
208 output tables', until the backward ends, as HyperParallel's does.
209 reshard_params: Whether FSDP frees a layer's gathered parameters
210 once the layer has run, forward or backward, and gathers them
211 again when it runs next, as HyperParallel's does by default;
212 the root's, the embedding and output tables, stay gathered.
213 param_casts: Whether a layer keeps, beside each matmul's input, the
214 cast the formulas price where the optimizer does not shard
215 (``n_attParamCast``, ``n_ffParamCast``); unstated, it does
216 exactly there. HyperParallel's FSDP gathers its weights in the
217 compute dtype, and its layers keep none at any ``dp_shard``.
218 pp_schedule: The pipeline schedule, such as ``1f1b``.
219 offset: How many layers each pipeline stage holds beyond an even
220 split, per stage or per chunk and stage.
221 seq_split: How many parts a sequence is split into for a sequence
222 pipeline.
223 pp_partition: The layers each stage holds, when stated outright.
224 mtp_in_offset, emb_out_in_offset: Whether the MTP layers, and the
225 embedding and output layers, count in the offset's layers.
226 full_recompute, selective_recompute: Which layers are recomputed in
227 full or selectively: true for all, or counts per stage.
228 recompute: How each layer recomputes, as :class:`RecomputeRange`
229 ranges over the layers in model order. Stated, it replaces the
230 per-stage forms above; a layer no range covers is not
231 recomputed.
232 selective_comm_recompute: Which layers recompute their
233 sequence-parallel all-gather, under the MindFormers rule.
234 selective_rule: Whose selective recompute the run uses, one of
235 :data:`SELECTIVE_RULES`.
236 recompute_slice_activation: Whether a recomputed layer keeps its
237 input sliced over tensor parallelism.
238 param_bytes, compute_bytes, softmax_bytes, grad_bytes,
239 optimizer_state_bytes, norm_bytes, dropout_bytes: Bytes per
240 element of the parameters, activations, softmax outputs,
241 gradients, optimizer states, norm activations and dropout masks.
242 optimizer_states: How many states the optimizer keeps per parameter
243 of a layer: AdamW's two moments, or Muon's one momentum. The
244 embedding and output tables keep AdamW's two.
245 main_param_bytes: Bytes per element of the copy of the parameters
246 the optimizer keeps, 0 when it keeps none, as HyperParallel's
247 fp32 main parameters keep one.
248 flash_attention, grad_clip, grouped_gemm, tie_embeddings: Kernels
249 and features the run uses.
250 vocab_emb_dp: Whether the vocabulary embedding runs data parallel.
251 emb_dp_sharded: Whether the embedding table is sharded over data
252 parallelism.
253 cp_algo: The context-parallel algorithm.
254 capacity_factor: The experts' token capacity factor.
255 shard_mtp_param: Whether the MTP layers' parameters are sharded.
256 frozen: Whether the model's parameters are frozen.
257 optimizer: The optimizer's name.
258 seq_length: The training sequence length.
259 device_memory: The device's memory, such as ``"64GB"``.
260 """
262 dp: Optional[int] = None
263 tp: Optional[int] = None
264 pp: Optional[int] = None
265 vpp: Optional[int] = None
266 cp: Optional[int] = None
267 ep: Optional[int] = None
268 etp: Optional[int] = None
269 sequence_parallel: Optional[bool] = None
270 shard_activations: Optional[bool] = None
271 loss_parallel: Optional[bool] = None
273 micro_batch_size: Optional[int] = None
274 micro_batch_num: Optional[int] = None
275 global_batch_size: Optional[int] = None
277 optimizer_parallel: Optional[bool] = None
278 optimizer_shard: Optional[int] = None
279 expert_shard: Optional[int] = None
280 expert_shard_group: Optional[bool] = None
281 grad_shard: Optional[bool] = None
282 grad_shard_as_params: Optional[bool] = None
283 shard_spans_cp: Optional[bool] = None
284 grad_accumulation: Optional[bool] = None
285 deferred_grad_accumulation: Optional[bool] = None
286 overlapped_grad_reduce: Optional[bool] = None
287 reshard_params: Optional[bool] = None
288 param_casts: Optional[bool] = None
290 pp_schedule: Optional[str] = None
291 offset: Optional[Union[int, list]] = None
292 seq_split: Optional[int] = None
293 pp_partition: Optional[list] = None
294 mtp_in_offset: Optional[bool] = None
295 emb_out_in_offset: Optional[bool] = None
297 full_recompute: Optional[Union[bool, int, list]] = None
298 selective_recompute: Optional[Union[bool, list]] = None
299 selective_comm_recompute: Optional[Union[bool, list]] = None
300 selective_rule: Optional[str] = None
301 recompute_slice_activation: Optional[bool] = None
302 recompute: Optional[Tuple[RecomputeRange, ...]] = None
304 param_bytes: Optional[int] = None
305 compute_bytes: Optional[int] = None
306 softmax_bytes: Optional[int] = None
307 grad_bytes: Optional[int] = None
308 optimizer_state_bytes: Optional[int] = None
309 optimizer_states: Optional[int] = None
310 main_param_bytes: Optional[int] = None
311 norm_bytes: Optional[int] = None
312 dropout_bytes: Optional[int] = None
314 flash_attention: Optional[bool] = None
315 grad_clip: Optional[bool] = None
316 grouped_gemm: Optional[bool] = None
317 vocab_emb_dp: Optional[bool] = None
318 emb_dp_sharded: Optional[bool] = None
319 tie_embeddings: Optional[bool] = None
320 cp_algo: Optional[str] = None
321 capacity_factor: Optional[float] = None
322 shard_mtp_param: Optional[bool] = None
323 frozen: Optional[bool] = None
324 optimizer: Optional[str] = None
326 seq_length: Optional[int] = None
327 device_memory: Optional[Any] = None
329 # ---- validation ----------------------------------------------------
331 def validate(self) -> "ExecSpec":
332 """Raise :class:`ExecSpecError` unless every stated value can be right.
334 Returns *self*, so a producer can ``return spec.validate()``.
335 """
336 for name in _KINDS["count"]:
337 value = getattr(self, name)
338 if value is not None and value < 1:
339 raise ExecSpecError(f"{name} must be at least 1, got {value!r}")
340 for name in _KINDS["size"]:
341 value = getattr(self, name)
342 if value is not None and value < 0:
343 raise ExecSpecError(f"{name} must not be negative, got {value!r}")
344 if self.selective_rule is not None and self.selective_rule not in SELECTIVE_RULES:
345 raise ExecSpecError(
346 f"selective_rule must be one of {list(SELECTIVE_RULES)}, got {self.selective_rule!r}"
347 )
348 if self.capacity_factor is not None and self.capacity_factor <= 0:
349 raise ExecSpecError(f"capacity_factor must be positive, got {self.capacity_factor!r}")
350 if self.recompute is not None:
351 _check_ranges(self.recompute)
352 return self
354 # ---- serialised form -----------------------------------------------
356 def to_dict(self) -> Dict[str, Any]:
357 """Return the stated fields as a plain mapping, ready for YAML."""
358 out: Dict[str, Any] = {}
359 for spec_field in fields(self):
360 value = getattr(self, spec_field.name)
361 if value is None:
362 continue
363 if spec_field.name == "recompute":
364 value = [item.to_dict() for item in value]
365 out[spec_field.name] = value if spec_field.name != "device_memory" else str(value)
366 return out
368 @classmethod
369 def from_dict(cls, data: Mapping[str, Any]) -> "ExecSpec":
370 """Build a validated spec from a mapping, refusing an unknown key.
372 Raises:
373 ExecSpecError: If a key is unknown, a value has the wrong type,
374 or a stated value cannot be right.
375 """
376 if not isinstance(data, Mapping):
377 raise ExecSpecError(f"an execution spec must be a mapping, got {data!r}")
378 known = [spec_field.name for spec_field in fields(cls)]
379 unknown = sorted(set(data) - set(known))
380 if unknown:
381 raise ExecSpecError(f"unknown execution spec keys {unknown}; the keys are {known}")
382 kwargs = {key: _coerce(key, value) for key, value in data.items()}
383 return cls(**kwargs).validate()
386# The fields by kind, for coercion from YAML and for validation; a field in
387# none of them is kept as it is.
388_KINDS: Dict[str, tuple] = {
389 "count": ("dp", "tp", "pp", "vpp", "cp", "ep", "micro_batch_size", "micro_batch_num",
390 "global_batch_size", "optimizer_shard", "expert_shard", "seq_split", "seq_length",
391 "optimizer_states"),
392 "size": ("etp", "param_bytes", "compute_bytes", "softmax_bytes", "grad_bytes",
393 "optimizer_state_bytes", "main_param_bytes", "norm_bytes", "dropout_bytes"),
394 "flag": ("sequence_parallel", "shard_activations", "loss_parallel", "optimizer_parallel", "grad_shard",
395 "grad_shard_as_params", "shard_spans_cp", "grad_accumulation", "deferred_grad_accumulation",
396 "overlapped_grad_reduce", "reshard_params", "param_casts", "expert_shard_group",
397 "mtp_in_offset",
398 "emb_out_in_offset",
399 "recompute_slice_activation", "flash_attention", "grad_clip", "grouped_gemm",
400 "vocab_emb_dp", "emb_dp_sharded", "tie_embeddings", "shard_mtp_param", "frozen"),
401 "name": ("pp_schedule", "selective_rule", "cp_algo", "optimizer", "device_memory"),
402}
405def _as_number(value: Any) -> float:
406 """Coerce the capacity factor to a number."""
407 try:
408 return float(value)
409 except (TypeError, ValueError) as exc:
410 raise ExecSpecError(f"capacity_factor must be a number, got {value!r}") from exc
413def _as_ranges(value: Any) -> Tuple[RecomputeRange, ...]:
414 """Coerce a list of recompute ranges, each a range or its mapping."""
415 if not isinstance(value, (list, tuple)):
416 raise ExecSpecError(f"recompute must list recompute ranges, got {value!r}")
417 return tuple(item if isinstance(item, RecomputeRange) else RecomputeRange.from_dict(item) for item in value)
420# The fields with a coercion of their own.
421_COERCIONS = {"capacity_factor": _as_number, "recompute": _as_ranges}
424def _coerce(key: str, value: Any) -> Any:
425 """Coerce one field read from YAML to the type the schema gives it."""
426 if value is None:
427 return None
428 if key in _KINDS["count"] or key in _KINDS["size"]:
429 if isinstance(value, bool) or (isinstance(value, float) and not value.is_integer()):
430 raise ExecSpecError(f"{key} must be a whole number, got {value!r}")
431 try:
432 return int(value)
433 except (TypeError, ValueError) as exc:
434 raise ExecSpecError(f"{key} must be a whole number, got {value!r}") from exc
435 if key in _KINDS["flag"]:
436 if not isinstance(value, bool):
437 raise ExecSpecError(f"{key} must be true or false, got {value!r}")
438 return value
439 if key in _KINDS["name"]:
440 return str(value)
441 if key in _COERCIONS:
442 return _COERCIONS[key](value)
443 return value