Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / _op_profiles.py: 95%

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"""Op profiles: the layer kinds of each known architecture family, as data. 

16 

17Every file under ``op_profiles/`` names a family and declares its layer kinds. 

18A kind states how many times one layer runs each op the cost model prices, 

19the numbers the ND arch hooks used to assign in code, and two flavours from a 

20closed set that say how the layer is shaped: 

21 

22- ``attention``: ``full`` (the default) or ``linear``, a gated-DeltaNet layer 

23 priced on the model's linear-attention dimensions; 

24- ``ffn``: ``dense`` or ``moe``, or absent to keep the model's own 

25 feed-forward, for families whose layers do not differ in it. 

26 

27A profile with more than one kind may name a ``default``, the kind a layer 

28takes when nothing says otherwise; a profile with one kind defaults to it. 

29 

30A profile may also state what a model of its family has when its producer 

31does not say: under ``run``, the run's byte widths, whether its gradients 

32take memory without pipeline parallelism and whether tensor parallelism 

33shards the activations between layers (:data:`DEFAULT_RUN` otherwise); under 

34``model``, an MLA family's value-head width. The cost model's ``derive`` 

35reads them. 

36 

37A model spec names its profile with ``arch``, and may declare its own counts 

38in ``ops`` instead. The cost model reads only what the spec declares. A 

39producer that has nothing better than a free-text model name, such as a 

40MindFormers ``trainer.model_name``, falls back on :func:`infer_arch`. 

41""" 

42import functools 

43import logging 

44import os 

45from dataclasses import dataclass, field 

46from types import MappingProxyType 

47from typing import Any, Dict, Mapping, Optional, Tuple 

48 

49import yaml 

50 

51from hyper_parallel.auto_parallel._exec_spec import ExecSpec, ExecSpecError 

52from hyper_parallel.auto_parallel._model_spec import ModelSpecError, OpCounts 

53 

54logger = logging.getLogger(__name__) 

55 

56PROFILE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "op_profiles") 

57DEFAULT_ARCH = "default" 

58# The profile every vision tower is priced with. 

59VISION_ARCH = "vision" 

60 

61ATTENTION_FLAVOURS = ("full", "linear") 

62FFN_FLAVOURS = ("dense", "moe") 

63 

64# The run of a family's model when its producer states none of it, and the 

65# only execution-spec fields a profile's ``run`` may change: the byte widths 

66# of gradients, optimizer states, norm activations and dropout masks, whether 

67# gradients take memory without pipeline parallelism, whether tensor 

68# parallelism shards the activations between layers, and whether the loss 

69# runs on logits sharded over the vocabulary. 

70DEFAULT_RUN = MappingProxyType({ 

71 "grad_bytes": 4, 

72 "optimizer_state_bytes": 4, 

73 "optimizer_states": 2, 

74 "main_param_bytes": 0, 

75 "norm_bytes": 4, 

76 "dropout_bytes": 0, 

77 "grad_accumulation": False, 

78 "shard_activations": False, 

79 "loss_parallel": True, 

80 "reshard_params": False, 

81 "deferred_grad_accumulation": False, 

82 "overlapped_grad_reduce": False, 

83}) 

84 

85# The model facts a profile's ``model`` may give a model that states none. 

86MODEL_DEFAULT_KEYS = ("v_head_dim",) 

87 

88# What a model name opens with, with the family it means. The longest that 

89# a name opens with wins, so Qwen3.5, under both the Transformers spelling 

90# and the release name, is not taken for the Qwen family it opens with too, 

91# and the order they are listed in decides nothing. 

92_NAME_PREFIXES = { 

93 "llama2": "llama2", "mixtral": "mixtral", "t5": "t5", "pangualpha": "pangualpha", 

94 "deepseek": "deepseek", "qwen3_5": "qwen3_5", "qwen3.5": "qwen3_5", "qwen": "qwen", 

95} 

96# The names the default profile was taken for, each said once a process. 

97_DEFAULTED_NAMES = set() 

98 

