Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / _hf_model_spec.py: 99%
256 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"""Shared Transformers-config resolution for the Auto Parallel cost model.
17Both the SAPP-ND parser (``CostModelParserHyperV2``) and the config-adapter
18loader need the same model dimensions out of an AutoModels ``model`` section.
19This module owns that translation exactly once, so the two callers cannot
20drift apart on field names, aliases, or fallback behaviour.
22It is also the single place that imports ``transformers``. The import is
23function-local because ``transformers`` is not a hard dependency of
24``hyper_parallel`` (``requirements.txt`` only pins numpy) and the non-Hyper
25cost-model backends must keep working without it.
26"""
27import copy
28import dataclasses
29import hashlib
30import importlib
31import json
32import logging
33import os
34import sys
35from typing import Any, Dict, Mapping, Optional, Tuple
37from hyper_parallel.auto_parallel._layer_census import (
38 CensusUnavailable,
39 census_activations,
40 census_output_activations,
41)
42from hyper_parallel.auto_parallel._layer_stack import spec_layer_stack, tower_layer_stack
43from hyper_parallel.auto_parallel._model_spec import ModelSpec
44from hyper_parallel.auto_parallel._op_profiles import (
45 infer_arch,
46 infer_qk_norm,
47 infer_qkv_bias,
48 infer_shared_expert_gate,
49 resolve_ops,
50)
51from hyper_parallel.auto_parallel._spec_census import census_model_spec
53logger = logging.getLogger(__name__)
55# AutoModels root sections that identify the current Trainer schema.
56_AUTO_MODELS_SECTIONS = ("training", "accelerator", "fsdp_config")
58# ``model`` keys forwarded verbatim to ``AutoConfig.from_pretrained``.
59_CONFIG_KWARG_NAMES = (
60 "cache_dir", "local_files_only", "revision", "subfolder",
61 "token", "trust_remote_code",
62)
64# Where the model directory is named. A from_config recipe builds the model
65# through a factory (``_target_`` plus ``config_path``) and never loads a
66# checkpoint, so it carries no pretrained path.
67_MODEL_PATH_ALIASES = ("pretrained_model_name_or_path", "config_path")
69# The config a recipe's own factory builds, for the factories whose model
70# ``AutoConfig`` cannot resolve: the path beside ``_target_`` is then a
71# directory Transformers refuses, and the shape the run really trains is the
72# one this function returns. DeepSeek-V4.1 is the case that forced it:
73# Transformers 5.14.1 has no ``deepseek_v41``, and the factory's arguments are
74# a recipe for deriving the shape (a divisor and an expert count) rather than
75# the canonical fields :func:`_explicit_overrides` can harvest, so nothing
76# short of calling it knows what the run builds. A factory is consulted
77# BEFORE ``AutoConfig``, because where both answer they answer differently and
78# the factory is the one that says what the run trains: V4.1's crops the
79# released model eightfold, and the released config would price a model nobody
80# runs. One line a family, and the entry is the declaration that the two are
81# a pair.
82_CONFIG_FACTORIES = {
83 "hyper_parallel.models.deepseek_v41.adapter.validation.cropped_model."
84 "build_cropped_deepseek_v41":
85 "hyper_parallel.models.deepseek_v41.adapter.validation.cropped_model."
86 "build_deepseek_v41_validation_config",
87}
89# Canonical text-tower fields, with the Transformers aliases that carry them.
90# First populated alias wins, so Hyper-internal names keep priority over the
91# Hugging Face spelling when a config happens to define both.
92_TEXT_FIELD_ALIASES: Dict[str, Tuple[str, ...]] = {
93 "hidden_size": ("hidden_size",),
94 "num_hidden_layers": ("num_hidden_layers",),
95 "num_attention_heads": ("num_attention_heads",),
96 "num_key_value_heads": ("num_key_value_heads",),
97 "head_dim": ("head_dim",),
98 "intermediate_size": ("intermediate_size",),
99 "vocab_size": ("vocab_size",),
100 "max_position_embeddings": ("max_position_embeddings",),
101 "num_experts": ("num_experts", "n_routed_experts"),
102 "num_experts_per_tok": ("num_experts_per_tok",),
103 "num_shared_experts": ("n_shared_experts", "num_shared_experts"),
104 "moe_intermediate_size": ("moe_intermediate_size",),
105 "shared_expert_intermediate_size": ("shared_expert_intermediate_size",),
106 "first_k_dense_replace": ("first_k_dense_replace",),
107 "mtp_depth": ("mtp_depth", "num_nextn_predict_layers", "mtp_num_hidden_layers"),
108 "multiple_of": ("multiple_of",),
109 "ffn_dim_multiplier": ("ffn_dim_multiplier",),
110 "kv_lora_rank": ("kv_lora_rank",),
111 "q_lora_rank": ("q_lora_rank",),
112 "qk_rope_head_dim": ("qk_rope_head_dim",),
113 "qk_nope_head_dim": ("qk_nope_head_dim",),
114 "v_head_dim": ("v_head_dim",),
115 "attn_output_gate": ("attn_output_gate",),
116 "tie_word_embeddings": ("tie_word_embeddings",),
117 # Llama's attention_bias biases all four projections; Qwen2's configs
118 # state none and bias the query, key and value projections alone.
119 "qkv_bias": ("qkv_bias", "attention_bias"),
120 "o_bias": ("o_bias", "attention_bias"),
121 "mlp_bias": ("mlp_bias",),
122 # Stated by a few families; Qwen3's config does not state its own.
123 "qk_norm": ("qk_norm", "use_qk_norm", "qk_layernorm"),
124 # A hybrid stack states its layers; without this every layer is costed
125 # as full attention, which is quadratic in the sequence length.
126 "layer_types": ("layer_types",),
127 "linear_num_key_heads": ("linear_num_key_heads",),
128 "linear_key_head_dim": ("linear_key_head_dim",),
129 "linear_num_value_heads": ("linear_num_value_heads",),
130 "linear_value_head_dim": ("linear_value_head_dim",),
131 "linear_conv_kernel_dim": ("linear_conv_kernel_dim",),
132 # The dtype the checkpoint's weights are saved in, which the runtime
133 # loads them in under its default torch_dtype of auto; newer
134 # Transformers name it dtype.
135 "torch_dtype": ("dtype", "torch_dtype"),
136}
138# Vision towers use their own spelling for the shared concepts.
139_VISION_FIELD_ALIASES: Dict[str, Tuple[str, ...]] = {
140 "hidden_size": ("hidden_size",),
141 "num_hidden_layers": ("depth", "num_hidden_layers"),
142 "num_attention_heads": ("num_heads", "num_attention_heads"),
143 "intermediate_size": ("intermediate_size",),
144 "out_hidden_size": ("out_hidden_size",),
145 "patch_size": ("patch_size",),
146 "spatial_merge_size": ("spatial_merge_size",),
147 "num_position_embeddings": ("num_position_embeddings",),
148}
150# Fallback when a vision tower declares no positional-embedding grid.
151_DEFAULT_VISUAL_SEQ_LEN = 1024
153# Keys of ``model.config_overrides`` that state the run, not the model. The
154# model spec never sees them; the Hyper parser reads them into its ExecSpec.
155EXEC_OVERRIDE_KEYS = (
156 "offset", "full_rec", "sel_rec", "capacity_factor", "cap_fact", "use_gmm", "gmm", "seq_length",
157)
160def _declares_pretrained_path(mapping: Any) -> bool:
161 """Return whether the ``model`` section names a Transformers checkpoint.
163 This is the one key only this schema has: the older HyperParallel yamls
164 name a checkpoint as ``model.weights_path``. Recognising it means a
165 config needs no empty ``training``/``accelerator``/``fsdp_config``
166 section just to be identified.
167 """
168 if isinstance(mapping, Mapping):
169 model = mapping.get("model")
170 else:
171 model = getattr(mapping, "model", None)
172 if model is None:
173 return False
174 if isinstance(model, Mapping):
175 return bool(model.get("pretrained_model_name_or_path"))
176 return bool(getattr(model, "pretrained_model_name_or_path", None))
179def is_auto_models_schema(mapping: Any) -> bool:
180 """Return whether *mapping* looks like an AutoModels Trainer config.
182 Accepts a plain mapping or any object exposing the root sections as
183 attributes, so the SAPP-ND ``Config`` tree and a parsed YAML dict can be
184 tested with the same call.
185 """
186 if _declares_pretrained_path(mapping):
187 return True
188 if isinstance(mapping, Mapping):
189 return any(name in mapping for name in _AUTO_MODELS_SECTIONS)
190 holder = getattr(mapping, "__dict__", None)
191 if isinstance(holder, dict):
192 return any(name in holder for name in _AUTO_MODELS_SECTIONS)
193 return False
196def _first_attr(config: Any, names: Tuple[str, ...], default: Any = None) -> Any:
197 """Return the first populated attribute of *config* among *names*."""
198 for name in names:
199 value = getattr(config, name, None)
200 if value is not None:
201 return value
202 return default
205def _text_tower(model_config: Any) -> Any:
206 """Return the language-model sub-config of a composite config.
208 Multimodal Transformers configs (``Qwen3VLMoeConfig`` and friends) keep
209 the language model under ``text_config`` and expose none of its fields at
210 the top level, so reading ``hidden_size`` directly raises AttributeError.
211 """
212 return getattr(model_config, "text_config", None) or model_config
215def _at_stated_depth(model_config: Any, explicit: Mapping[str, Any]) -> Any:
216 """*model_config* at the depth the model section states, as a crop of the released model is built.
218 A ``num_hidden_layers`` stated beside the checkpoint crops it to its first layers. The formulas read that
219 count as stated, but the census builds its layers from the config, so the config it builds from has to
220 carry the crop too, its ``layer_types`` cut to the same layers; otherwise the census lists the release's
221 layers against the crop's count.
223 Returns:
224 *model_config* itself where the section states no depth or the release's own, else a copy at that depth.
225 """
226 depth = explicit.get("num_hidden_layers")
227 if depth is None or getattr(_text_tower(model_config), "num_hidden_layers", None) == int(depth):
228 return model_config
229 cropped = copy.deepcopy(model_config)
230 text = _text_tower(cropped)
231 text.num_hidden_layers = int(depth)
232 layer_types = getattr(text, "layer_types", None)
233 if isinstance(layer_types, (list, tuple)):
234 text.layer_types = list(layer_types)[:int(depth)]
235 return cropped
238def _spec_from_aliases(config: Any, aliases: Dict[str, Tuple[str, ...]]) -> Dict[str, Any]:
239 """Collect the canonical fields declared by *config*."""
240 spec = {}
241 for canonical, names in aliases.items():
242 value = _first_attr(config, names)
243 if value is not None:
244 spec[canonical] = value
245 return spec
248def _derive_shared_experts(spec: Dict[str, Any]) -> None:
249 """Infer the shared-expert count from its total feed-forward width.
251 Qwen2-MoE style configs declare one wide shared expert instead of a
252 count. The cost model only ever uses ``n_shared_exp * hff_exp``, so
253 encoding that width as several narrow experts is width-equivalent.
254 """
255 if spec.get("num_shared_experts"):
256 return
257 moe_inter = int(spec.get("moe_intermediate_size", 0) or 0)
258 shared_inter = int(spec.get("shared_expert_intermediate_size", 0) or 0)
259 if moe_inter and shared_inter:
260 spec["num_shared_experts"] = max(1, shared_inter // moe_inter)
263def _derive_dense_ffn_width(spec: Dict[str, Any]) -> None:
264 """Fill the dense feed-forward width for an all-MoE config.
266 A model whose every layer is MoE (Qwen3.5-MoE) declares no
267 ``intermediate_size``, but the cost model prices the shared expert at
268 that width, so leaving it absent drops the shared expert entirely.
269 The shared expert's own width is the value meant there.
270 """
271 if spec.get("intermediate_size"):
272 return
273 width = spec.get("shared_expert_intermediate_size") or spec.get("moe_intermediate_size")
274 if width:
275 spec["intermediate_size"] = int(width)
278def _visual_seq_len(vision_spec: Dict[str, Any], override: Optional[int]) -> int:
279 """Resolve the encoder sequence length in merged visual tokens.
281 The true count depends on the image resolution served by the dataset,
282 which no configuration we read carries. The positional-embedding grid
283 divided by the spatial merge is the tightest bound available from the
284 model alone, so it is the default and *override* exists to correct it.
285 """
286 if override:
287 logger.info("visual sequence length taken from context.visual_seq_len: %d", override)
288 return int(override)
289 grid = int(vision_spec.get("num_position_embeddings", 0) or 0)
290 merge = max(1, int(vision_spec.get("spatial_merge_size", 1) or 1))
291 if not grid:
292 logger.warning(
293 "vision tower declares no num_position_embeddings; assuming %d visual tokens. "
294 "Set context.visual_seq_len to the value your dataset produces.",
295 _DEFAULT_VISUAL_SEQ_LEN,
296 )
297 return _DEFAULT_VISUAL_SEQ_LEN
298 derived = max(1, grid // (merge * merge))
299 logger.info(
300 "visual sequence length derived from the vision config: %d "
301 "(num_position_embeddings=%d, spatial_merge_size=%d). "
302 "Set context.visual_seq_len to override.",
303 derived, grid, merge,
304 )
305 return derived
308def _get_hf_config(model_raw: Mapping[str, Any]) -> Any:
309 """Call the AutoModels Transformers entry point for *model_raw*."""
310 # Transformers is an optional dependency; import only when a config
311 # actually needs it so the other cost-model backends stay importable.
312 from hyper_parallel.models._transformers.config_resolver import get_hf_config # pylint: disable=C0415
314 config_kwargs = {
315 name: model_raw[name]
316 for name in _CONFIG_KWARG_NAMES
317 if model_raw.get(name) is not None
318 }
319 return get_hf_config(
320 str(_model_path(model_raw)),
321 str(model_raw.get("attn_implementation", "sdpa")),
322 model_raw.get("torch_dtype", "auto"),
323 **config_kwargs,
324 )
327def checkpoint_configs(model_raw: Mapping[str, Any]) -> Tuple[Any, Any]:
328 """The Transformers config of the checkpoint *model_raw* names, and its language model's."""
329 config = _get_hf_config(model_raw)
330 return config, _text_tower(config)
333def _factory_config(model_raw: Mapping[str, Any]) -> Any:
334 """The config the recipe's own factory builds, or None where it names none.
336 The model section's keys are offered to the factory by name, filtered to
337 the ones it accepts, so a recipe states its shape once and the cost model
338 reads the same statement the trainer does. Anything the factory cannot
339 take, such as a dtype or a placement switch, is left out rather than
340 refused.
342 Returns:
343 A Transformers config, or None where ``_target_`` names no factory
344 this knows.
346 Raises:
347 ValueError: If the factory is known but cannot build its config, since
348 a recipe that names one and then fails is a broken recipe rather
349 than a model to price by formula.
350 """
351 import inspect # pylint: disable=C0415
353 target = str(model_raw.get("_target_") or "")
354 path = _CONFIG_FACTORIES.get(target)
355 if path is None:
356 return None
357 module_path, _, name = path.rpartition(".")
358 try:
359 factory = getattr(importlib.import_module(module_path), name)
360 accepted = inspect.signature(factory).parameters
361 offered = {key: value for key, value in model_raw.items() if key in accepted}
362 logger.info("building the config %s states, through %s", target.rpartition(".")[2], name)
363 return factory(**offered)
364 except (ImportError, AttributeError, TypeError, OSError, KeyError, ValueError) as exc:
365 raise ValueError(
366 f"model._target_ names {target}, whose config factory {name} could not build its "
367 f"config ({exc}); the cost model prices what that factory builds, so fix the recipe "
368 f"rather than let the formulas price a different model"
369 ) from exc
372def _model_path(model_raw: Mapping[str, Any]) -> Optional[str]:
373 """Return the model directory the section names, under any accepted key."""
374 for alias in _MODEL_PATH_ALIASES:
375 value = model_raw.get(alias)
376 if value:
377 return str(value)
378 return None
381def _explicit_overrides(model_raw: Mapping[str, Any]) -> Dict[str, Any]:
382 """Return the model fields the model section states itself, as a plain dict.
384 ``config_overrides`` is the declared form, less its run keys
385 (:data:`EXEC_OVERRIDE_KEYS`). A from_config recipe instead passes the shape
386 it builds as factory arguments beside ``_target_``, such as a
387 ``num_hidden_layers`` that crops the released model; those name canonical
388 fields, so they are overrides too, and the released config must not win
389 over them. ``config_overrides`` keeps priority where both are present. A
390 value of ``auto``, which ``torch_dtype`` takes to load the checkpoint's
391 own, states nothing over the checkpoint.
392 """
393 overrides = model_raw.get("config_overrides")
394 explicit = {}
395 if isinstance(overrides, Mapping):
396 explicit = {key: value for key, value in overrides.items() if key not in EXEC_OVERRIDE_KEYS}
397 for field in _TEXT_FIELD_ALIASES:
398 value = model_raw.get(field)
399 if value is not None and value != "auto":
400 explicit.setdefault(field, value)
401 return explicit
404def exec_overrides(model_raw: Mapping[str, Any]) -> Dict[str, Any]:
405 """Return the run keys of the ``config_overrides`` fallback.
407 Args:
408 model_raw: The ``model`` section, as a plain mapping.
410 Returns:
411 The entries whose keys are in :data:`EXEC_OVERRIDE_KEYS`.
412 """
413 overrides = model_raw.get("config_overrides")
414 if not isinstance(overrides, Mapping):
415 return {}
416 return {key: value for key, value in overrides.items() if key in EXEC_OVERRIDE_KEYS}
419def _validated(spec: Dict[str, Any]) -> Dict[str, Any]:
420 """Check *spec* against the model IR before any consumer reads it.
422 Parsing and re-serialising is what makes the schema load-bearing rather
423 than advisory: a field the model declares but the IR does not know would
424 be dropped here, and an incoherent one raises by name. The op profile
425 and the layer stack are settled here too, so the serialised spec states
426 what it will be priced with and which kind each layer is, even when the
427 family was only inferred from the model name and the stack from
428 ``layer_types`` or ``first_k_dense_replace``; and so is whether its
429 attention normalizes queries and keys, which a Transformers config does
430 not state, and the biases and norms its layers hold: a Transformers
431 decoder layer holds two RMSNorms, no bias unless its config states one
432 or its family has one, and a shared expert gated where its family
433 gates it. ``layer_types`` is
434 consumed: ``layers`` says the same, checked against the profile. A
435 vision tower states its stack likewise.
436 """
437 typed = ModelSpec.from_dict(spec)
438 typed = dataclasses.replace(
439 typed, arch=typed.arch or infer_arch(typed.name),
440 qk_norm=typed.qk_norm if typed.qk_norm is not None else infer_qk_norm(typed.name),
441 qkv_bias=typed.qkv_bias if typed.qkv_bias is not None else infer_qkv_bias(typed.name),
442 o_bias=bool(typed.o_bias), mlp_bias=bool(typed.mlp_bias), norm_bias=bool(typed.norm_bias),
443 layer_norms=typed.layer_norms or 2,
444 shared_expert_gate=(typed.shared_expert_gate if typed.shared_expert_gate is not None
445 else infer_shared_expert_gate(typed.name)),
446 )
447 resolve_ops(typed.arch, typed.ops)
448 layers = spec_layer_stack(typed).to_layers()
449 vision = typed.vision
450 if vision is not None:
451 vision = dataclasses.replace(vision, layers=tower_layer_stack(vision).to_layers())
452 return dataclasses.replace(typed, layers=layers, vision=vision, layer_types=None).to_dict()
453# Canonical fields that say how a model is held, not what it is: a section
454# stating only these describes no model.
455_HOLDING_FIELDS = frozenset({"torch_dtype"})
458def _describes_a_model(explicit: Mapping[str, Any]) -> bool:
459 """Whether the fields a section states describe a model, beyond how it is held."""
460 return any(key not in _HOLDING_FIELDS for key in explicit)
463def _dtype_name(value: Any) -> Any:
464 """A dtype as its plain name, ``bfloat16`` for ``torch.bfloat16``; anything else as it is."""
465 return str(value).replace("torch.", "") if value is not None else None
468# The censuses this process has run, by config, layer stack and length: a
469# harness pricing several runs of one model runs one.
470_CENSUSES: Dict[Tuple[Any, ...], Dict[str, Any]] = {}
472# Where a census is kept between runs, which it then costs nothing:
473# ND_CENSUS_CACHE names another folder, or is "off" to keep none.
474_CENSUS_FOLDER = os.path.join(os.path.expanduser("~"), ".cache", "hyper_parallel", "census")
475# The package whose files a kept census is checked against.
476_PACKAGE_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
477# The search's own code, which runs after a census and hands it only what its
478# key holds: a change there leaves the kept censuses valid.
479_NOT_CENSUS_CODE = (
480 "hyper_parallel.auto_parallel.sapp_nd",
481 "hyper_parallel.auto_parallel.sapp_ppb",
482 "hyper_parallel.auto_parallel.config_adapter",
483)
486def _census_file(key: Tuple[Any, ...]) -> Optional[str]:
487 """The file a census of *key* is kept in, named by the key and the torch and Transformers it runs on.
489 ``None`` where ``ND_CENSUS_CACHE`` is ``off``.
490 """
491 folder = os.environ.get("ND_CENSUS_CACHE") or _CENSUS_FOLDER
492 if folder.lower() == "off":
493 return None
494 versions = [str(getattr(sys.modules.get(name), "__version__", None)) for name in ("torch", "transformers")]
495 name = hashlib.sha256(json.dumps([repr(key), *versions]).encode("utf-8")).hexdigest()
496 return os.path.join(folder, name[:32] + ".json")
499def _source_digest(path: str) -> str:
500 """The SHA-256 of the file at *path*, its line endings read as Python reads them.
502 Two checkouts of one commit can hold a file with Windows line endings in
503 one and Unix ones in the other, which is the same code.
504 """
505 with open(path, "rb") as handle:
506 return hashlib.sha256(handle.read().replace(b"\r\n", b"\n")).hexdigest()
509def _census_sources() -> Dict[str, str]:
510 """The digest of each file of the package this process has loaded, the search's left out, by its path.
512 A census runs on them: its own code, the kernels' contracts, the
513 checkpoint policy and the modules a run's replacements build.
514 """
515 sources = {}
516 for name, module in list(sys.modules.items()):
517 path = getattr(module, "__file__", None)
518 if path and name.split(".")[0] == "hyper_parallel" and not name.startswith(_NOT_CENSUS_CODE):
519 sources[os.path.relpath(path, _PACKAGE_ROOT).replace(os.sep, "/")] = _source_digest(path)
520 return sources
523def _kept_census(path: str) -> Optional[Dict[str, Any]]:
524 """The census kept in *path*, or ``None`` where there is none or a file it ran on has changed since."""
525 try:
526 with open(path, encoding="utf-8") as handle:
527 kept = json.load(handle)
528 if all(_source_digest(os.path.join(_PACKAGE_ROOT, source)) == digest
529 for source, digest in kept["sources"].items()):
530 return kept["census"]
531 except (OSError, ValueError, KeyError, TypeError, AttributeError):
532 pass
533 return None
536def _keep_census(path: str, census: Dict[str, Any]) -> None:
537 """Keep *census* in *path* with the digest of each file it ran on; failing costs the next run a census."""
538 partial = f"{path}.{os.getpid()}"
539 try:
540 os.makedirs(os.path.dirname(path), exist_ok=True)
541 with open(partial, "w", encoding="utf-8") as handle:
542 json.dump({"sources": _census_sources(), "census": census}, handle)
543 os.replace(partial, path)
544 except (OSError, TypeError, ValueError) as exc:
545 logger.warning("the census could not be kept in %s: %s", path, exc)
548def _census(text_config: Any, layers: Any, seq_length: int, replacements: Tuple[Any, ...] = ()) -> Dict[str, Any]:
549 """The spec's census of each layer kind of *layers* and of the output layer, once per config and length.
551 A census is kept on disk (:func:`_census_file`) and read back by a later
552 run of the same model, length and replacements, on the same torch and
553 Transformers, while no file of the package it ran on has changed.
554 """
555 key = (
556 text_config.to_json_string(),
557 tuple((group["kind"], int(group["count"]), bool(group.get("mtp"))) for group in layers),
558 int(seq_length),
559 tuple((f"{spec.factory.__module__}.{spec.factory.__qualname__}",
560 f"{spec.module_type.__module__}.{spec.module_type.__qualname__}",
561 tuple(spec.match), bool(spec.exact_type)) for spec in replacements),
562 )
563 if key not in _CENSUSES:
564 path = _census_file(key)
565 kept = _kept_census(path) if path is not None else None
566 if kept is not None:
567 logger.info("census of each layer kind and of the output layer at %d tokens read from %s",
568 seq_length, path)
569 _CENSUSES[key] = kept
570 return copy.deepcopy(kept)
571 logger.info("census of each layer kind and of the output layer at %d tokens%s", seq_length,
572 f", {len(replacements)} module replacements installed" if replacements else "")
573 try:
574 kinds = census_activations(text_config, layers, seq_length, replacements)
575 _CENSUSES[key] = {
576 "activations": {kind: record.to_dict() for kind, record in kinds.items()},
577 "output_activations": census_output_activations(text_config, seq_length, replacements).to_dict(),
578 }
579 except CensusUnavailable as exc:
580 # The formulas price the layers, as they do without a census: a
581 # census of Transformers' own modules would price a layer the run
582 # does not build.
583 logger.warning("%s", _NOT_A_BOUND % exc)
584 _CENSUSES[key] = {}
585 else:
586 if path is not None:
587 _keep_census(path, _CENSUSES[key])
588 return copy.deepcopy(_CENSUSES[key])
591# What a run is told when a census it asked for did not happen. Measured
592# without one, the formulas read 37 to 55% below Qwen3.5's peak and 51.1%
593# below the DeepSeek-V4.1 crop's, so a memory a search filters on is an
594# estimate and not an upper bound, and a strategy it calls feasible may not
595# be. Said once, in the run's own log, where the number is read.
596_NOT_A_BOUND = (
597 "no census of the layers: %s. The formulas price them instead, and without a census they have "
598 "measured 37 to 55%% below a real peak, so this run's memory is an estimate and NOT a bound: a "
599 "strategy it reports as fitting may still run out of memory"
600)
603def _no_census(census_seq_len: int, explicit: Mapping[str, Any]) -> None:
604 """Warn that a census asked for cannot run, unless the overrides state its records.
606 A census builds its layers from the checkpoint's config; a search hands
607 ND a spec whose overrides state the records its reader measured.
608 """
609 if census_seq_len and not explicit.get("activations"):
610 logger.warning("%s", _NOT_A_BOUND % "it needs the checkpoint's Transformers config")
613def _stated_spec(model_raw: Mapping[str, Any], explicit: Dict[str, Any], census_seq_len: int) -> Dict[str, Any]:
614 """The spec of a model its section states by hand, derived and settled as a resolved one is.
616 A model no Transformers config describes, every model Transformers
617 cannot build, is stated field by field, and it takes the derivations a
618 resolved spec does: a shared expert stated by its width counts as that
619 many experts, and an all-MoE model's dense width is its shared
620 expert's. Without them a hand-described all-MoE model kept no
621 ``intermediate_size``, the width its shared expert is priced at, and
622 its shared expert was dropped.
623 """
624 explicit.setdefault("name", model_raw.get("name", "custom"))
625 _derive_shared_experts(explicit)
626 _derive_dense_ffn_width(explicit)
627 _no_census(census_seq_len, explicit)
628 return _validated(explicit)
631def resolve_hf_model_spec(
632 model_raw: Mapping[str, Any],
633 visual_seq_len: Optional[int] = None,
634 census_seq_len: int = 0,
635 census_spec: bool = False,
636 replacements: Tuple[Any, ...] = (),
637) -> Dict[str, Any]:
638 """Return canonical cost-model fields for a Trainer ``model`` section.
640 Resolves ``model.pretrained_model_name_or_path`` through the same
641 Transformers path AutoModels uses. A composite (vision-language) config
642 contributes its language tower under the canonical keys and its vision
643 tower under ``"vision"``; a plain config leaves ``"vision"`` absent.
645 ``model.config_overrides`` stays supported for standalone cost-model
646 search files, and doubles as the fallback when the Transformers config
647 cannot be reached (offline node, unreachable repository). Explicit
648 overrides always win over resolved values. Its run keys,
649 :data:`EXEC_OVERRIDE_KEYS`, are not the model's: :func:`exec_overrides`
650 returns them.
652 Args:
653 model_raw: The ``model`` section, as a plain mapping.
654 visual_seq_len: Optional override for the encoder sequence length.
655 census_seq_len: The tokens to run a census of each layer kind of
656 the language model at, and of its output layer
657 (:mod:`hyper_parallel.auto_parallel._layer_census`), which the
658 spec states as ``"activations"`` and ``"output_activations"``;
659 0 runs none. The
660 census builds its layers from the checkpoint's config: a spec
661 from ``config_overrides`` alone gets none, and an override of a
662 model field does not reach it.
663 census_spec: Whether the language model's fields, its layer stack
664 and its op counts are the census's, measured on the layers
665 Transformers builds of the checkpoint
666 (:func:`~hyper_parallel.auto_parallel._spec_census.census_model_spec`),
667 rather than read from its config's fields. A spec from
668 ``config_overrides`` alone is the overrides' either way.
669 replacements: The module replacements the run installs
670 (:func:`~hyper_parallel.auto_parallel._layer_census.replacement_specs`),
671 which the census runs its layers with, so it measures what the
672 trainer saves.
674 Returns:
675 A dict of canonical model fields, always carrying ``"name"``.
677 Raises:
678 ValueError: If neither a pretrained path nor overrides can supply
679 the model dimensions.
680 """
681 explicit = _explicit_overrides(model_raw)
682 model_path = _model_path(model_raw)
684 if not model_path:
685 if _describes_a_model(explicit):
686 return _stated_spec(model_raw, explicit, census_seq_len)
687 raise ValueError(
688 "AutoModels train.yaml requires model.pretrained_model_name_or_path, "
689 "model.config_path or model.config_overrides for Auto Parallel search"
690 )
692 # A recipe that builds its model through a factory is asked first: the
693 # path beside _target_ is a directory that factory reads, not a checkpoint
694 # to price, and where the two disagree the factory is what the run trains.
695 model_config = _factory_config(model_raw)
696 try:
697 model_config = model_config if model_config is not None else _get_hf_config(model_raw)
698 except (ImportError, OSError, ValueError, TypeError, AttributeError, KeyError) as exc:
699 if _describes_a_model(explicit):
700 logger.warning(
701 "Transformers config resolution failed (%s); "
702 "falling back to model.config_overrides", exc,
703 )
704 return _stated_spec(model_raw, explicit, census_seq_len)
705 raise ValueError(
706 f"cannot resolve model.pretrained_model_name_or_path '{model_path}'; "
707 "install transformers, set model.config_overrides, or make the config "
708 "available offline (warm the HF_HOME cache, or pass local_files_only)"
709 ) from exc
711 name = str(getattr(model_config, "model_type", None) or model_path)
712 if census_spec:
713 spec = census_model_spec(_at_stated_depth(model_config, explicit), explicit.get("arch"), name)
714 else:
715 spec = _spec_from_aliases(_text_tower(model_config), _TEXT_FIELD_ALIASES)
716 _derive_shared_experts(spec)
717 _derive_dense_ffn_width(spec)
718 spec["name"] = name
719 # The dtype the checkpoint is saved in, by name, which the census's
720 # fields do not state.
721 dtype = _first_attr(_text_tower(model_config), _TEXT_FIELD_ALIASES["torch_dtype"])
722 if dtype is not None:
723 spec["torch_dtype"] = _dtype_name(dtype)
725 vision_config = getattr(model_config, "vision_config", None)
726 if vision_config is not None:
727 vision_spec = _spec_from_aliases(vision_config, _VISION_FIELD_ALIASES)
728 vision_spec["name"] = f"{spec['name']}_vision"
729 vision_spec["max_position_embeddings"] = _visual_seq_len(vision_spec, visual_seq_len)
730 spec["vision"] = vision_spec
732 spec.update(explicit)
733 resolved = _validated(spec)
734 if census_seq_len:
735 resolved.update(_census(_text_tower(model_config), resolved["layers"], census_seq_len, replacements))
736 resolved = ModelSpec.from_dict(resolved).to_dict()
737 return resolved
740def resolve_model_spec(
741 model_raw: Mapping[str, Any],
742 visual_seq_len: Optional[int] = None,
743 census_seq_len: int = 0,
744 census_spec: bool = False,
745 replacements: Tuple[Any, ...] = (),
746) -> ModelSpec:
747 """Return the typed :class:`ModelSpec` for a Trainer ``model`` section.
749 The typed entry point behind :func:`resolve_hf_model_spec`, which returns
750 the same facts as a plain mapping for callers that still read one.
752 Args:
753 model_raw: The ``model`` section, as a plain mapping.
754 visual_seq_len: Optional override for the encoder sequence length.
755 census_seq_len, census_spec, replacements: As
756 :func:`resolve_hf_model_spec` takes them.
758 Returns:
759 A validated spec.
761 Raises:
762 ModelSpecError: If the resolved fields are incomplete or incoherent.
763 ValueError: If neither a pretrained path nor overrides can supply
764 the model dimensions.
765 """
766 return ModelSpec.from_dict(
767 resolve_hf_model_spec(model_raw, visual_seq_len, census_seq_len, census_spec, replacements))