Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / _model_spec.py: 98%
357 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 model IR: a typed, serialisable description of what is being trained.
17``memory_estimation/README.md`` already calls ``CostModelConfig`` "the
18interface between framework input files and the memory module". It cannot
19be one, because it carries derived values, fuses model shape with parallel
20strategy, and has no serialised form -- so the interface gets faked, most
21visibly by writing a temporary YAML to disk and parsing it back.
23:class:`ModelSpec` is that interface made real for the model half: the facts
24about an architecture that hold whatever strategy it is trained under. It is
25a schema, not a grammar. Three producers fill it -- a framework parser, a
26census of an instantiated model, or an engineer writing YAML by hand -- and
27every consumer downstream reads the same fields.
29Two rules make it safe to hand-fill, both learned from Qwen3.5-35B-A3B, which
30reached the cost model with its MTP layer dropped, its shared expert priced at
31width zero, and half its attention parameters missing, all silently:
33- **Absent is not zero.** Every optional field defaults to ``None``. A width
34 nothing supplied can then be told apart from a width declared as ``0``,
35 which is what let an all-MoE config price its shared expert at no width at
36 all and then divide by it.
37- **Refuse rather than default.** :meth:`ModelSpec.validate` rejects a spec
38 whose fields contradict each other, at parse time and by name, instead of
39 resolving to a number that is wrong further downstream.
41The strategy half (degrees, sharding, recompute, precision) is the execution
42spec's, :class:`~hyper_parallel.auto_parallel._exec_spec.ExecSpec`. A key that
43states neither the model nor the run is refused.
44"""
45from dataclasses import dataclass, fields
46from typing import Any, Dict, Mapping, Optional, Tuple
49# The model facts a spec states as true or false.
50_FLAGS = ("attn_output_gate", "tie_word_embeddings", "qk_norm", "qkv_bias", "o_bias", "mlp_bias", "norm_bias",
51 "shared_expert_gate")
53# The legacy keys that state a run's precision. Refused as model facts like
54# any unknown key, but named with the keys a run reads instead: a description
55# stating one among its facts was once priced at the default width.
56_RUN_PRECISION_KEYS = ("compute_dtype", "param_init_type")
59class ModelSpecError(ValueError):
60 """A model spec is missing a required fact or contradicts itself.
62 Raised at parse time, by field name, so a producer learns what it failed
63 to supply instead of a consumer inheriting a zero.
64 """
67def _as_int(value: Any, name: str) -> Optional[int]:
68 """Coerce *value* to ``int``, or raise naming the field it came from."""
69 if value is None:
70 return None
71 try:
72 return int(value)
73 except (TypeError, ValueError) as exc:
74 raise ModelSpecError(f"{name} must be an integer, got {value!r}") from exc
77def _as_count(value: Any, name: str) -> int:
78 """Coerce *value* to a non-negative whole number, or raise naming it."""
79 if isinstance(value, bool) or (isinstance(value, float) and not value.is_integer()):
80 raise ModelSpecError(f"{name} must be a whole number, got {value!r}")
81 count = _as_int(value, name)
82 if count is None or count < 0:
83 raise ModelSpecError(f"{name} must be a non-negative count, got {value!r}")
84 return count
87@dataclass(frozen=True)
88class OpCounts:
89 """How many times one layer kind runs each op the cost model prices.
91 The estimators already price a layer as a count times a per-op cost
92 (``perf_estimation/getters.py``); this is the count half, declared as data
93 instead of assigned by a per-family callback. It belongs to a layer
94 *kind*, not to a model: an encoder-decoder, or a stack mixing linear and
95 full attention, runs different ops in different layers.
97 Every op is required. A count left out would be priced as zero, which is
98 exactly the silent failure the model IR exists to refuse, so a layer that
99 does not run an op declares ``0``. The parameter-cast counts are absent on
100 purpose: whether parameters are cast depends on optimizer sharding, which
101 is strategy, so the consumer derives them from ``attMM`` and ``ffMM``.
102 """
104 # The fields are the op names of ND's op vector, spelled as the op
105 # profiles' yaml, the recompute switches and every census record spell
106 # them (attMM, normOp, ...): renaming one renames those keys.
107 # pylint: disable=invalid-name
108 attMM: int # attention projections: q, k, v, o
109 attBMM: int # attention batched matmuls: QK^T and AV
110 ffMM: int # feed-forward projections, e.g. 3 for a gated MLP
111 softmax: int
112 dropout: int
113 normOp: int
114 gather: int # tensor-parallel gathers
115 headCast: int # casts of the per-head score tensor
116 ffAct: int # feed-forward activation functions
117 linrec: int # linear-attention state updates: the delta rule's recurrence
118 # pylint: enable=invalid-name
120 @classmethod
121 def from_dict(cls, data: Mapping[str, Any], where: str = "ops") -> "OpCounts":
122 """Build the vector from a mapping, refusing a missing or unknown op.
124 Args:
125 data: ``{op name: count}``, one entry per op.
126 where: Where the mapping came from, for the error message.
128 Raises:
129 ModelSpecError: If an op is missing, unknown, or not a count.
130 """
131 if not isinstance(data, Mapping):
132 raise ModelSpecError(f"{where} must map op names to counts, got {data!r}")
133 names = [f.name for f in fields(cls)]
134 unknown = sorted(set(data) - set(names))
135 if unknown:
136 raise ModelSpecError(f"{where} declares unknown ops {unknown}; the ops are {names}")
137 missing = [name for name in names if name not in data]
138 if missing:
139 raise ModelSpecError(
140 f"{where} is missing counts for {missing}; declare 0 for an op "
141 "the layer does not run"
142 )
143 return cls(**{name: _as_count(data[name], f"{where}.{name}") for name in names})
145 def to_dict(self) -> Dict[str, int]:
146 """Return the vector as ``{op name: count}``, in declaration order."""
147 return {f.name: getattr(self, f.name) for f in fields(self)}
149 def __deepcopy__(self, memo: Dict[int, Any]) -> "OpCounts":
150 """Frozen and made of immutable fields: a deep copy is the record itself."""
151 del memo
152 return self
155def ops_from_dict(data: Any) -> Dict[str, OpCounts]:
156 """Parse ``{layer kind: {op: count}}``, the serialised form of ``ModelSpec.ops``."""
157 if not isinstance(data, Mapping) or not data:
158 raise ModelSpecError(
159 f"ops must map at least one layer kind to its op counts, got {data!r}"
160 )
161 return {str(kind): OpCounts.from_dict(counts, f"ops.{kind}") for kind, counts in data.items()}
164@dataclass(frozen=True)
165class LayerGroup:
166 """Consecutive layers of one kind, in model order.
168 A model's stack is a tuple of groups: DeepSeek-V3 is three dense layers,
169 58 MoE layers and one MTP layer; Qwen3.5 alternates groups of linear- and
170 full-attention layers. ``kind`` names a kind of the model's op profile.
171 MTP layers are body-shaped layers after the body, marked so that nothing
172 has to count them from the end.
173 """
175 kind: str
176 count: int
177 mtp: bool = False
179 @classmethod
180 def from_dict(cls, data: Mapping[str, Any], where: str = "layers") -> "LayerGroup":
181 """Build a group from ``{kind, count, mtp}``, refusing anything else.
183 Raises:
184 ModelSpecError: If the kind is missing, the count is not a
185 positive whole number, or an unknown key is present.
186 """
187 if not isinstance(data, Mapping):
188 raise ModelSpecError(f"{where} must be a mapping of kind and count, got {data!r}")
189 unknown = sorted(set(data) - {"kind", "count", "mtp"})
190 if unknown:
191 raise ModelSpecError(f"{where} declares unknown keys {unknown}; the keys are kind, count, mtp")
192 kind = data.get("kind")
193 if not isinstance(kind, str) or not kind:
194 raise ModelSpecError(f"{where}.kind must name a layer kind, got {kind!r}")
195 count = _as_count(data.get("count"), f"{where}.count")
196 if count == 0:
197 raise ModelSpecError(f"{where}.count must be positive; leave out a group with no layers")
198 mtp = data.get("mtp", False)
199 if not isinstance(mtp, bool):
200 raise ModelSpecError(f"{where}.mtp must be true or false, got {mtp!r}")
201 return cls(kind=kind, count=count, mtp=mtp)
203 def to_dict(self) -> Dict[str, Any]:
204 """Return the group as a mapping, omitting ``mtp`` when it is false."""
205 out: Dict[str, Any] = {"kind": self.kind, "count": self.count}
206 if self.mtp:
207 out["mtp"] = True
208 return out
211def layers_from_list(data: Any, where: str = "layers") -> Tuple[LayerGroup, ...]:
212 """Parse a list of groups, the serialised form of ``ModelSpec.layers``."""
213 if not isinstance(data, (list, tuple)) or not data:
214 raise ModelSpecError(f"{where} must list at least one group, got {data!r}")
215 return tuple(LayerGroup.from_dict(group, f"{where}[{index}]") for index, group in enumerate(data))
218def check_layer_counts(layers: Tuple[LayerGroup, ...], num_layers: int, mtp_depth: int,
219 where: str = "layers") -> None:
220 """Raise unless *layers* covers exactly the body and the MTP layers, in order.
222 Args:
223 layers: The groups, in model order.
224 num_layers: The number of body layers the model declares.
225 mtp_depth: The number of MTP layers the model declares.
226 where: The field the groups were read from, for the message.
228 Raises:
229 ModelSpecError: If a body group follows an MTP group, or either
230 part does not sum to its declared count.
231 """
232 seen_mtp = False
233 for index, group in enumerate(layers):
234 if seen_mtp and not group.mtp:
235 raise ModelSpecError(f"{where}[{index}] is a body group after an MTP group; MTP layers come last")
236 seen_mtp = seen_mtp or group.mtp
237 body = sum(group.count for group in layers if not group.mtp)
238 mtp = sum(group.count for group in layers if group.mtp)
239 if body != num_layers:
240 raise ModelSpecError(f"{where} list {body} body layers, but num_hidden_layers is {num_layers}")
241 if mtp != mtp_depth:
242 raise ModelSpecError(f"{where} list {mtp} MTP layers, but mtp_depth is {mtp_depth}")
245# The fields a census record states in pairs: what a layer keeps under
246# HyperParallel's selective activation checkpointing, the shares of its
247# matmul FLOPs that recomputes, and what it keeps for each op.
248_PAIRED = (("selective", "selective_tp"), ("selective_attention_mm", "selective_ffn_mm"), ("ops", "ops_tp"))
250# The census record's fields that map ops to bytes per token.
251_BY_OP = ("ops", "ops_tp")
254@dataclass(frozen=True)
255class KindActivations:
256 """What one layer of a kind keeps for its backward, and the most its backward holds, per token.
258 Bytes per token at micro-batch 1, as a census measures them on one layer
259 (:mod:`hyper_parallel.auto_parallel._layer_census`) at ``seq_length``
260 tokens under the runtime's kernels: the part no tensor-parallel rank
261 splits, and the part it splits, of which a layer at TP 2 holds half. The
262 backward's working set leaves out the parameters and their gradients,
263 which the memory model counts on its own. ``selective`` and
264 ``selective_tp`` are what the layer keeps under HyperParallel's selective
265 activation checkpointing, whose backward recomputes the rest and holds
266 the same working set, and ``selective_attention_mm`` and
267 ``selective_ffn_mm`` the shares of the FLOPs of its attention's
268 projections and of the rest of its matmuls that recomputes. ``ops`` and
269 ``ops_tp`` are what it keeps for each op of the op vector, the two parts
270 of ``saved`` and ``saved_tp``, and ``other`` for its own code: the op
271 records of shared decision S1, as the census fills them. A record
272 states each pair whole or not at all.
273 """
275 saved: float
276 saved_tp: float
277 working: float
278 working_tp: float
279 seq_length: int
280 selective: Optional[float] = None
281 selective_tp: Optional[float] = None
282 selective_attention_mm: Optional[float] = None
283 selective_ffn_mm: Optional[float] = None
284 ops: Optional[Mapping[str, float]] = None
285 ops_tp: Optional[Mapping[str, float]] = None
287 def to_dict(self) -> Dict[str, Any]:
288 """Return the record as a plain mapping, the selective part and the ops only when stated."""
289 out = {spec_field.name: getattr(self, spec_field.name) for spec_field in fields(self)
290 if getattr(self, spec_field.name) is not None}
291 for name in _BY_OP:
292 if name in out:
293 out[name] = dict(out[name])
294 return out
296 @staticmethod
297 def _check_keys(data: Mapping[str, Any], names: Tuple[str, ...], where: str) -> None:
298 """Raise unless *data* states the fields of *names*, each of :data:`_PAIRED` whole or not at all."""
299 unknown = sorted(set(data) - set(names))
300 missing = [name for name in names if name not in data and all(name not in pair for pair in _PAIRED)]
301 if unknown or missing:
302 raise ModelSpecError(f"{where} has unknown keys {unknown} and lacks {missing}")
303 for pair in _PAIRED:
304 if len({data.get(name) is None for name in pair}) > 1:
305 raise ModelSpecError(f"{where} states one of {list(pair)} without the other")
307 @classmethod
308 def from_dict(cls, data: Mapping[str, Any], where: str = "activations") -> "KindActivations":
309 """Build a record, refusing a key it does not know, a missing one, a negative size or a share above 1."""
310 if not isinstance(data, Mapping):
311 raise ModelSpecError(f"{where} must map the record's fields to their values, got {data!r}")
312 names = tuple(spec_field.name for spec_field in fields(cls))
313 cls._check_keys(data, names, where)
314 sizes = {name: float(data[name]) for name in names
315 if name != "seq_length" and name not in _BY_OP and data.get(name) is not None}
316 if any(size < 0 for size in sizes.values()):
317 raise ModelSpecError(f"{where}: bytes per token cannot be negative, got {sizes}")
318 if any(sizes.get(name, 0) > 1 for name in _PAIRED[1]):
319 raise ModelSpecError(f"{where}: a share of FLOPs cannot exceed 1, got {sizes}")
320 by_op = {name: _op_bytes(data[name], f"{where}.{name}") for name in _BY_OP if data.get(name) is not None}
321 return cls(seq_length=_as_count(data["seq_length"], f"{where}.seq_length"), **sizes, **by_op)
324def _op_bytes(data: Any, where: str) -> Dict[str, float]:
325 """Parse ``{op: bytes per token}``, the ops of the op vector and ``other``, refusing a negative size."""
326 known = [spec_field.name for spec_field in fields(OpCounts)] + ["other"]
327 if not isinstance(data, Mapping) or any(op not in known for op in data):
328 raise ModelSpecError(f"{where} must map ops of {known} to bytes per token, got {data!r}")
329 sizes = {str(op): float(size) for op, size in data.items()}
330 if any(size < 0 for size in sizes.values()):
331 raise ModelSpecError(f"{where}: bytes per token cannot be negative, got {sizes}")
332 return sizes
335def activations_from_dict(data: Any) -> Dict[str, KindActivations]:
336 """Parse a spec's ``activations``, a mapping of layer kind to its :class:`KindActivations`."""
337 if not isinstance(data, Mapping):
338 raise ModelSpecError(f"activations must map layer kinds to their records, got {data!r}")
339 return {str(kind): KindActivations.from_dict(record, f"activations.{kind}") for kind, record in data.items()}
342@dataclass(frozen=True)
343class VisionSpec:
344 """The vision tower of a multimodal model.
346 Towers are nested rather than flattened with a prefix, so a spec can carry
347 more than one and a consumer can ask whether there is one at all.
348 ``max_position_embeddings`` here is the encoder sequence length in merged
349 visual tokens, which depends on the images the dataset serves and is
350 therefore an input, not a property of the checkpoint. ``layers`` is the
351 tower's stack, groups of kinds of the vision profile; a producer states
352 one encoder group when the config states none.
353 """
355 hidden_size: int
356 num_hidden_layers: int
357 num_attention_heads: int
358 name: str = "vision"
359 intermediate_size: Optional[int] = None
360 out_hidden_size: Optional[int] = None
361 patch_size: Optional[int] = None
362 spatial_merge_size: Optional[int] = None
363 num_position_embeddings: Optional[int] = None
364 max_position_embeddings: Optional[int] = None
365 layers: Optional[Tuple[LayerGroup, ...]] = None
367 def validate(self) -> None:
368 """Raise :class:`ModelSpecError` if the tower cannot be costed."""
369 for name in ("hidden_size", "num_hidden_layers", "num_attention_heads"):
370 value = getattr(self, name)
371 if not value or value <= 0:
372 raise ModelSpecError(
373 f"vision.{name} is required and must be positive, got {value!r}"
374 )
375 if self.layers is not None:
376 if any(group.mtp for group in self.layers):
377 raise ModelSpecError("vision.layers marks MTP layers; a tower has none")
378 check_layer_counts(self.layers, self.num_hidden_layers, 0, "vision.layers")
380 def to_dict(self) -> Dict[str, Any]:
381 """Return the tower as a plain mapping, unset fields omitted."""
382 out: Dict[str, Any] = {}
383 for spec_field in fields(self):
384 value = getattr(self, spec_field.name)
385 if value is None:
386 continue
387 out[spec_field.name] = [group.to_dict() for group in value] if spec_field.name == "layers" else value
388 return out
391@dataclass(frozen=True)
392class ModelSpec:
393 """What is being trained, independent of how it is parallelised.
395 The four required fields are the ones every cost path dereferences
396 unconditionally; the rest are optional because a dense model genuinely has
397 no expert count, and ``None`` says so without claiming the count is zero.
399 Attributes:
400 arch: The op profile the model is priced with, one of the files in
401 ``auto_parallel/op_profiles``. A producer that knows the family
402 declares it; otherwise it is inferred from ``name`` when the spec
403 is resolved.
404 ops: Per-layer-kind op counts that replace the profile's own, for a
405 model no profile describes. Its kinds must be the profile's.
406 layers: The layer stack: groups of consecutive layers of one kind of
407 the profile, body first and MTP layers last. A producer derives
408 it from ``layer_types``, ``first_k_dense_replace`` and
409 ``mtp_depth`` when the config does not state it.
410 layer_types: The attention flavour of each layer, as a hybrid config
411 such as Qwen3.5 lists them; a producer states ``layers`` from
412 them and drops them.
413 tie_word_embeddings: Whether the output head shares the embedding's
414 table, as a Transformers config states it.
415 qk_norm: Whether attention normalizes each head's queries and keys
416 over the head width before the scores, as Qwen3's ``q_norm`` and
417 ``k_norm`` do. A Transformers config does not state it; its
418 producer infers it from the model's name.
419 qkv_bias, o_bias, mlp_bias: Whether the query, key and value
420 projections, the output projection and the feed-forward's hold a
421 bias. Unstated, the parameter formulas count their own.
422 norm_bias, layer_norms: Whether each norm holds a bias beside its
423 weight, as a LayerNorm does, and how many norms a layer holds,
424 its attention's per-head query and key norms apart.
425 shared_expert_gate: Whether a MoE layer gates its shared expert's
426 output with a weight of its own, as Qwen2-MoE's does.
427 torch_dtype: The dtype the checkpoint's weights are saved in, by name,
428 which the runtime loads them in under its default ``auto``; how a
429 run holds and computes them is its execution spec's.
430 activations: What a layer of each kind keeps for its backward and
431 holds in it, per token, as a census measures them; the memory
432 model prices a layer of a kind stated here with them rather
433 than its formulas.
434 output_activations: The same record for the output layer, the final
435 norm, the output projection and the loss; its TP part is the
436 part a vocabulary-parallel loss splits.
437 """
439 name: str
440 hidden_size: int
441 num_hidden_layers: int
442 num_attention_heads: int
443 vocab_size: int
445 num_key_value_heads: Optional[int] = None
446 head_dim: Optional[int] = None
447 intermediate_size: Optional[int] = None
448 max_position_embeddings: Optional[int] = None
450 num_experts: Optional[int] = None
451 num_experts_per_tok: Optional[int] = None
452 num_shared_experts: Optional[int] = None
453 moe_intermediate_size: Optional[int] = None
454 shared_expert_intermediate_size: Optional[int] = None
455 first_k_dense_replace: Optional[int] = None
457 mtp_depth: Optional[int] = None
459 multiple_of: Optional[int] = None
460 ffn_dim_multiplier: Optional[float] = None
461 kv_lora_rank: Optional[int] = None
462 q_lora_rank: Optional[int] = None
463 qk_rope_head_dim: Optional[int] = None
464 qk_nope_head_dim: Optional[int] = None
465 v_head_dim: Optional[int] = None
466 # Qwen3.5 fuses an output gate into the query projection, doubling its width.
467 attn_output_gate: Optional[bool] = None
468 tie_word_embeddings: Optional[bool] = None
469 qk_norm: Optional[bool] = None
470 qkv_bias: Optional[bool] = None
471 o_bias: Optional[bool] = None
472 mlp_bias: Optional[bool] = None
473 norm_bias: Optional[bool] = None
474 layer_norms: Optional[int] = None
475 shared_expert_gate: Optional[bool] = None
476 torch_dtype: Optional[str] = None
478 # Gated DeltaNet linear attention (Qwen3.5), for the layers a linear kind prices.
479 linear_num_key_heads: Optional[int] = None
480 linear_key_head_dim: Optional[int] = None
481 linear_num_value_heads: Optional[int] = None
482 linear_value_head_dim: Optional[int] = None
483 linear_conv_kernel_dim: Optional[int] = None
485 arch: Optional[str] = None
486 ops: Optional[Dict[str, OpCounts]] = None
487 layers: Optional[Tuple[LayerGroup, ...]] = None
488 layer_types: Optional[Tuple[str, ...]] = None
489 activations: Optional[Dict[str, KindActivations]] = None
490 output_activations: Optional[KindActivations] = None
492 vision: Optional[VisionSpec] = None
494 _REQUIRED = ("hidden_size", "num_hidden_layers", "num_attention_heads", "vocab_size")
496 # ---- derived reads -------------------------------------------------
498 @property
499 def effective_head_dim(self) -> int:
500 """Per-head width, explicit when declared and derived when not.
502 Costing attention as ``hidden_size x hidden_size`` is only correct
503 when ``head_dim == hidden_size / num_attention_heads``. Qwen3.5
504 declares ``head_dim: 256`` against ``h=2048, a=16``, so the two
505 differ by 2x and half the projection parameters go missing.
506 """
507 if self.head_dim:
508 return int(self.head_dim)
509 return self.hidden_size // self.num_attention_heads
511 @property
512 def attention_width(self) -> int:
513 """Total width of the query projection, ``a * head_dim``."""
514 return self.num_attention_heads * self.effective_head_dim
516 @property
517 def is_moe(self) -> bool:
518 """Whether the model routes tokens to more than one expert."""
519 return bool(self.num_experts and self.num_experts > 1)
521 @property
522 def routed_expert_width(self) -> int:
523 """Feed-forward width of one routed expert, ``0`` when none resolves.
525 Mixtral-style configs run their experts at the dense feed-forward
526 width and declare no ``moe_intermediate_size``, so the dense width is
527 a real fallback rather than a guess.
528 """
529 width = self.moe_intermediate_size or self.intermediate_size
530 return int(width or 0)
532 @property
533 def shared_expert_width(self) -> int:
534 """Feed-forward width of one shared expert, ``0`` when there is none.
536 A config may declare the shared expert's own width or inherit the
537 routed one. An all-MoE config declares no dense ``intermediate_size``
538 at all, which is the case that used to price this at zero.
539 """
540 if not self.num_shared_experts:
541 return 0
542 width = self.shared_expert_intermediate_size or self.routed_expert_width
543 return int(width or 0)
545 # ---- validation ----------------------------------------------------
547 def validate(self) -> "ModelSpec":
548 """Raise :class:`ModelSpecError` unless the spec is internally coherent.
550 Returns *self*, so a producer can ``return spec.validate()``.
551 """
552 for name in self._REQUIRED:
553 value = getattr(self, name)
554 if not value or value <= 0:
555 raise ModelSpecError(
556 f"{name} is required and must be positive, got {value!r}. "
557 "Supply it in the model config or in model.config_overrides."
558 )
560 kv_heads = self.num_key_value_heads
561 if kv_heads and self.num_attention_heads % kv_heads:
562 raise ModelSpecError(
563 f"num_attention_heads ({self.num_attention_heads}) must be a multiple "
564 f"of num_key_value_heads ({kv_heads})"
565 )
567 if self.is_moe and not self.routed_expert_width:
568 raise ModelSpecError(
569 f"no feed-forward width resolves for {self.num_experts} routed "
570 "experts; declare moe_intermediate_size, or intermediate_size if "
571 "the experts run at the dense width"
572 )
573 if self.num_shared_experts and not self.shared_expert_width:
574 raise ModelSpecError(
575 "num_shared_experts is set but no width resolves for it; declare "
576 "shared_expert_intermediate_size or moe_intermediate_size"
577 )
578 if self.layers is not None:
579 check_layer_counts(self.layers, self.num_hidden_layers, self.mtp_depth or 0)
580 self._check_activations()
581 if self.vision is not None:
582 self.vision.validate()
583 return self
585 def _check_activations(self) -> None:
586 """Raise unless every kind the activations state is a kind of the stack."""
587 stated = set(self.activations or {}) - {group.kind for group in self.layers or ()}
588 if stated:
589 raise ModelSpecError(f"activations of kinds {sorted(stated)} no layer of the stack is")
591 # ---- serialised form -----------------------------------------------
593 def to_dict(self) -> Dict[str, Any]:
594 """Return the spec as a plain mapping, ready for YAML.
596 Unset optional fields are omitted rather than written as ``null``, so
597 a hand-written file and a dumped one look alike and a reader cannot
598 mistake "not declared" for "declared empty".
599 """
600 out: Dict[str, Any] = {}
601 for spec_field in fields(self):
602 if spec_field.name in ("vision", "ops", "layers", "layer_types", "activations", "output_activations"):
603 continue
604 value = getattr(self, spec_field.name)
605 if value is not None:
606 out[spec_field.name] = value
607 if self.ops is not None:
608 out["ops"] = {kind: counts.to_dict() for kind, counts in self.ops.items()}
609 if self.layers is not None:
610 out["layers"] = [group.to_dict() for group in self.layers]
611 if self.layer_types is not None:
612 out["layer_types"] = list(self.layer_types)
613 if self.activations is not None:
614 out["activations"] = {kind: record.to_dict() for kind, record in self.activations.items()}
615 if self.output_activations is not None:
616 out["output_activations"] = self.output_activations.to_dict()
617 if self.vision is not None:
618 out["vision"] = self.vision.to_dict()
619 return out
621 @classmethod
622 def from_dict(cls, data: Mapping[str, Any], strict: bool = True) -> "ModelSpec":
623 """Build a spec from a mapping, refusing a key it does not know.
625 Args:
626 data: A parsed YAML mapping, or whatever a producer assembled.
627 :func:`model_fields` takes the model's keys from a mapping
628 that also states the run.
629 strict: Validate before returning. Only a test constructing a
630 deliberately partial spec should pass ``False``.
632 Raises:
633 ModelSpecError: If a key is unknown, a required field is missing,
634 a field is not an integer, or the declared fields contradict
635 each other.
636 """
637 unknown = sorted(set(data) - MODEL_KEYS)
638 if unknown:
639 precision = [key for key in unknown if key in _RUN_PRECISION_KEYS]
640 raise ModelSpecError(
641 f"unknown model spec keys {unknown}: a model spec states the model only, "
642 "and a run's strategy, such as its offset or recompute, belongs in its "
643 "execution spec" + (
644 f"; {precision} state the run's precision, which a HyperParallel run takes from "
645 "fsdp_config.mix_precision.param_dtype, else model_init_dtype, else model.torch_dtype"
646 if precision else ""
647 )
648 )
649 kwargs: Dict[str, Any] = {}
650 for key, value in data.items():
651 if key != "vision":
652 kwargs[key] = cls._field_value(key, value)
654 for name in cls._REQUIRED:
655 if name not in kwargs:
656 raise ModelSpecError(
657 f"{name} is required and was not supplied. "
658 "Supply it in the model config or in model.config_overrides."
659 )
660 kwargs.setdefault("name", "custom")
662 raw_vision = data.get("vision")
663 if isinstance(raw_vision, Mapping):
664 kwargs["vision"] = cls._vision_from_dict(raw_vision)
666 spec = cls(**kwargs)
667 return spec.validate() if strict else spec
669 @staticmethod
670 def _field_value(key: str, value: Any) -> Any:
671 """Coerce one declared field to the type the schema gives it."""
672 if key == "name":
673 return str(value)
674 if value is None:
675 return None
676 if key == "arch":
677 return str(value)
678 if key == "torch_dtype":
679 return str(value).replace("torch.", "")
680 if key == "ops":
681 return ops_from_dict(value)
682 if key == "layers":
683 return layers_from_list(value)
684 if key == "activations":
685 return activations_from_dict(value)
686 if key == "output_activations":
687 return KindActivations.from_dict(value, key)
688 if key == "ffn_dim_multiplier":
689 return float(value)
690 if key in _FLAGS:
691 if not isinstance(value, bool):
692 raise ModelSpecError(f"{key} must be true or false, got {value!r}")
693 return value
694 if key == "layer_types":
695 if not isinstance(value, (list, tuple)) or not all(isinstance(kind, str) for kind in value):
696 raise ModelSpecError(f"layer_types must list layer type names, got {value!r}")
697 return tuple(value)
698 return _as_int(value, key)
700 @staticmethod
701 def _vision_from_dict(data: Mapping[str, Any]) -> VisionSpec:
702 """Build a :class:`VisionSpec`, dropping keys it does not declare."""
703 known = {f.name for f in fields(VisionSpec)}
704 kwargs: Dict[str, Any] = {}
705 for key, value in data.items():
706 if key not in known:
707 continue
708 if key == "layers":
709 kwargs[key] = None if value is None else layers_from_list(value, "vision.layers")
710 else:
711 kwargs[key] = str(value) if key == "name" else _as_int(value, f"vision.{key}")
712 missing = [
713 name for name in ("hidden_size", "num_hidden_layers", "num_attention_heads")
714 if name not in kwargs
715 ]
716 if missing:
717 raise ModelSpecError(
718 f"vision tower is missing required fields: {', '.join(missing)}"
719 )
720 return VisionSpec(**kwargs)
723# The keys a model spec reads.
724MODEL_KEYS = frozenset(spec_field.name for spec_field in fields(ModelSpec))
727def model_fields(data: Mapping[str, Any]) -> Dict[str, Any]:
728 """Return the entries of *data* a model spec reads.
730 For a mapping that states the model and the run together, such as a
731 search config's model section.
732 """
733 return {key: value for key, value in data.items() if key in MODEL_KEYS}