99 

100@dataclass(frozen=True) 

101class LayerKind: 

102 """One kind of layer: the ops it runs and how it is shaped. 

103 

104 Attributes: 

105 name: The kind's name in its profile, which a layer stack refers to. 

106 ops: How many times one layer of this kind runs each op. 

107 attention: ``full``, or ``linear`` for a gated-DeltaNet layer. 

108 ffn: ``dense``, ``moe``, or ``None`` to keep the model's own. 

109 """ 

110 

111 name: str 

112 ops: OpCounts 

113 attention: str = "full" 

114 ffn: Optional[str] = None 

115 

116 def __deepcopy__(self, memo: Dict[int, Any]) -> "LayerKind": 

117 """Frozen and made of immutable fields: a deep copy is the record itself.""" 

118 del memo 

119 return self 

120 

121 

122@dataclass(frozen=True) 

123class OpProfile: 

124 """The layer kinds of one architecture family. 

125 

126 Attributes: 

127 arch: The family, the profile's file name. 

128 layer_kinds: Every kind of the family, by name. 

129 default: The kind a layer takes when nothing says otherwise, or 

130 ``None`` for a family whose stack always states its kinds. 

131 run: The run of the family's models where their producer states 

132 none, every key of :data:`DEFAULT_RUN`. 

133 model: The model facts the family gives a model that states none. 

134 """ 

135 

136 arch: str 

137 layer_kinds: Dict[str, LayerKind] 

138 default: Optional[str] = None 

139 run: Mapping[str, Any] = field(default_factory=lambda: DEFAULT_RUN) 

140 model: Mapping[str, int] = field(default_factory=lambda: MappingProxyType({})) 

141 

142 @property 

143 def kinds(self) -> Dict[str, OpCounts]: 

144 """Return each kind's op counts, by kind name.""" 

145 return {name: kind.ops for name, kind in self.layer_kinds.items()} 

146 

147 def counts(self, kind: str) -> OpCounts: 

148 """Return one layer kind's counts, or raise naming the kinds there are.""" 

149 return self.kind(kind).ops 

150 

151 def kind(self, name: str) -> LayerKind: 

152 """Return one layer kind, or raise naming the kinds there are.""" 

153 if name not in self.layer_kinds: 

154 raise ModelSpecError( 

155 f"op profile {self.arch!r} has no layer kind {name!r}; " 

156 f"it has {sorted(self.layer_kinds)}" 

157 ) 

158 return self.layer_kinds[name] 

159 

160 

161@functools.lru_cache(maxsize=None) 

162def known_archs() -> Tuple[str, ...]: 

163 """Return the name of every family that has a profile file.""" 

164 return tuple(sorted( 

165 name[:-len(".yaml")] for name in os.listdir(PROFILE_DIR) if name.endswith(".yaml") 

166 )) 

167 

168 

169def _flavour(value: Any, allowed: Tuple[str, ...], where: str, optional: bool) -> Optional[str]: 

170 """Return a flavour value, refusing one outside *allowed*.""" 

171 if value is None and optional: 

172 return None 

173 if value not in allowed: 

174 raise ModelSpecError(f"{where} must be one of {list(allowed)}, got {value!r}") 

175 return str(value) 

176 

177 

178def _layer_kind(arch: str, name: str, data: Any) -> LayerKind: 

179 """Build one kind of a profile from its ``{ops, attention, ffn}`` mapping.""" 

180 where = f"{arch}.kinds.{name}" 

181 if not isinstance(data, Mapping): 

182 raise ModelSpecError(f"{where} must map ops and flavours, got {data!r}") 

183 unknown = sorted(set(data) - {"ops", "attention", "ffn"}) 

184 if unknown: 

185 raise ModelSpecError(f"{where} has unknown keys {unknown}; a kind has ops, attention and ffn") 

