models#

Per-model descriptors, one package per HF model type.

Directory names are the HF config.model_type they describe; a model whose code ships with the checkpoint uses its remote-code model_type and records that as ModelSpec.modeling_source, while a model defined in transformers records the release its definitions come from as ModelSpec.min_transformers_version. Each package’s specs.py registers one global ModelSpec (built from the section classes in this package’s own specs.py) at import time. Importing this package registers them all. Consumers resolve a spec via the registry lookups (get_spec / match_moe_block) and read its sections.

The per-model file is named for what it holds, not for who reads it: a model’s spec is general model data, and export is only its first consumer.

Classes

ExportSpec

Subsystem section: per-model data of the unified HF export path.

MoESpec

Topic section: the model's MoE-block layout.

ModelSpec

The one global per-model descriptor, holding each section as an attribute.

SpecSection

Base class for the sections composing a ModelSpec.

Functions

get_spec

Return the spec registered for model_type, or None.

get_specs

Return all registered specs, in registration order.

hf_model_type

Return the root HF model type (model.config.model_type), or None.

is_moe

Return whether module is an MoE block.

list_all_possible

List a spec attribute's values across all registered specs, deduplicated in order.

match_class_names

Return True if any of names equals a class name in module's MRO.

match_moe_block

Return the MoE layout layout for module, resolved by model type.

match_moe_model

Return the spec whose MoE section matches module, or None.

register

Register a model spec and return it.

class ExportSpec#

Bases: SpecSection

Subsystem section: per-model data of the unified HF export path.

Architecture facts (MoE block classes, expert naming) live in MoESpec; this section holds data consumed by the export algorithms only.

__init__(*, grouped_expert_export=False, pqs_fuse_rules=(), weight_plus_one_norm_names=())#
Parameters:
  • grouped_expert_export (bool)

  • pqs_fuse_rules (tuple[tuple[tuple[str, ...], str, str], ...])

  • weight_plus_one_norm_names (tuple[str, ...])

Return type:

None

grouped_expert_export: bool = False#

Whether get_experts_list may group this model’s experts for the AWQ / NVFP4-SVDQuant resmoothing pass.

A statement about what modelopt has validated, not about the model: qwen3_5_moe is architecturally identical to qwen3_moe here and is still False, because the pre-registry code keyed off the root class name and "qwen3_5moeforcausallm" matched none of its qwen substrings. That is why it lives in the export section rather than on MoESpec, which holds only facts about the architecture.

Whether the experts are actually iterable is separate again, and read off the module: a model listed here still exports fine when a newer transformers fuses its experts, because there is then simply nothing to group.

pqs_fuse_rules: tuple[tuple[tuple[str, ...], str, str], ...] = ()#

AWQ pre_quant_scale fusion rules, each a (module_class_substrings, fuse_into, fuse_from) triple: for a module whose class name contains one of the substrings, the pre_quant_scale on fuse_from is folded into fuse_into (e.g. attention o_proj -> v_proj, MLP down_proj -> up_proj). A rule asserts mathematical equivalence for that model’s modules, so it is declared per model rather than applied generically.

weight_plus_one_norm_names: tuple[str, ...] = ()#

Class names of norm layers whose stored weight is w - 1 (the effective scale is weight + 1), e.g. Gemma’s RMSNorm layouts and LayerNorm1P. Matched against a norm module’s MRO (case-insensitive exact names). Export must account for the +1 when folding scales into the norm weight (AWQ pre_quant_scale fusion). A structural fallback (zero_centered_gamma) stays in the engine.

class MoESpec#

Bases: SpecSection

Topic section: the model’s MoE-block layout.

Describes what a model’s MoE blocks are – which class, what the expert projections are called, how they are stored – so any modelopt subsystem (export, quantization, speculative decoding, …) can read it instead of keeping its own per-model MoE table.

One layout per model. block_names is a tuple, so a model whose MoE appears under several class names is covered as long as they share a layout (gpt_oss’s GptOssMLP/GptOssMoE; Qwen3-Omni’s Thinker and Talker blocks). No model in transformers needs two different layouts today. If one ever does, this section becomes a tuple on ModelSpec; nothing here is shaped to prevent that.

A layout is a fact about the architecture. Whether modelopt’s grouped export path may use it is policy and lives in ExportSpec; whether the experts are iterable right now depends on the installed transformers and is read off the module.

__init__(*, block_names=(), expert_linear_names=None, fused_expert_names=False, gate_up_pair=None)#
Parameters:
  • block_names (tuple[str, ...])

  • expert_linear_names (tuple[str, ...] | None)

  • fused_expert_names (bool)

  • gate_up_pair (tuple[str, str] | None)

Return type:

None

block_names: tuple[str, ...] = ()#

The matching key – MoE block class names, matched against the module’s MRO (case-insensitive exact names, not substrings; see match_class_names).

expert_linear_names: tuple[str, ...] | None = None#

Expert linear projection names, e.g. ("gate_proj", "down_proj", "up_proj"). For layouts modelopt rewrites (e.g. quantized DBRX), these are the names on the rewritten module.

expert_linear_names_for(module, fused=None)#

Resolve module’s expert linear names, or None when this layout cannot.

The module’s class is not consulted: a spec may provide naming without declaring block_names at all, and with one layout per model there is nothing to disambiguate.

fused is the layout the caller observed on the module. Passing it makes the answer conditional on describing that same layout, so a spec that only describes the per-expert form declines for a fused container rather than returning names that do not apply there. None accepts either.

Parameters:

fused (bool | None)

Return type:

tuple[str, …] | None

fused_expert_names: bool = False#

True when expert_linear_names name a fused container’s own 3-D parameters (gate_up_proj/down_proj) rather than per-expert sub-modules.

