diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md new file mode 100644 index 00000000000..9ae044b9e2a --- /dev/null +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -0,0 +1,253 @@ +# Export: Model-Specific Logic Refactor — Action Plan + +**Goal:** Extract model-specific logic out of `modelopt/torch/export` into a new +modeling library (`modelopt/torch/export/modeling/`), organized by model family, +leaving the export code generic and model-agnostic. + +**Scope (this effort):** HF / TRT-LLM export path only. The Megatron path +(`plugins/mcore_*`) already follows the target registry pattern and is kept as the +reference north star, not refactored here. + +--- + +## 1. Inventory of Model-Specific Logic + +Counts are grep-derived and reproducible. "Branches" = decision points +(`decoder_type ==/in`, class-name / `module_match_name_list` string matches). +"Line footprint" = lines referencing a model-family name literal (approximate +measure of model-specific contamination). + +### 1.1 HF / TRT-LLM path (refactor target) + +| Category | Branches | Line footprint | Main location | +|---|---|---|---| +| Model ID table `MODEL_NAME_TO_TYPE` | 41 entries | ~50 | `model_utils.py` | +| `decoder_type ==/in` imperative branches | 51 | — | layer_utils 25, model_config_export 15, tensorrt_llm_utils 11 | +| Class-name / arch string branches (`type().__name__`, `module_match_name_list`) | ~74 | — | layer_utils 42, unified_export_hf 20, quant_utils 4, diffusers_utils 5, model_utils 2, transformer_engine 1 | +| Config field map `HF_CONFIG_MAP` | 36 entries | ~36 | `hf_config_map.py` | +| AWQ fusion table `PQS_FUSE_MODULE_MAPPING` | 2 groups | ~12 | `quant_utils.py:1141` | +| Speculative-decoding templates | 2 configs + 3 Exporter classes | ~150 | `plugins/hf_spec_configs.py`, `plugins/hf_spec_export.py` | + +**Totals:** ≈ 125 imperative branch points + ≈ 79 data-table entries. +Concentrated in 4 files — `layer_utils.py` (67, the heaviest), `model_config_export.py`, +`tensorrt_llm_utils.py`, `unified_export_hf.py` — plus two tables in +`model_utils.py` and `hf_config_map.py`. + +### 1.2 Breakdown by functional category (HF / TRT-LLM path) + +| Functional category | ~Branches | Representative location | +|---|---|---| +| MoE expert naming / router / expert stacking | ~20 | `get_expert_linear_names`, `build_moe_config` | +| Attention variants (MQA/GQA/MLA/cross-attn/clip_qkv) | ~15 | `build_decoder_config`, `build_attention_config` | +| LayerNorm fixups (weight+1, etc.) | ~3 | `build_layernorm_config` | +| Activation overrides per family | ~4 | `build_mlp_config` | +| MLP gate/fc naming swap, fused chunk order | ~8 | `build_mlp_config` | +| Positional / rope / config field translation | ~36 + a few | `HF_CONFIG_MAP`, Phi3 longrope, etc. | +| Embedding scaling / share / tied-weight dedup | ~6 | `model_config_export`, `model_utils` | +| Encoder-decoder routing (T5/BART/Whisper) | ~12 | `model_config_export`, `tensorrt_llm_utils` | +| Multimodal VLM language-tower extraction | ~5 | `model_utils`, `unified_export_hf` | +| Quant-format specific (AWQ fusion / NVFP4 3D experts / GPT-OSS) | ~8 | `quant_utils`, `unified_export_hf` | +| Speculative-decoding export | ~5 classes/templates | `plugins/hf_spec_*` | + +### 1.3 Megatron path (out of scope — reference only) + +Already a declarative registry; not refactored here. + +| Item | Amount | Location | +|---|---|---| +| Per-family mapping files (≈100% model-specific) | 6 files / 771 lines | `plugins/mcore_{llama,qwen,qwen3vl,deepseek,gptoss,nemotron}.py` | +| Family registry | 14 export + 10 import entries | `plugins/mcore_common.py:44` | +| `MCORE_CONFIG_MAP` field map | 3 entries | `mcore_config_map.py` | +| Arch branches in export driver (MLA / VLM / packed-expert) | ~3 | `unified_export_megatron.py` | + +### 1.4 Key takeaways + +- The real backlog to absorb is the HF/TRT-LLM path: **≈125 scattered imperative + branches + ≈79 data-table entries**. +- ~80% of the imperative branches live in a single file: `layer_utils.py` (67). +- ~79 entries are **already data tables** (`MODEL_NAME_TO_TYPE`, `HF_CONFIG_MAP`, + `PQS_FUSE_MODULE_MAPPING`) — moving them into the modeling lib is near-mechanical. +- Of the ~125 imperative branches, only a minority are true control-flow + (attention variants, enc-dec); most are "look up a value by model". + +--- + +## 2. Rationale & Trade-offs + +This is a `operation × model` matrix. Organize by operation (today) and changing an +operation is cheap but adding a model is shotgun surgery; organize by model and the +reverse. The right axis is the one that changes most. + +**Evidence — export files touched per recent model-support PR:** + +| PR | export files touched | +|---|---| +| NemotronH non-gated MoE (#1756) | 5 (layer_utils, moe_utils, quant_utils, unified_export_hf, vllm_fakequant) | +| DiffusionGemma enc-dec (#1707) | 4 (model_utils, moe_utils, quant_utils, unified_export_hf) | +| Gemma4 MoE (#1219) | 1 (layer_utils) | +| Llama4 MoE fix (#1744) | 1 (unified_export_hf) | + +Adding a model is frequent and often spans 4–5 files — the cost a by-model split removes. + +**Scope discipline (the red line):** move per-model *deltas*, not generic algorithms. + +| Move into modeling lib (by model) | Keep in functional files (by operation) | +|---|---| +| Data: expert names, norm+1, activation, embed scale, share_embed, AWQ fusion table, config field maps | Generic algorithms: scale compute, weight packing, TP/PP split | +| Most "look up a value by model" branches | True control-flow oddballs: enc-dec routing, MLA | + +**Net:** function-cohesion for generic code is correct and stays; only the per-model +deltas leaking into those files get extracted. Target = generic engine (by operation) +- model-descriptor registry (by model), mirroring the Megatron `plugins/mcore_*` pattern. + +## 3. Target Architecture + +Two layers, split along the §2 red line: + +- **Generic engine** (the existing files, organized *by operation*): `layer_utils`, + `quant_utils`, `moe_utils`, `model_config_export`, … — keep all algorithms here. +- **Modeling library** (`modeling/`, organized *by model*): per-model **data only**, + read by the engine through registry lookups. + +### 3.1 Layout + +```text +modeling/ + base.py # ModelSpec dataclass — the per-model contract + registry.py # register() + lookups; returns None when unmatched + __init__.py # re-exports; importing it registers all families + families/ + __init__.py # imports each family module (import == registration) + qwen.py deepseek.py gemma.py mixtral.py dbrx.py gptoss.py nemotron.py ... +``` + +`modeling/` depends only on `torch.nn` + stdlib. It must **never** import from the +parent export package (no `from ..layer_utils import ...`), so it stays at the bottom +of the dependency graph and cannot reintroduce the `unified_export_hf ↔ plugins` cycle. + +### 3.2 `ModelSpec` — the contract + +A per-model-family record with three kinds of members: + +- **Matching keys** (how a spec is resolved): `architectures`, `decoder_types`, + `moe_block_names` (case-insensitive substring), `mlp_block_names` (exact). +- **Per-model data** (preferred whenever a value suffices): `expert_linear_names`, + `has_iterable_experts`, `forced_activation`, `force_share_embedding_table`, + `mlp_keyword_roles`, `pqs_fuse_rules`. +- **Per-model hooks** (`hooks: ModelHooks`, default no-op) — used only where behavior + genuinely diverges and no value can express it (see §3.6). + +Principle: the engine owns the common algorithm; a family overrides only **named, +fine-grained seams** via data or hooks with a default fallback — it never forks a whole +function. Prefer a data field; reach for a hook only when the divergence is behavioral. + +### 3.3 Lookup conventions + +Resolvers in `registry.py`, each returning `ModelSpec | None` (plus `iter_pqs_fuse_rules` +which aggregates a rule list across all specs): + +- `match_by_architecture(arch)` — keyed on `model.config.architectures[0]`. +- `match_by_decoder_type(decoder_type)` — keyed on the ModelOpt `decoder_type` string. +- `match_moe_block(module)` / `match_mlp_block(module)` — keyed on a sub-module class name. + +Specs are matched in registration order; families keep their match keys mutually +exclusive, so order is never load-bearing. A family should ideally be a **single spec** +carrying all of its keys, data, and hooks (consolidate split specs as hooks are added). + +### 3.4 Call-site integration pattern (fallback-first) + +Every migrated call site follows the same shape — look up, use if present, else fall +through to the **unchanged** legacy code: + +```python +spec = match_moe_block(module) # or match_by_architecture(arch) +if spec is not None and spec. is not None: + return spec. +# ... legacy branch preserved verbatim as fallback ... +``` + +This is what makes migration incremental: a family not yet in `modeling/` returns +`None` and behaves exactly as before. Once every family for a category is migrated, the +legacy branch is dead but kept until the category is fully retired. + +### 3.5 Conventions for contributors + +- **Add a model:** create/extend `families/.py`, register a `ModelSpec`, done — + no engine edits. +- **Add a data category:** add a field to `ModelSpec`, populate it in the families, + convert the one engine call site to the fallback-first pattern, add an equivalence test. +- **Add a behavioral seam:** add a default no-op method to `ModelHooks`, call it at the + engine seam with fallback, override it in the families that need it (see §3.6). +- Keep the common algorithm in the engine; families supply only per-model values or + override specific seams. Do not fork whole functions into `modeling/`. + +### 3.6 Per-model hooks + +Some divergence is behavioral (which sub-module to read, which config slot to assign, +how to traverse a family's module tree) and cannot be a value. These are handled by +`ModelHooks`, not by forking the engine. + +Two simplifying decisions keep this small: + +1. **Data first** — a hook is used only when no value can express the divergence. +2. **Engine builds, the hook places** — the engine still classifies a sub-module and + builds its config object; the hook only routes that object into a family-specific + slot. So most hooks need no builders and the injected context stays tiny. + +Pieces: + +- **`ModelHooks`** (`modeling/hooks.py`) — a base class with one default **no-op** method + per seam. A family subclasses it and overrides only the seams it needs; + `ModelSpec.hooks` defaults to a shared no-op instance. +- **`ExportContext`** (`modeling/context.py`) — a small struct carrying export metadata + plus the few builder callables the heavier seams need (only `build_moe` and Gemma3 + q/k-norm extraction). Hooks call `ctx.build_*` instead of importing `layer_utils`, so + `modeling/` stays at the bottom of the dependency graph (no cycle). +- **Engine seams** call the hook and fall back when it declines (no-op sentinel): + + ```python + if is_moe(layer): + config.mlp = hooks.build_moe(layer, ctx) or build_moe_config(layer, decoder_type) + else: + built = build_layernorm_config(child) # or build_attention_config(...) + if not hooks.place_submodule(name, built, config, ctx): + ... # default placement + ``` + +The seam set is **finite and known** (the library was audited end to end): + +| Seam | Covers | +|---|---| +| `unwrap_decoder_layer(module, ctx)` | module-tree shape differences (DBRX, ExaOne, Deci) | +| `place_submodule(name, built, layer_cfg, ctx)` | family-specific config slots — Gemma2/3 norms, MLlama/enc-dec self↔cross attention, Gemma3 q/k-norm | +| `build_moe(module, ctx)` | per-family router + experts + shared expert | + +Everything else stays data: `head_is_first_dim` (bool) and the TRT-LLM target-config +extras (a per-family field dict) are values, not hooks. + +## 4. Migration Plan & Priority + +Priority is set by two tests: **(a) how pure-data it is** (purer → sooner) and +**(b) how central it is to add-a-model shotgun surgery** (the §2 evidence points at +MoE). Each step keeps a fallback to the legacy path, so migration is incremental and +never breaks un-migrated models. + +| Step | What | Why this order | Risk | +|---|---|---|---| +| **P0** ✅ | Registry skeleton: `modeling/` + `ModelSpec` + lookups (`match_moe_block` / `match_mlp_block` / `match_by_decoder_type` / `match_by_architecture`, return `None` → fallback). Do **not** touch `MODEL_NAME_TO_TYPE`. | Everything keys off it; purely additive, no behavior change. | very low | +| **P1** (pilot) ✅ | MoE expert naming: `get_expert_linear_names` if-chain → `spec.expert_linear_names`. | Purest data + the #1 shotgun-surgery driver. Proves the mechanism end-to-end with near-zero blast radius. | very low | +| **P2** ✅ | The duplicate expert-naming table in `get_experts_list` → `spec.expert_linear_names` + `spec.has_iterable_experts` (the grouped-export capability flag that reproduces the legacy "supported 4 families, else `NotImplementedError`"). | Removes the second copy of expert-naming data → "add a MoE model" stops touching this site. | low | +| **P3** (partial) | Non-MoE per-model data. Done: activation override (`forced_activation` ✅), shared-embedding (`force_share_embedding_table` ✅), MLP keyword roles (`mlp_keyword_roles` ✅ — Arctic/InternLM2 remap, MPT/Phi3Small/TLGv4/Nemotron up_proj→fc). Deferred: norm+1 (mixes the generic `LayerNorm1P` which stays in the engine; needs a case-sensitive norm-class resolver) and embed √scale (Gemma1-only + `is_tensorrt_llm_0_8_or_9()` version gate) — low value. | Mechanical bool/str/dict moves; clears scattered `layer_utils` / `model_config_export` branches. | low | +| **P4** (partial) | `PQS_FUSE_MODULE_MAPPING` (AWQ fusion) → `spec.pqs_fuse_rules`, aggregated via `iter_pqs_fuse_rules` (✅, llama/qwen families). `HF_CONFIG_MAP` deliberately **kept**: it is a shared *alias* table applied generically to every config (`for keys, name in HF_CONFIG_MAP`), not a per-model branch — splitting it by family would be artificial and more fragile. | PQS rules are genuinely per-model (keyed by module class); the alias table is not. | low | +| **P5** (hooks) | Behavioral per-model logic via `ModelHooks` + `ExportContext` (§3.6), 3 seams. **P5a** (pilot): hook machinery + `place_submodule` for decoder slots (Gemma2/3 norms, MLlama self/cross, Gemma3 q/k-norm). **P5b**: `unwrap_decoder_layer` (DBRX/ExaOne/Deci) + the `head_is_first_dim` data flag. **P5c**: `build_moe` per family (Llama/Phi3, DBRX, DeepSeek, Qwen). **P5d** (later, mostly data): TRT-LLM target-config extras as a per-family field dict; enc-dec routing and VLM extraction remain separate paths. | Moves the structural `build_decoder_config` / `build_moe_config` branches that data fields cannot express; biggest remaining per-model cluster. | medium (hooks run real building logic — equivalence-test each seam) | +| **OUT** | Shared/generic tables that are not per-model (`HF_CONFIG_MAP`); generic algorithms (`_GATE_UP_PAIRS`, amax fallback); the ChatGLM/Phi3 fused gate/fc chunk-swap (a fused-weight reshape, not a per-model value); speculative-decoding export (already its own modular subsystem under `plugins/hf_spec_*`). | Not per-model branches, pure algorithm, or already modular — they stay where they are. | — | + +**Guardrails:** fallback-first (un-migrated models keep the old path); one seam/category +per PR with an export-test equivalence check; the engine keeps the common algorithm — +families supply only per-model values or override named seams, never fork whole functions. + +> **P2 note:** an earlier draft scoped P2 as "expert amax fallback / gate-up amax sync." +> On inspection `set_expert_quantizer_amax` and `sync_moe_gate_up_amax` are generic +> *algorithms* (the gate/up pairs in `_GATE_UP_PAIRS` are a flat, non-per-model list), so +> by the red line they stay in the engine. The only per-model *data* left in this area was +> the duplicate naming table in `get_experts_list`, which is what P2 actually migrated. diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index d1bd0d7121f..4078af07910 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -60,6 +60,13 @@ RgLruConfig, ) from .model_config_utils import pad_weights +from .modeling import ( + NULL_HOOKS, + ExportContext, + match_by_decoder_type, + match_mlp_block, + match_moe_block, +) from .postprocess import view_as_float8_e4m3fn_if_needed, view_as_uint8_if_needed from .quant_utils import ( get_activation_scaling_factor, @@ -93,25 +100,14 @@ def get_experts_list( """ experts_list = [] - # Define linear layer names for different model types - if "mixtralforcausallm" in model_type: - linear_names = ["w1", "w2", "w3"] - elif any( - qwen_variant in model_type - for qwen_variant in [ - "qwenmoeforcausallm", - "qwen2moeforcausallm", - "qwen3moeforcausallm", - "qwen3nextforcausallm", - ] - ): - linear_names = ["gate_proj", "down_proj", "up_proj"] - elif "nemotronhforcausallm" in model_type: - linear_names = ["up_proj", "down_proj"] - elif "gemma4" in model_type: - linear_names = ["gate_proj", "down_proj", "up_proj"] - else: + # Expert linear names live in modeling/families/*. The grouped export path only + # supports families whose experts are iterable per-expert sub-modules (Mixtral, + # Qwen MoE, NemotronH, Gemma4); stacked/fused layouts (DBRX, GptOss, ...) raise + # NotImplementedError here and are handled by other paths. + spec = match_moe_block(module) + if spec is None or not spec.has_iterable_experts or spec.expert_linear_names is None: raise NotImplementedError(f" {model_type} not supported") + linear_names = list(spec.expert_linear_names) # Common logic for all supported model types experts_list.extend( @@ -770,26 +766,18 @@ def _split_gate_from_fc(decoder_type, module, fc_name, fc_layer): "c_fc_1", # exaone } - # Arctic (llama-based MoE, decoder_type is "llama") has MLP keyword conflicts with Qwen - # Arctic's residual MLP use w1 for fc, w2 for proj, w3 for gate - if type(module).__name__ in ["ArcticMLP", "InternLM2MLP"]: - fc_keywords.discard("w2") - gate_keywords.discard("w1") - fc_keywords.add("w1") - proj_keywords.add("w2") - gate_keywords.add("w3") - - if decoder_type == "mpt": - fc_keywords.add("up_proj") - gate_keywords.discard("up_proj") - - if type(module).__name__ in [ - "TLGv4MLP", - "Phi3SmallMLP", - "NemotronMLP", - ]: # for TLGv4ForCausalLM - fc_keywords.add("up_proj") - gate_keywords.discard("up_proj") + # Per-family MLP keyword role overrides live in modeling/families/* (e.g. Arctic/ + # InternLM2 w1/w2/w3 remap; MPT/Phi3Small/TLGv4/Nemotron up_proj→fc). Resolve by the + # MLP module class and by decoder_type; apply on top of the shared keyword sets. + _role_sets = {"fc": fc_keywords, "gate": gate_keywords, "proj": proj_keywords} + for _spec in (match_mlp_block(module), match_by_decoder_type(decoder_type)): + if _spec is None or not _spec.mlp_keyword_roles: + continue + for _kw, _role in _spec.mlp_keyword_roles.items(): + fc_keywords.discard(_kw) + gate_keywords.discard(_kw) + proj_keywords.discard(_kw) + _role_sets[_role].add(_kw) fc_linear: nn.Module = None gate_linear: nn.Module = None @@ -865,11 +853,11 @@ def _split_gate_from_fc(decoder_type, module, fc_name, fc_layer): assert config.proj is not None and config.fc is not None, "proj or fc can not be found" - # Override hidden_act based on decoder_type - if decoder_type in ["bloom", "glm"]: - hidden_act = "gelu" - if decoder_type == "phi3": - hidden_act = "swiglu" + # Per-family activation override lives in modeling/families/* (e.g. bloom/glm → gelu, + # phi3 → swiglu). Unmatched families fall through to activation detection below. + _spec = match_by_decoder_type(decoder_type) + if _spec is not None and _spec.forced_activation is not None: + hidden_act = _spec.forced_activation if hidden_act is None: if hasattr(module, "activation"): @@ -992,6 +980,12 @@ def module_match_name_list(module, name_list): if hasattr(module.experts, f"{first_proj_attr}_weight_quantizers"): return [first_proj_attr, "down_proj"] + # Resolve expert names from the model-family registry; fall back to the mapping + # below when no family spec matches the MoE block. + spec = match_moe_block(module) + if spec is not None and spec.expert_linear_names is not None: + return list(spec.expert_linear_names) + if module_match_name_list( module, [ @@ -1581,6 +1575,11 @@ def build_decoder_config( """Builds the full decoder config from the module.""" quantization = get_quantization_format(module) config = DecoderLayerConfig(decoder_type=decoder_type) + + # Per-model behavioral hooks (default no-op) and the context passed into them. + _spec = match_by_decoder_type(decoder_type) + hooks = _spec.hooks if (_spec is not None and _spec.hooks is not None) else NULL_HOOKS + hook_ctx = ExportContext(build_layernorm_config=build_layernorm_config) # Supporting different attention layer config in MCoreGPTModel. If per layer config # exists, override the global config. if hasattr(module, "self_attention") and hasattr(module.self_attention, "config"): @@ -1641,26 +1640,17 @@ def __init__(self): # and residual_layernorm could be after post_layernorm if is_layernorm(layer): layernorm_config = build_layernorm_config(layer) + # Family-specific layernorm slots (e.g. Gemma2/3) are placed by the hook. + if hooks.place_submodule(name, layer, layernorm_config, config, hook_ctx): + pass # Special attributes naming for Encoder-Decoder models - if model_type_is_enc_dec(decoder_type): + elif model_type_is_enc_dec(decoder_type): _update_encoder_decoder_layernorm_config( model_metadata_config, config, layernorm_config ) # For all decoder only models elif name in ["ln_mlp"]: config.mlp_layernorm = layernorm_config - elif ( - config.decoder_type in ["gemma2", "gemma3"] - ) and "post_attention_layernorm" in name: - config.post_layernorm = layernorm_config - elif ( - config.decoder_type in ["gemma2", "gemma3"] - ) and "pre_feedforward_layernorm" in name: - config.pre_feedforward_layernorm = layernorm_config - elif ( - config.decoder_type in ["gemma2", "gemma3"] - ) and "post_feedforward_layernorm" in name: - config.post_feedforward_layernorm = layernorm_config elif config.input_layernorm is None: config.input_layernorm = layernorm_config elif config.post_layernorm is None: @@ -1681,8 +1671,12 @@ def __init__(self): attention_config = build_attention_config( layer, model_metadata_config, config, tp_size=tp_size ) + # Family-specific attention placement (self/cross routing, extra norms) + # is handled by the hook; otherwise use the default placement below. + if hooks.place_submodule(name, layer, attention_config, config, hook_ctx): + pass # For decoder of Encoder-Decoder model with self, cross attention - if ( + elif ( model_type_is_enc_dec(decoder_type) and model_metadata_config["enc_dec"] == "dec" ): @@ -1691,19 +1685,8 @@ def __init__(self): config.self_attention = attention_config else: config.cross_attention = attention_config - elif decoder_type == "mllama": - if "cross" in type(layer).__name__.lower(): - config.cross_attention = attention_config - else: - config.self_attention = attention_config else: config.attention = attention_config - if decoder_type == "gemma3": - # Gemma3 has q_norm and k_norm within self_attn - q_layernorm_config = build_layernorm_config(layer.q_norm) - k_layernorm_config = build_layernorm_config(layer.k_norm) - config.attention.q_layernorm = q_layernorm_config - config.attention.k_layernorm = k_layernorm_config elif is_moe(layer): if quantization not in [QUANTIZATION_NONE, QUANTIZATION_FP8]: diff --git a/modelopt/torch/export/model_config_export.py b/modelopt/torch/export/model_config_export.py index ae92e2776f5..fb3d46bc14e 100644 --- a/modelopt/torch/export/model_config_export.py +++ b/modelopt/torch/export/model_config_export.py @@ -60,6 +60,7 @@ pack_linear_weights, split_config_and_weights, ) +from .modeling import match_by_decoder_type from .postprocess import ( check_weight_shape_valid, pad_embedding_lm_head, @@ -312,8 +313,10 @@ def torch_to_tensorrt_llm_checkpoint( # the model head weight is None so we just skip processing config.share_embedding_table = True continue - # TRT LLM forces the embedding table to be shared for the following models. - force_share_embedding_table = decoder_type in ["gemma", "gemma2", "gemma3"] + # TRT LLM forces the embedding table to be shared for some families + # (see modeling/families/*); the equality check below still gates it. + _spec = match_by_decoder_type(decoder_type) + force_share_embedding_table = bool(_spec and _spec.force_share_embedding_table) if force_share_embedding_table and torch.equal( module.weight, config.vocab_embedding.weight ): diff --git a/modelopt/torch/export/modeling/__init__.py b/modelopt/torch/export/modeling/__init__.py new file mode 100644 index 00000000000..8c0172dfaf9 --- /dev/null +++ b/modelopt/torch/export/modeling/__init__.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Model-family export descriptors. + +Holds per-model export data organized by model family. The export code resolves a +``ModelSpec`` via the registry lookups and reads its fields; an unmatched lookup +returns ``None`` so callers fall back to their default behavior. +""" + +from . import families +from .base import ModelSpec +from .context import ExportContext +from .hooks import NULL_HOOKS, ModelHooks +from .registry import ( + iter_pqs_fuse_rules, + match_by_architecture, + match_by_decoder_type, + match_mlp_block, + match_moe_block, + register, +) + +__all__ = [ + "NULL_HOOKS", + "ExportContext", + "ModelHooks", + "ModelSpec", + "iter_pqs_fuse_rules", + "match_by_architecture", + "match_by_decoder_type", + "match_mlp_block", + "match_moe_block", + "register", +] diff --git a/modelopt/torch/export/modeling/base.py b/modelopt/torch/export/modeling/base.py new file mode 100644 index 00000000000..7c9f7ad0a15 --- /dev/null +++ b/modelopt/torch/export/modeling/base.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-model-family export descriptor. + +A ``ModelSpec`` declares how one model family differs from the generic export path, +so the export code can read these values instead of branching on model names. Each +spec holds per-model data only, not export logic. +""" + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .hooks import ModelHooks + +__all__ = ["ModelSpec"] + + +@dataclass +class ModelSpec: + """Per-model-family export data. + + A spec is resolved from a model (or one of its sub-modules) via any of its matching + keys: + architectures: exact HuggingFace architecture class names (e.g. "Qwen3MoeForCausalLM"). + decoder_types: ModelOpt ``decoder_type`` strings (e.g. "bloom", "phi3"). Exact match. + moe_block_names: MoE block class-name substrings, matched case-insensitively + (e.g. "Qwen3MoeSparseMoeBlock"). + mlp_block_names: MLP module class names, matched by exact equality + (e.g. "ArcticMLP", "Phi3SmallMLP"). + + Per-model fields: + expert_linear_names: expert linear projection names for this family's MoE block, + e.g. ("gate_proj", "down_proj", "up_proj"). ``None`` if not applicable. + has_iterable_experts: True when experts are stored as per-expert iterable + sub-modules (Mixtral, Qwen MoE, NemotronH, Gemma4) and can be grouped by + ``get_experts_list``. False for stacked or fused layouts (DBRX, GptOss). + forced_activation: activation that overrides MLP activation detection + (e.g. Bloom/GLM → "gelu", Phi3 → "swiglu"). ``None`` for no override. + force_share_embedding_table: True for families that share the embedding/output + table (Gemma/Gemma2/Gemma3); still gated by a weight-equality check. + mlp_keyword_roles: overrides mapping a child-module name to its MLP role, e.g. + {"up_proj": "fc"} (MPT/Phi3Small) or {"w1": "fc", "w2": "proj", "w3": "gate"} + (Arctic/InternLM2). Each keyword is removed from the default role sets and + added to its target role. ``None`` for no override. + pqs_fuse_rules: 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). + hooks: optional :class:`ModelHooks` for behavioral seams that no value can express + (custom config-slot placement, module-tree unwrap, MoE building). ``None`` uses + the engine defaults. + """ + + name: str + + # Matching keys. + architectures: tuple[str, ...] = () + decoder_types: tuple[str, ...] = () + moe_block_names: tuple[str, ...] = () + mlp_block_names: tuple[str, ...] = () + + # MoE expert layout. + expert_linear_names: tuple[str, ...] | None = None + has_iterable_experts: bool = False + + # MLP and activation. + forced_activation: str | None = None + mlp_keyword_roles: dict[str, str] | None = None + + # Embedding. + force_share_embedding_table: bool = False + + # AWQ pre_quant_scale fusion. + pqs_fuse_rules: tuple[tuple[tuple[str, ...], str, str], ...] = () + + # Behavioral overrides for structural export seams (None = use the engine default). + hooks: "ModelHooks | None" = None + + # Free-form extension slot for future per-model fields. + _extra: dict = field(default_factory=dict, repr=False) diff --git a/modelopt/torch/export/modeling/context.py b/modelopt/torch/export/modeling/context.py new file mode 100644 index 00000000000..10b1e0f414e --- /dev/null +++ b/modelopt/torch/export/modeling/context.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Context passed from the export engine into per-model hooks. + +The engine fills this with the builder callables a hook may need and passes it in, so +hooks never import the engine modules directly (keeping ``modeling/`` free of cycles). +""" + +from collections.abc import Callable +from dataclasses import dataclass + +__all__ = ["ExportContext"] + + +@dataclass +class ExportContext: + """Builders and metadata available to ``ModelHooks`` methods.""" + + build_layernorm_config: Callable diff --git a/modelopt/torch/export/modeling/families/__init__.py b/modelopt/torch/export/modeling/families/__init__.py new file mode 100644 index 00000000000..2c911181868 --- /dev/null +++ b/modelopt/torch/export/modeling/families/__init__.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-model-family export specs. Importing each module registers its specs.""" + +from . import ( + arctic, + bloom, + dbrx, + deepseek, + gemma, + glm, + gptoss, + internlm, + llama, + mixtral, + mllama, + mpt, + nemotron, + phi, + qwen, +) diff --git a/modelopt/torch/export/modeling/families/arctic.py b/modelopt/torch/export/modeling/families/arctic.py new file mode 100644 index 00000000000..cd66644858b --- /dev/null +++ b/modelopt/torch/export/modeling/families/arctic.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Arctic family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +# Arctic (llama-based MoE) residual MLP uses w1 for fc, w2 for proj, w3 for gate, +# which conflicts with Qwen's w1/w2 defaults; remap the roles for this MLP. +register( + ModelSpec( + name="arctic", + mlp_block_names=("ArcticMLP",), + mlp_keyword_roles={"w1": "fc", "w2": "proj", "w3": "gate"}, + ) +) diff --git a/modelopt/torch/export/modeling/families/bloom.py b/modelopt/torch/export/modeling/families/bloom.py new file mode 100644 index 00000000000..908a1da5ea0 --- /dev/null +++ b/modelopt/torch/export/modeling/families/bloom.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Bloom family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="bloom", + decoder_types=("bloom",), + forced_activation="gelu", + ) +) diff --git a/modelopt/torch/export/modeling/families/dbrx.py b/modelopt/torch/export/modeling/families/dbrx.py new file mode 100644 index 00000000000..5d53ddb753a --- /dev/null +++ b/modelopt/torch/export/modeling/families/dbrx.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DBRX family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="dbrx", + moe_block_names=("DBRXMoeSparseMoeBlock",), + expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), + ) +) diff --git a/modelopt/torch/export/modeling/families/deepseek.py b/modelopt/torch/export/modeling/families/deepseek.py new file mode 100644 index 00000000000..fa0a0589f60 --- /dev/null +++ b/modelopt/torch/export/modeling/families/deepseek.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DeepSeek family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="deepseek", + moe_block_names=("DeepseekMoE",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + ) +) diff --git a/modelopt/torch/export/modeling/families/gemma.py b/modelopt/torch/export/modeling/families/gemma.py new file mode 100644 index 00000000000..171eb75f617 --- /dev/null +++ b/modelopt/torch/export/modeling/families/gemma.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Gemma family export specs.""" + +from ...model_config import AttentionConfig +from ..base import ModelSpec +from ..hooks import ModelHooks +from ..registry import register + + +class GemmaHooks(ModelHooks): + """Gemma2/3 carry extra layernorms; Gemma3 also has per-head q/k norms.""" + + def place_submodule(self, name, module, built, layer_config, ctx): + """Place Gemma2/3 extra layernorms and Gemma3 attention q/k norms.""" + if layer_config.decoder_type in ("gemma2", "gemma3"): + if "post_attention_layernorm" in name: + layer_config.post_layernorm = built + return True + if "pre_feedforward_layernorm" in name: + layer_config.pre_feedforward_layernorm = built + return True + if "post_feedforward_layernorm" in name: + layer_config.post_feedforward_layernorm = built + return True + if layer_config.decoder_type == "gemma3" and isinstance(built, AttentionConfig): + layer_config.attention = built + built.q_layernorm = ctx.build_layernorm_config(module.q_norm) + built.k_layernorm = ctx.build_layernorm_config(module.k_norm) + return True + return False + + +register( + ModelSpec( + name="gemma4_moe", + # Gemma4 MoE experts are unfused into per-expert nn.Linear layers. + moe_block_names=("Gemma4TextDecoderLayer",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + has_iterable_experts=True, + ) +) + +# Dense Gemma 1/2/3: shared embedding/output table, plus Gemma2/3 layernorm layout. +register( + ModelSpec( + name="gemma", + decoder_types=("gemma", "gemma2", "gemma3"), + force_share_embedding_table=True, + hooks=GemmaHooks(), + ) +) diff --git a/modelopt/torch/export/modeling/families/glm.py b/modelopt/torch/export/modeling/families/glm.py new file mode 100644 index 00000000000..0ab3cff80d8 --- /dev/null +++ b/modelopt/torch/export/modeling/families/glm.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GLM family export specs. + +Note: only ``decoder_type == "glm"`` forces gelu; "chatglm" deliberately does not (it +relies on activation detection), matching the legacy behavior. +""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="glm", + decoder_types=("glm",), + forced_activation="gelu", + ) +) diff --git a/modelopt/torch/export/modeling/families/gptoss.py b/modelopt/torch/export/modeling/families/gptoss.py new file mode 100644 index 00000000000..1c1cc342e1f --- /dev/null +++ b/modelopt/torch/export/modeling/families/gptoss.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPT-OSS family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="gptoss", + # GPT-OSS fuses gate and up into a single gate_up_proj. + moe_block_names=("GptOssMoE",), + expert_linear_names=("gate_up_proj", "down_proj"), + ) +) diff --git a/modelopt/torch/export/modeling/families/internlm.py b/modelopt/torch/export/modeling/families/internlm.py new file mode 100644 index 00000000000..2aca5631564 --- /dev/null +++ b/modelopt/torch/export/modeling/families/internlm.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""InternLM family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +# InternLM2 MLP shares Arctic's w1/w2/w3 role layout (w1=fc, w2=proj, w3=gate). +register( + ModelSpec( + name="internlm2", + mlp_block_names=("InternLM2MLP",), + mlp_keyword_roles={"w1": "fc", "w2": "proj", "w3": "gate"}, + ) +) diff --git a/modelopt/torch/export/modeling/families/llama.py b/modelopt/torch/export/modeling/families/llama.py new file mode 100644 index 00000000000..3f83e7559f9 --- /dev/null +++ b/modelopt/torch/export/modeling/families/llama.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Llama family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="llama", + # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. + pqs_fuse_rules=( + (("LlamaAttention",), "v_proj", "o_proj"), + (("LlamaMLP",), "up_proj", "down_proj"), + ), + ) +) diff --git a/modelopt/torch/export/modeling/families/mixtral.py b/modelopt/torch/export/modeling/families/mixtral.py new file mode 100644 index 00000000000..ade34693a2b --- /dev/null +++ b/modelopt/torch/export/modeling/families/mixtral.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Mixtral family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +# Mixtral with iterable experts uses w1/w2/w3. Fused experts (transformers 5.0+) are +# detected from their per-expert quantizer attributes and need no naming override here. +register( + ModelSpec( + name="mixtral", + moe_block_names=("MixtralSparseMoeBlock",), + expert_linear_names=("w1", "w2", "w3"), + has_iterable_experts=True, + ) +) + +# Older transformers naming for Mixtral. +register( + ModelSpec( + name="mixtral_mcore", + moe_block_names=("MixtralMoeSparseMoeBlock",), + expert_linear_names=("linear_fc1", "linear_fc2"), + ) +) diff --git a/modelopt/torch/export/modeling/families/mllama.py b/modelopt/torch/export/modeling/families/mllama.py new file mode 100644 index 00000000000..1d8ed0012a7 --- /dev/null +++ b/modelopt/torch/export/modeling/families/mllama.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MLlama family export specs.""" + +from ...model_config import AttentionConfig +from ..base import ModelSpec +from ..hooks import ModelHooks +from ..registry import register + + +class MLlamaHooks(ModelHooks): + """MLlama interleaves self- and cross-attention layers; route by class name.""" + + def place_submodule(self, name, module, built, layer_config, ctx): + """Route attention into the self- or cross-attention slot by module class name.""" + if isinstance(built, AttentionConfig): + if "cross" in type(module).__name__.lower(): + layer_config.cross_attention = built + else: + layer_config.self_attention = built + return True + return False + + +register( + ModelSpec( + name="mllama", + decoder_types=("mllama",), + hooks=MLlamaHooks(), + ) +) diff --git a/modelopt/torch/export/modeling/families/mpt.py b/modelopt/torch/export/modeling/families/mpt.py new file mode 100644 index 00000000000..0b450838336 --- /dev/null +++ b/modelopt/torch/export/modeling/families/mpt.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MPT family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +# MPT treats up_proj as the fc projection (not gate). +register( + ModelSpec( + name="mpt", + decoder_types=("mpt",), + mlp_keyword_roles={"up_proj": "fc"}, + ) +) diff --git a/modelopt/torch/export/modeling/families/nemotron.py b/modelopt/torch/export/modeling/families/nemotron.py new file mode 100644 index 00000000000..1999b65c7b4 --- /dev/null +++ b/modelopt/torch/export/modeling/families/nemotron.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Nemotron family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="nemotron_h", + # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). + moe_block_names=("NemotronHMOE",), + expert_linear_names=("up_proj", "down_proj"), + has_iterable_experts=True, + ) +) + +# Dense Nemotron MLP treats up_proj as the fc projection (not gate). +register( + ModelSpec( + name="nemotron", + mlp_block_names=("NemotronMLP",), + mlp_keyword_roles={"up_proj": "fc"}, + ) +) diff --git a/modelopt/torch/export/modeling/families/phi.py b/modelopt/torch/export/modeling/families/phi.py new file mode 100644 index 00000000000..38e9d750779 --- /dev/null +++ b/modelopt/torch/export/modeling/families/phi.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Phi family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="phi3", + decoder_types=("phi3",), + forced_activation="swiglu", + ) +) + +# Phi3Small and the related TLGv4 MLP treat up_proj as the fc projection (not gate). +register( + ModelSpec( + name="phi3_small", + mlp_block_names=("Phi3SmallMLP", "TLGv4MLP"), + mlp_keyword_roles={"up_proj": "fc"}, + ) +) diff --git a/modelopt/torch/export/modeling/families/qwen.py b/modelopt/torch/export/modeling/families/qwen.py new file mode 100644 index 00000000000..9f866e7063f --- /dev/null +++ b/modelopt/torch/export/modeling/families/qwen.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="qwen_moe", + moe_block_names=( + "Qwen2MoeSparseMoeBlock", + "Qwen3MoeSparseMoeBlock", + "Qwen3NextSparseMoeBlock", + "Qwen3_5MoeSparseMoeBlock", + ), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + has_iterable_experts=True, + ) +) + +# Qwen3 (dense + MoE) AWQ pre_quant_scale fusion: fold o_proj into v_proj, +# down_proj into up_proj. +register( + ModelSpec( + name="qwen3", + pqs_fuse_rules=( + (("Qwen3Attention", "Qwen3MoeAttention"), "v_proj", "o_proj"), + (("Qwen3MLP", "Qwen3MoeMLP"), "up_proj", "down_proj"), + ), + ) +) diff --git a/modelopt/torch/export/modeling/hooks.py b/modelopt/torch/export/modeling/hooks.py new file mode 100644 index 00000000000..c783091ef54 --- /dev/null +++ b/modelopt/torch/export/modeling/hooks.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-model behavioral hooks for the export engine. + +A family subclasses :class:`ModelHooks` and overrides only the seams it needs; every +default method declines (returns the no-op sentinel) so the engine falls back to its +generic path. The engine builds each config object and the hook only routes it, so most +hooks need nothing from :class:`ExportContext`. +""" + +import torch.nn as nn + +from .context import ExportContext + +__all__ = ["NULL_HOOKS", "ModelHooks"] + + +class ModelHooks: + """Optional per-model overrides for the export engine's structural seams.""" + + def unwrap_decoder_layer(self, module: nn.Module, ctx: ExportContext) -> nn.Module | None: + """Return the real decoder layer when a family wraps it (e.g. DBRX/ExaOne/Deci). + + Returns ``None`` to use ``module`` unchanged. + """ + return None + + def place_submodule( + self, name: str, module: nn.Module, built, layer_config, ctx: ExportContext + ) -> bool: + """Route an already-built sub-module config into a family-specific slot. + + ``built`` is the config object the engine produced for ``module`` (e.g. a + ``LayernormConfig`` or ``AttentionConfig``); assign it onto ``layer_config`` and + return ``True`` if handled, else return ``False`` to use the default placement. + """ + return False + + def build_moe(self, module: nn.Module, ctx: ExportContext): + """Build and return this family's MoE config, or ``None`` to use the default.""" + return None + + +# Shared no-op instance used when a model has no overrides. +NULL_HOOKS = ModelHooks() diff --git a/modelopt/torch/export/modeling/registry.py b/modelopt/torch/export/modeling/registry.py new file mode 100644 index 00000000000..068c9d1462d --- /dev/null +++ b/modelopt/torch/export/modeling/registry.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Registry that resolves a model (or one of its sub-modules) to its ``ModelSpec``. + +Families register their specs at import time (see ``families/``). Lookups return +``None`` when nothing matches, so callers can fall back to their default behavior. +""" + +import torch.nn as nn + +from .base import ModelSpec + +__all__ = [ + "iter_pqs_fuse_rules", + "match_by_architecture", + "match_by_decoder_type", + "match_mlp_block", + "match_moe_block", + "register", +] + +_SPECS: list[ModelSpec] = [] + + +def register(spec: ModelSpec) -> ModelSpec: + """Register a model-family spec and return it.""" + _SPECS.append(spec) + return spec + + +def match_by_architecture(architecture: str) -> ModelSpec | None: + """Return the spec whose ``architectures`` contains an exact ``architecture`` match.""" + for spec in _SPECS: + if architecture in spec.architectures: + return spec + return None + + +def iter_pqs_fuse_rules(): + """Yield every ``(module_class_substrings, fuse_into, fuse_from)`` AWQ fusion rule. + + Aggregated across all registered specs (the consumer matches each model module + against the substrings, so the order across families does not matter). + """ + for spec in _SPECS: + yield from spec.pqs_fuse_rules + + +def match_by_decoder_type(decoder_type: str) -> ModelSpec | None: + """Return the spec whose ``decoder_types`` contains ``decoder_type`` (exact match).""" + for spec in _SPECS: + if decoder_type in spec.decoder_types: + return spec + return None + + +def match_moe_block(module: nn.Module) -> ModelSpec | None: + """Return the spec matching ``module``'s class name against ``moe_block_names``. + + Case-insensitive substring match against ``type(module).__name__``. + """ + cls_name = type(module).__name__.lower() + for spec in _SPECS: + if any(name.lower() in cls_name for name in spec.moe_block_names): + return spec + return None + + +def match_mlp_block(module: nn.Module) -> ModelSpec | None: + """Return the spec whose ``mlp_block_names`` exactly equals ``module``'s class name.""" + cls_name = type(module).__name__ + for spec in _SPECS: + if cls_name in spec.mlp_block_names: + return spec + return None diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 2af5f6eab0b..5413e914680 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -70,6 +70,7 @@ QUANTIZATION_W4A8_NVFP4_FP8, QUANTIZATION_W4A16_NVFP4, ) +from .modeling import iter_pqs_fuse_rules logger = logging.getLogger(__name__) @@ -1139,19 +1140,14 @@ def _update_svdquant(modules, new_pre_quant_scale): finish_stats_collection(module.weight_quantizer) -# Format: (list of target modules, tuple of (linear_to_fuse_into, linear_from_with_scale)) -PQS_FUSE_MODULE_MAPPING = [ - # Attention: Fuse o_proj's pre_quant_scale into v_proj's output dimension - # Mathematical equivalence: - # Before: o_proj_out = [attn @ (v_proj_in @ v_proj.W^T)^T * scale] @ o_proj.W^T - # After: o_proj_out = [attn @ (v_proj_in @ (v_proj.W * scale)^T)^T] @ o_proj.W^T - (["LlamaAttention", "Qwen3Attention", "Qwen3MoeAttention"], ("v_proj", "o_proj")), - # MLP: Fuse down_proj's pre_quant_scale into up_proj's output dimension - # Mathematical equivalence: - # Before: down_proj_out = {[act_fn(self.gate_proj(x)) * up_proj(x)] * scale} @ down_proj.W^T - # After: down_proj_out = {[act_fn(self.gate_proj(x)) * (up_proj(x) * scale)]} @ down_proj.W^T - (["LlamaMLP", "Qwen3MLP", "Qwen3MoeMLP"], ("up_proj", "down_proj")), -] +# AWQ pre_quant_scale fusion rules are per-model data and live in modeling/families/*: +# - Attention: fold o_proj's pre_quant_scale into v_proj's output dimension. +# Before: o_proj_out = [attn @ (v_proj_in @ v_proj.W^T)^T * scale] @ o_proj.W^T +# After: o_proj_out = [attn @ (v_proj_in @ (v_proj.W * scale)^T)^T] @ o_proj.W^T +# - MLP: fold down_proj's pre_quant_scale into up_proj's output dimension. +# Before: down_proj_out = {[act_fn(gate_proj(x)) * up_proj(x)] * scale} @ down_proj.W^T +# After: down_proj_out = {[act_fn(gate_proj(x)) * (up_proj(x) * scale)]} @ down_proj.W^T +# Each rule is (module_class_substrings, fuse_into, fuse_from); see ``iter_pqs_fuse_rules``. def fuse_prequant_to_linear(model: torch.nn.Module, fuse_grouped_heads=False): @@ -1167,12 +1163,10 @@ def fuse_prequant_to_linear(model: torch.nn.Module, fuse_grouped_heads=False): """ # Fuse pre_quant_scale to the linear weights for _, module in model.named_modules(): - for module_map in PQS_FUSE_MODULE_MAPPING: - target_module_list = module_map[0] - linear_pair = module_map[1] + for target_module_list, fuse_into, fuse_from in iter_pqs_fuse_rules(): if any(module_name in type(module).__name__ for module_name in target_module_list): - linear_fuse_into = module.get_submodule(linear_pair[0]) - linear_pqs_from = module.get_submodule(linear_pair[1]) + linear_fuse_into = module.get_submodule(fuse_into) + linear_pqs_from = module.get_submodule(fuse_from) if hasattr(linear_pqs_from, "input_quantizer") and hasattr( linear_pqs_from.input_quantizer, "_pre_quant_scale" ):