186 return LayerKind( 

187 name=str(name), 

188 ops=OpCounts.from_dict(data.get("ops"), f"{where}.ops"), 

189 attention=_flavour(data.get("attention", "full"), ATTENTION_FLAVOURS, f"{where}.attention", False), 

190 ffn=_flavour(data.get("ffn"), FFN_FLAVOURS, f"{where}.ffn", True), 

191 ) 

192 

193 

194def _run_defaults(arch: str, data: Any) -> Mapping[str, Any]: 

195 """Return a family's run: :data:`DEFAULT_RUN`, with what its profile states.""" 

196 stated = {} if data is None else data 

197 if not isinstance(stated, Mapping): 

198 raise ModelSpecError(f"{arch}.run must map run facts to values, got {data!r}") 

199 unknown = sorted(set(stated) - set(DEFAULT_RUN)) 

200 if unknown: 

201 raise ModelSpecError(f"{arch}.run has unknown keys {unknown}; a family states {sorted(DEFAULT_RUN)}") 

202 try: 

203 spec = ExecSpec.from_dict(stated) 

204 except ExecSpecError as exc: 

205 raise ModelSpecError(f"{arch}.run: {exc}") from exc 

206 return MappingProxyType({**DEFAULT_RUN, **spec.to_dict()}) 

207 

208 

209def _model_defaults(arch: str, data: Any) -> Mapping[str, int]: 

210 """Return the model facts a family gives a model that states none.""" 

211 stated = {} if data is None else data 

212 if not isinstance(stated, Mapping): 

213 raise ModelSpecError(f"{arch}.model must map model facts to values, got {data!r}") 

214 unknown = sorted(set(stated) - set(MODEL_DEFAULT_KEYS)) 

215 if unknown: 

216 raise ModelSpecError( 

217 f"{arch}.model has unknown keys {unknown}; a family states {list(MODEL_DEFAULT_KEYS)}" 

218 ) 

219 for key, value in stated.items(): 

220 if isinstance(value, bool) or not isinstance(value, int) or value < 1: 

221 raise ModelSpecError(f"{arch}.model.{key} must be a positive whole number, got {value!r}") 

222 return MappingProxyType(dict(stated)) 

223 

224 

225@functools.lru_cache(maxsize=None) 

226def load_op_profile(arch: str) -> OpProfile: 

227 """Read and validate the profile of one family. 

228 

229 Raises: 

230 ModelSpecError: If no profile has that name, or the file declares no 

231 layer kind, an unknown key, an incomplete op vector, an unknown 

232 flavour, a default that is not one of its kinds, or a run or 

233 model default it cannot state. 

234 """ 

235 if arch not in known_archs(): 

236 raise ModelSpecError( 

237 f"unknown arch {arch!r}; the op profiles are {list(known_archs())}" 

238 ) 

239 with open(os.path.join(PROFILE_DIR, f"{arch}.yaml"), encoding="utf-8") as handle: 

240 data = yaml.safe_load(handle) or {} 

241 unknown = sorted(set(data) - {"kinds", "default", "run", "model"}) 

242 if unknown: 

243 raise ModelSpecError(f"op profile {arch!r} has unknown keys {unknown}") 

244 kinds = data.get("kinds") 

245 if not isinstance(kinds, Mapping) or not kinds: 

246 raise ModelSpecError(f"op profile {arch!r} declares no layer kind") 

247 layer_kinds = {str(name): _layer_kind(arch, name, kind) for name, kind in kinds.items()} 

248 default = data.get("default") 

249 if default is None and len(layer_kinds) == 1: 

250 default = next(iter(layer_kinds)) 

251 if default is not None and default not in layer_kinds: 

252 raise ModelSpecError( 

253 f"op profile {arch!r} names default kind {default!r}, but its kinds are {sorted(layer_kinds)}" 

254 ) 

255 return OpProfile(arch, layer_kinds, default, _run_defaults(arch, data.get("run")), 

256 _model_defaults(arch, data.get("model"))) 

257 

258 

259def family_profile(arch: Optional[str]) -> OpProfile: 

260 """Return the profile of *arch*, or the default family's when no profile has that name.""" 

