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.

A package may also hold modeling_ptq.py, the model-specific PTQ wrappers. It is not imported here; the HF quantization plugin imports it (see README.md).

Classes

`EPSpec`	Subsystem section: per-model data of the expert-parallel (DP=EP) PTQ path.
`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 EPSpec#

Bases: SpecSection

Subsystem section: per-model data of the expert-parallel (DP=EP) PTQ path.

Calibrating a large MoE data-parallel with its experts sharded needs three things from the model, and all three are facts about the architecture rather than about any run: which modules are the routed experts and how they shard, which transformers release first ships that information, and – for a VLM – which nested config carries it.

Most models need only the data below. A checkpoint whose modeling code it ships itself may need actual code run against it before it can be loaded expert-parallel (patching a remote-code class, decompressing packed experts); that is what prepare is for, and it is deliberately the exception rather than the interface.

__init__(*, base_model_ep_plan=None, requires_transformers=None, plan_config_attr=None, prepare=None, prepared_ep_plan=None, unsupported=None)#

Parameters:

base_model_ep_plan (dict[str, str] | None)
requires_transformers (str | None)
plan_config_attr (str | None)
prepare (Callable[..., Any] | None)
prepared_ep_plan (dict[str, str] | None)
unsupported (str | None)

Return type:

None

base_model_ep_plan: dict[str, str] | None = None#

The base_model_ep_plan transformers needs to shard this model’s experts, applied to the config when it does not already carry one.

Stock transformers ships a plan for some MoE models and not others; where it does, leave this None and let the model’s own plan win. Entries map a parameter glob to a parallel style, e.g. {"layers.*.mlp.experts.gate_up_proj": "grouped_gemm"}.

A router entry is not always wanted. DeepSeek-V3’s router is logits-only and its MoE block computes global top-k ids, so the experts already receive unmasked global ids – which is exactly what the all-to-all dispatch needs, and adding a gate entry would mask them.

plan_config_attr: str | None = None#

The nested config attribute carrying the expert-parallel plan, for a VLM whose MoE is its text sub-model ("text_config", "llm_config"). None means the root config.

The plan has to sit on whichever config describes the module tree the globs are written against, so a plan on the wrong config silently matches nothing.

prepare: Callable[..., Any] | None = None#

Escape hatch for a model that cannot be made expert-parallel by data alone.

Called as prepare(hf_config, ckpt_path=..., trust_remote_code=...) and returns the config to load. Reserved for checkpoints carrying their own modeling code, where the work is patching classes or installing on-the-fly weight decompression. Runs after the declarative fields above have been applied.

prepared_ep_plan: dict[str, str] | None = None#

The base_model_ep_plan prepare attaches, for a model whose plan only prepare supplies.

Stated, not applied: prepare still attaches it. A run that rules expert parallelism out skips prepare, which patches the model’s code, yet must still tell an expert-parallel MoE it would hold whole on every rank from a dense model. It finds the routed experts through this plan, in the model as the checkpoint’s own code builds it: prepare replaces the experts in place, so the plan names their containers there too.

requires_transformers: str | None = None#

Earliest transformers release that supports THIS model expert-parallel, when it is later than the framework floor.

Distinct from ModelSpec.min_transformers_version, which is about the modeling code matching the spec. A model can be fully modelled and still not be shardable: qwen3_moe’s own base_model_ep_plan lands in 5.7, while the shared EP framework is 5.6.

unsupported: str | None = None#

Why this model cannot be quantized expert-parallel yet, when its config would route it there.

Stock transformers ships a base_model_ep_plan for some MoEs the expert-parallel path cannot handle, and the plan alone makes a model expert-parallel capable. Naming the reason here refuses the run before any weight is loaded instead of failing deep in quantization.

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 into per-expert modules, 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).

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, ep_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)
ep_spec (EPSpec | None)

Return type:

None

ep_spec: EPSpec | None = None#: Per-model data of the expert-parallel PTQ path, or None when the model is dense or needs nothing beyond what stock transformers already ships.

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)#

Parameters:: spec (ModelSpec)
Return type:: ModelSpec