The same model type materializes both ways across transformers releases, so a consumer must know which layout a naming describes: applied to the other one it is wrong, not merely generic.

gate_up_pair: tuple[str, str] | None = None#

The (gate, up) pair among expert_linear_names that serving engines fuse into a single gate_up_proj, e.g. ("gate_proj", "up_proj") or ("w1", "w3"). None for non-gated experts (NemotronH) and already-fused layouts (GptOss, DBRX).

property gate_up_pairs: tuple[tuple[str, str], ...]#

This layout’s (gate, up) pair, as a tuple so it can join a global vocabulary.

matches(module)#

Whether module’s class is one this layout describes.

Return type:

bool

class ModelSpec#

Bases: object

The one global per-model descriptor, holding each section as an attribute.

Resolved by HF model type (see get_spec below); a model registers exactly one instance, filling only the sections it customizes and leaving the rest None. Consumers must handle an absent section; match_moe_block and list_all_possible already do for the lookups they cover.

__init__(*, model_type, min_transformers_version=None, modeling_source='transformers', moe_spec=None, export_spec=None)#
Parameters:
  • model_type (str)

  • min_transformers_version (str | None)

  • modeling_source (Literal['transformers', 'remote_code'])

  • moe_spec (MoESpec | None)

  • export_spec (ExportSpec | None)

Return type:

None

export_spec: ExportSpec | None = None#

Per-model data of the unified HF export path, or None when export needs nothing model-specific.

min_transformers_version: str | None = None#

Earliest transformers release whose definitions match this spec, or None when the question does not apply (modeling_source="remote_code").

Clamped below at the repo’s minimum supported transformers (tf_min in noxfile.py): a model that predates the floor records the floor, since nothing older is ever installed or tested. A model added after the floor records its own release, which is what lets the test suite tell an expected absence on an older transformers apart from a spec that no longer matches reality.

model_type: str#

The HF model type this spec describes (config.model_type, e.g. "qwen3_moe"). Unique across the registry.

modeling_source: Literal['transformers', 'remote_code'] = 'transformers'#

shipped inside transformers, or carried by the checkpoint and loaded with trust_remote_code=True.

A fact about the model, not about any one subsystem: it decides whether the classes named in this spec can be imported at all, so model_type is a remote-code spelling rather than a transformers one, and loading the model needs the flag.

Type:

Where the model’s modeling code lives

moe_spec: MoESpec | None = None#

The model’s MoE architecture facts, or None for a dense model.

class SpecSection#

Bases: object

Base class for the sections composing a ModelSpec.

Marks a class as a section so _spec_sections can find ModelSpec’s section fields from its annotations. Without it the lookups would need a hand-maintained list of section names, which would silently fall out of date the first time a section was added and the list was not.

__init__()#
Return type:

None

get_spec(model_type)#

Return the spec registered for model_type, or None.

Parameters:

model_type (str)

Return type:

ModelSpec | None

get_specs()#

Return all registered specs, in registration order.

Return type:

list[ModelSpec]

hf_model_type(model)#

Return the root HF model type (model.config.model_type), or None.

Accepts a model or a config object (duck-typed, no transformers import). This is the key for get_spec / match_moe_block.

Return type:

str | None

is_moe(module, model_type=None)#

Return whether module is an MoE block.

model_type (model.config.model_type) scopes the registry lookup to the model’s own spec; None searches every spec.

The model’s own spec is consulted first, so per-model data always outranks the generic name and structural fallbacks. The fallbacks matter: a model with no spec still has to be recognised, which is what keeps export working for MoE architectures nobody has registered.

Parameters:
  • module (Module)

  • model_type (str | None)

Return type:

bool

list_all_possible(attr)#

List a spec attribute’s values across all registered specs, deduplicated in order.

E.g. list_all_possible("gate_up_pairs"). Looks the name up on ModelSpec and then on its sections, so callers name the field ("pqs_fuse_rules") without naming the section that holds it. Models whose declaring section is None contribute nothing.

The result is a global vocabulary: consumers match it against any model’s modules, so adding a value to one spec affects all models the consumer walks — prefer get_spec(model_type) / match_moe_block wherever the owning model is identifiable.

attr must name a tuple-valued attribute. Scalars (model_type) would be iterated character by character, so they are rejected rather than silently producing nonsense.

Parameters:

attr (str)

Return type:

tuple

match_class_names(module, names)#

Return True if any of names equals a class name in module’s MRO.

Case-insensitive exact-name comparison against every class in type(module).__mro__: dynamically generated quantized classes match through their base class, and exact names avoid substring false positives.

Parameters:

names (tuple[str, ...])

Return type:

bool

match_moe_block(module, model_type=None)#

Return the MoE layout layout for module, resolved by model type.

model_type (the root model.config.model_type) is a strict filter: only that model’s own spec is consulted, and an unregistered model type resolves to None even if the module’s class names coincide with another model’s. model_type=None searches all specs. A composite model whose MoE lives under a sub-model type registers the root type too (see gemma4/specs.py). Within the spec, block_names matched against the module’s MRO decides.

Parameters:
  • module (nn.Module)

  • model_type (str | None)

Return type:

MoESpec | None

match_moe_model(module, model_type=None)#

Return the spec whose MoE section matches module, or None.

Same scoping as match_moe_block, which returns only the layout; use this when a consumer also needs the model’s other sections, such as the export policy that says whether its experts may be grouped.

Parameters:
  • module (nn.Module)

  • model_type (str | None)

Return type:

ModelSpec | None

register(spec)#

Register a model spec and return it. One spec per model type.

Parameters:

spec (ModelSpec)

Return type:

ModelSpec