261 return load_op_profile(arch if arch in known_archs() else DEFAULT_ARCH) 

262 

263 

264def infer_arch(name: Any) -> str: 

265 """Return the family a free-text model name belongs to. 

266 

267 The family the lower-cased name opens with, the longest first, else the 

268 default profile, said once a name. Matched anywhere in the name, in a 

269 fixed order, as the cost model matched names before the family became 

270 data, a family listed earlier won over the one the name opens with. 

271 Only a producer with no better source should use it; a spec that 

272 declares ``arch`` never gets here. 

273 """ 

274 lowered = str(name).lower() 

275 for prefix in sorted(_NAME_PREFIXES, key=len, reverse=True): 

276 if lowered.startswith(prefix): 

277 return _NAME_PREFIXES[prefix] 

278 if lowered not in _DEFAULTED_NAMES: 

279 # Once: a search parses a config more than once. 

280 _DEFAULTED_NAMES.add(lowered) 

281 logger.warning( 

282 "no op profile matches model %r, pricing it as %r; declare arch to choose one", 

283 name, DEFAULT_ARCH, 

284 ) 

285 return DEFAULT_ARCH 

286 

287 

288# What a lower-cased model name contains when its attention normalizes each 

289# head's queries and keys: the Qwen3 generation, Qwen3.5 and the Qwen3 

290# vision-language models included. 

291_QK_NORM_NAMES = ("qwen3",) 

292 

293 

294def infer_qk_norm(name: Any) -> bool: 

295 """Return whether a model named *name* normalizes each head's queries and keys. 

296 

297 For a producer whose config does not state it, as a Transformers config 

298 does not: the Qwen3 generation does, and no other family the cost model 

299 prices does. 

300 """ 

301 lowered = str(name).lower() 

302 return any(pattern in lowered for pattern in _QK_NORM_NAMES) 

303 

304 

305# What a lower-cased model name contains when its attention projects queries, 

306# keys and values with a bias and its output without, which its Transformers 

307# config does not state: the Qwen2 generation. 

308_QKV_BIAS_NAMES = ("qwen2",) 

309 

310# What a lower-cased model name contains when its MoE layer gates its shared 

311# expert's output with a weight of its own: Qwen2-MoE, Qwen3.5 and Qwen3-Next. 

312_SHARED_EXPERT_GATE_NAMES = ("qwen2_moe", "qwen3_5", "qwen3_next") 

313 

314 

315def infer_qkv_bias(name: Any) -> bool: 

316 """Return whether a model named *name* biases its query, key and value projections. 

317 

318 For a Transformers config that states no ``attention_bias``, as the 

319 Qwen2 generation's, whose projections have one. 

320 """ 

321 lowered = str(name).lower() 

322 return any(pattern in lowered for pattern in _QKV_BIAS_NAMES) 

323 

324 

325def infer_shared_expert_gate(name: Any) -> bool: 

326 """Return whether a model named *name* gates its shared expert's output with a weight of its own.""" 

327 lowered = str(name).lower() 

328 return any(pattern in lowered for pattern in _SHARED_EXPERT_GATE_NAMES) 

329 

330 

331def resolve_ops(arch: str, ops: Optional[Mapping[str, OpCounts]] = None) -> Dict[str, OpCounts]: 

332 """Return the op counts a model is priced with, per layer kind. 

333 

334 Declared *ops* replace the profile's, but must name the same layer kinds: 

335 the family's hooks decide which layers are of which kind, so counts for a 

336 kind they never assign would be dropped without a word. 

337 

338 Raises: 

339 ModelSpecError: If *arch* has no profile, or *ops* names other kinds. 

340 """ 

341 profile = load_op_profile(arch) 

342 if ops is None: 

343 return dict(profile.kinds) 

344 if set(ops) != set(profile.kinds): 

345 raise ModelSpecError( 

346 f"ops declares layer kinds {sorted(ops)}, but arch {arch!r} " 

347 f"prices {sorted(profile.kinds)}" 

348 ) 

349 return dict(ops)