Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / _layer_stack.py: 98%
126 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 layer stack: which kind of layer sits where, as data.
17A model's stack is a tuple of :class:`LayerGroup`, consecutive layers of one
18kind of its op profile, body first and MTP layers last. A spec may state it
19in ``layers``. Otherwise :func:`derive_layers` builds it from the fields every
20config already carries, by the rules the cost model's hooks used to encode:
22- ``layer_types`` names each layer's kind, when the profile has those kinds
23 (Qwen3.5 alternates ``linear_attention`` and ``full_attention``);
24- a profile with ``dense`` and ``moe`` kinds runs ``first_k_dense_replace``
25 dense layers and then MoE layers (DeepSeek);
26- a profile with ``encoder`` and ``decoder`` kinds runs one half of each (t5);
27- any other model runs the profile's default kind throughout;
28- the ``mtp_depth`` MTP layers repeat the kind of the last body layer.
30A vision tower's stack is the same, of kinds of the vision profile: one
31encoder group unless the tower states otherwise.
33:func:`resolve_layers` checks a stack against its profile and the model's
34dimensions and returns it as a :class:`LayerStack`, each kind resolved.
35"""
36import logging
37from dataclasses import dataclass
38from typing import Any, Dict, Mapping, Optional, Sequence, Tuple
40from hyper_parallel.auto_parallel._model_spec import (
41 LayerGroup,
42 ModelSpec,
43 ModelSpecError,
44 OpCounts,
45 VisionSpec,
46)
47from hyper_parallel.auto_parallel._op_profiles import (
48 VISION_ARCH,
49 LayerKind,
50 OpProfile,
51 load_op_profile,
52 resolve_ops,
53)
55logger = logging.getLogger(__name__)
57# The spec fields of a gated-DeltaNet layer, and the dimension each one sets.
58_LINEAR_FIELDS = (
59 ("linear_num_key_heads", "num_key_heads"),
60 ("linear_key_head_dim", "key_head_dim"),
61 ("linear_num_value_heads", "num_value_heads"),
62 ("linear_value_head_dim", "value_head_dim"),
63 ("linear_conv_kernel_dim", "conv_kernel_dim"),
64)
67@dataclass(frozen=True)
68class LinearAttentionDims:
69 """The dimensions a linear-attention (gated DeltaNet) layer is priced on."""
71 num_key_heads: int
72 key_head_dim: int
73 num_value_heads: int
74 value_head_dim: int
75 conv_kernel_dim: int
77 @classmethod
78 def from_fields(cls, values: Mapping[str, Any]) -> Optional["LinearAttentionDims"]:
79 """Build the dimensions from ``linear_*`` fields, or ``None`` if none is set.
81 Raises:
82 ModelSpecError: If some are set but not all, or one is not a
83 positive whole number.
84 """
85 found = {name: values.get(name) for name, _ in _LINEAR_FIELDS}
86 if all(value is None for value in found.values()):
87 return None
88 missing = [name for name, value in found.items() if value is None]
89 if missing:
90 raise ModelSpecError(f"linear attention is missing {missing}; declare all of {list(found)}")
91 dims: Dict[str, int] = {}
92 for name, attr in _LINEAR_FIELDS:
93 value = found[name]
94 if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
95 raise ModelSpecError(f"{name} must be a positive whole number, got {value!r}")
96 dims[attr] = value
97 return cls(**dims)
99 def __deepcopy__(self, memo: Dict[int, Any]) -> "LinearAttentionDims":
100 """Frozen and made of immutable fields: a deep copy is the record itself."""
101 del memo
102 return self
105@dataclass(frozen=True)
106class StackGroup:
107 """A group of a resolved stack: its kind, as the profile defines it."""
109 kind: LayerKind
110 count: int
111 mtp: bool = False
113 def __deepcopy__(self, memo: Dict[int, Any]) -> "StackGroup":
114 """Frozen and made of immutable fields: a deep copy is the record itself."""
115 del memo
116 return self
119@dataclass(frozen=True)
120class LayerStack:
121 """A model's layer stack, checked against its profile.
123 Attributes:
124 arch: The profile the kinds come from.
125 groups: The groups in model order, body first and MTP layers last.
126 linear: The linear-attention dimensions, when a kind needs them.
127 """
129 arch: str
130 groups: Tuple[StackGroup, ...]
131 linear: Optional[LinearAttentionDims] = None
133 def kinds(self) -> Tuple[LayerKind, ...]:
134 """Return the kind of every layer, in model order."""
135 return tuple(group.kind for group in self.groups for _ in range(group.count))
137 def distinct_kinds(self) -> Tuple[LayerKind, ...]:
138 """Return each kind the stack uses once, in order of first use."""
139 seen: Dict[str, LayerKind] = {}
140 for group in self.groups:
141 seen.setdefault(group.kind.name, group.kind)
142 return tuple(seen.values())
144 def to_layers(self) -> Tuple[LayerGroup, ...]:
145 """Return the stack in its serialised form, the groups of ``ModelSpec.layers``."""
146 return tuple(LayerGroup(group.kind.name, group.count, group.mtp) for group in self.groups)
148 def __deepcopy__(self, memo: Dict[int, Any]) -> "LayerStack":
149 """Frozen and made of immutable fields: a deep copy is the record itself."""
150 del memo
151 return self
154def _groups(kinds: Sequence[str]) -> Tuple[LayerGroup, ...]:
155 """Run-length group a sequence of kind names."""
156 groups = []
157 for kind in kinds:
158 if groups and groups[-1][0] == kind:
159 groups[-1][1] += 1
160 else:
161 groups.append([kind, 1])
162 return tuple(LayerGroup(kind, count) for kind, count in groups)
165def _body_from_layer_types(profile: OpProfile, num_layers: int,
166 layer_types: Sequence[Any]) -> Optional[Tuple[LayerGroup, ...]]:
167 """The body named by ``layer_types``, or ``None`` if the profile lacks its kinds."""
168 if len(layer_types) < num_layers:
169 raise ModelSpecError(
170 f"layer_types lists {len(layer_types)} layers, but num_hidden_layers is {num_layers}"
171 )
172 kinds = [str(kind) for kind in layer_types[:num_layers]]
173 unknown = sorted(set(kinds) - set(profile.layer_kinds))
174 if unknown:
175 logger.warning(
176 "layer_types names %s, which op profile %r does not declare; every layer is priced as %r",
177 unknown, profile.arch, profile.default,
178 )
179 return None
180 return _groups(kinds)
183def _structural_body(profile: OpProfile, num_layers: int, first_k_dense: int) -> Tuple[LayerGroup, ...]:
184 """The body the profile's kinds imply: dense then MoE, two halves, or one kind."""
185 names = set(profile.layer_kinds)
186 if {"dense", "moe"} <= names:
187 dense = min(max(int(first_k_dense), 0), num_layers)
188 groups = (LayerGroup("dense", dense), LayerGroup("moe", num_layers - dense))
189 return tuple(group for group in groups if group.count)
190 if {"encoder", "decoder"} <= names:
191 if num_layers % 2:
192 raise ModelSpecError(
193 f"an encoder-decoder stack needs an even layer count, got {num_layers}"
194 )
195 half = num_layers // 2
196 return (LayerGroup("encoder", half), LayerGroup("decoder", half)) if half else ()
197 if profile.default is None:
198 raise ModelSpecError(
199 f"op profile {profile.arch!r} has several kinds and names no default; state the stack in layers"
200 )
201 return (LayerGroup(profile.default, num_layers),) if num_layers else ()
204def derive_layers(profile: OpProfile, num_layers: int, mtp_depth: int = 0,
205 layer_types: Optional[Sequence[Any]] = None,
206 first_k_dense: int = 0) -> Tuple[LayerGroup, ...]:
207 """Build the stack a config implies, when it does not state one.
209 Args:
210 profile: The model's op profile, which names the kinds.
211 num_layers: The number of body layers.
212 mtp_depth: The number of MTP layers after the body.
213 layer_types: The per-layer kind names a config such as Qwen3.5 lists.
214 first_k_dense: How many leading layers are dense in a dense-then-MoE
215 family.
217 Returns:
218 The groups in model order, body first and MTP layers last.
220 Raises:
221 ModelSpecError: If ``layer_types`` lists fewer layers than the body,
222 or no rule gives the body a kind.
223 """
224 body = None
225 if layer_types:
226 body = _body_from_layer_types(profile, num_layers, layer_types)
227 if body is None:
228 body = _structural_body(profile, num_layers, first_k_dense)
229 if not mtp_depth:
230 return body
231 last = body[-1].kind if body else profile.default
232 if last is None:
233 raise ModelSpecError(f"op profile {profile.arch!r} gives the MTP layers no kind")
234 return body + (LayerGroup(last, int(mtp_depth), mtp=True),)
237def resolve_kinds(profile: OpProfile, ops: Optional[Mapping[str, OpCounts]] = None) -> Dict[str, LayerKind]:
238 """Return the profile's kinds, with declared *ops* replacing their counts.
240 Raises:
241 ModelSpecError: If *ops* names other kinds than the profile's.
242 """
243 counts = resolve_ops(profile.arch, ops)
244 return {
245 name: LayerKind(name, counts[name], kind.attention, kind.ffn)
246 for name, kind in profile.layer_kinds.items()
247 }
250def resolve_layers(arch: str, layers: Sequence[LayerGroup],
251 ops: Optional[Mapping[str, OpCounts]] = None,
252 linear: Optional[LinearAttentionDims] = None) -> LayerStack:
253 """Check a stack against its profile and resolve each group's kind.
255 Args:
256 arch: The model's op profile.
257 layers: The groups, in model order.
258 ops: Declared op counts that replace the profile's.
259 linear: The model's linear-attention dimensions, if it declares any.
261 Returns:
262 The stack with every kind resolved.
264 Raises:
265 ModelSpecError: If a group names a kind the profile lacks, or a
266 linear kind has no linear-attention dimensions to be priced on.
267 """
268 kinds = resolve_kinds(load_op_profile(arch), ops)
269 groups = []
270 for index, group in enumerate(layers):
271 if group.kind not in kinds:
272 raise ModelSpecError(
273 f"layers[{index}] names kind {group.kind!r}, but op profile {arch!r} "
274 f"has {sorted(kinds)}"
275 )
276 groups.append(StackGroup(kinds[group.kind], group.count, group.mtp))
277 needs_linear = any(group.kind.attention == "linear" for group in groups)
278 if needs_linear and linear is None:
279 raise ModelSpecError(
280 f"op profile {arch!r} prices linear-attention layers, but the model declares no "
281 f"{', '.join(name for name, _ in _LINEAR_FIELDS)}"
282 )
283 return LayerStack(arch, tuple(groups), linear if needs_linear else None)
286def spec_layer_stack(spec: ModelSpec) -> LayerStack:
287 """Return a spec's stack, as it states it or as its fields imply.
289 Raises:
290 ModelSpecError: If the spec has no arch, or its stack does not fit
291 its profile.
292 """
293 if spec.arch is None:
294 raise ModelSpecError(f"{spec.name}: a layer stack needs the spec's arch")
295 linear = LinearAttentionDims.from_fields(
296 {name: getattr(spec, name) for name, _ in _LINEAR_FIELDS}
297 )
298 layers = spec.layers
299 if layers is None:
300 layers = derive_layers(
301 load_op_profile(spec.arch),
302 spec.num_hidden_layers,
303 spec.mtp_depth or 0,
304 layer_types=spec.layer_types,
305 first_k_dense=spec.first_k_dense_replace or 0,
306 )
307 return resolve_layers(spec.arch, layers, spec.ops, linear)
310def tower_layer_stack(vision: VisionSpec) -> LayerStack:
311 """Return a vision tower's stack, as it states it or as one encoder group.
313 Raises:
314 ModelSpecError: If the tower names a kind the vision profile lacks.
315 """
316 layers = vision.layers
317 if layers is None:
318 layers = derive_layers(load_op_profile(VISION_ARCH), vision.num_hidden_layers)
319 return resolve_layers(VISION_ARCH, layers)