support-new-model

# Tutorial: Adding a New Model to LMDeploy (PyTorch Backend) This guide walks through adding a new LLM or VLM to LMDeploy's PyTorch backend. ______________________________________________________________________ ## Before Writing Any Code **Study the reference implementations before touching any files.** 1. Read the HF model's `config.json` to understand: `model_type`, `architectures`, layer counts, hidden dims, number of attention heads, MoE parameters (if applicable). 2. Identify which category the model falls into: - **LLM only** — pure text model - **VLM** — text + vision (needs an additional preprocessor in `vl/model/`) 3. Find the closest existing model in LMDeploy and read it thoroughly: | Reference model | File(s) | | ---------------------- | ----------------------------------------------------------------------------------------------------------------------------------------- | | LLM (dense) | `lmdeploy/pytorch/models/qwen3.py` | | LLM (MoE) | `lmdeploy/pytorch/models/qwen3_moe.py` | | VLM preprocessor | `lmdeploy/vl/model/qwen3.py` | | VLM (composite config) | `lmdeploy/pytorch/models/qwen3_omni_moe_thinker.py` + `lmdeploy/pytorch/configurations/qwen3_omni.py` + `lmdeploy/vl/model/qwen3_omni.py` | ______________________________________________________________________ ## Key Files Quick Reference | File | Purpose | | -------------------------------------------- | --------------------------------------------------------------- | | `lmdeploy/pytorch/models/<model>.py` | Attention, MLP, DecoderLayer, Model, ForCausalLM | | `lmdeploy/pytorch/models/module_map.py` | HF class name → LMDeploy class path mapping | | `lmdeploy/pytorch/configurations/<model>.py` | Config builder — only needed for non-standard/nested HF configs | | `lmdeploy/vl/model/<model>.py` | VLM: image/video preprocessing *(VLM only)* | | `lmdeploy/vl/model/base.py` | `VisionModel` base class + `VISION_MODELS` registry | | `lmdeploy/archs.py` | VLM: arch name → task mapping *(VLM only)* | | `lmdeploy/lite/apis/calibrate.py` | Quantization: layer/norm/head mappings *(optional)* | | `lmdeploy/lite/quantization/awq.py` | Quantization: AWQ scale mappings *(optional)* | ______________________________________________________________________ ## Step-by-Step: LLM (PyTorch Backend) ### Step 1 — Create the PyTorch model file **File:** `lmdeploy/pytorch/models/<model_name>.py` Implement the following class hierarchy (innermost → outermost): 1. **`<Model>Attention`** — QKV projection, rotary embedding, attention forward 2. **`<Model>MLP`** — gate-up linear, activation, down projection 3. **`<Model>DecoderLayer`** — wraps Attention + MLP with layer norms and residual connections 4. **`<Model>Model`** — embedding table, all decoder layers, final norm, rotary embedding 5. **`<Model>ForCausalLM`** — top-level class; inherits `nn.Module`, `DeployModelMixinV1`, `CudaGraphMixin` **Required imports:** ```python import torch import torch.nn as nn from lmdeploy.pytorch.model_inputs import StepContext, StepContextManager from lmdeploy.pytorch.nn import (ApplyRotaryEmb, Attention, RMSNorm, SiluAndMul, build_rotary_embedding_from_config) from lmdeploy.pytorch.nn.linear import (build_down_linear, build_gateup_linear, build_o_proj, build_qkv_proj) from lmdeploy.pytorch.weight_loader.model_weight_loader import load_weight from .patch import add_prefix from .utils.cudagraph import CudaGraphMixin from .utils.model import DeployModelMixinV1, build_embedding ``` **Attention skeleton:** ```python class MyModelAttention(nn.Module): def __init__(self, config, dtype=None, device=None, prefix=''): super().__init__() self.qkv_proj = build_qkv_proj( config.hidden_size, num_q_heads=config.num_attention_heads, num_kv_heads=config.num_key_value_heads, head_size=config.hidden_size // config.num_attention_heads, bias=False, dtype=dtype, device=device, prefix=add_prefix('qkv_proj', prefix)) self.apply_rotary_pos_emb = ApplyRotaryEmb() self.attn_fwd = Attention( config.num_attention_heads, config.hidden_size // config.num_attention_heads, num_kv_heads=config.num_key_value_heads) self.o_proj = build_o_proj( config.num_attention_heads, config.hidden_size // config.num_attention_heads, config.hidden_size, bias=False, dtype=dtype, device=device, prefix=add_prefix('o_proj', prefix)) def forward(self, hidden_states, rotary_pos_emb, past_key_value, attn_metadata): qkv_states = self.qkv_proj(hidden_states) # split q, k, v; apply rotary; call attn_fwd; project output ... ``` **MLP skeleton:** ```python class MyModelMLP(nn.Module): def __init__(self, config, dtype=None, device=None, prefix=''): super().__init__() self.gate_up_proj = build_gateup_linear( config.hidden_size, config.intermediate_size, bias=False, dtype=dtype, device=device, prefix=add_prefix('gate_up_proj', prefix)) self.down_proj = build_down_linear( config.intermediate_size, config.hidden_size, bias=False, dtype=dtype, device=device, prefix=add_prefix('down_proj', prefix)) self.act_fn = SiluAndMul() def forward(self, x): return self.down_proj(self.act_fn(self.gate_up_proj(x))) ``` **ForCausalLM skeleton (critical fields):** ```python class MyModelForCausalLM(nn.Module, DeployModelMixinV1, CudaGraphMixin): # Maps packed param name → list of original HF param suffixes packed_modules_mapping = { 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'], 'gate_up_proj': ['gate_proj', 'up_proj'], } def __init__(self, config, ctx_mgr=None, prefix='', **kwargs): super().__init__() self.model = MyModelModel(config, ...) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.ctx_mgr = ctx_mgr def get_input_embeddings(self): return self.model.embed_tokens def forward(self, input_ids, inputs_embeds, past_key_values, attn_metadata, **kwargs): hidden_states = self.model(input_ids, inputs_embeds, past_key_values, attn_metadata) return hidden_states def get_logits(self, hidden_states): return self.lm_head(hidden_states) # prepare_inputs_for_generation and load_weights: copy from qwen3.py, # update stacked_params_mapping to match this model's HF weight names. ``` ______________________________________________________________________ ### Step 2 — Register in `module_map.py` **File:** `lmdeploy/pytorch/models/module_map.py` Add an entry to `MODULE_MAP`. The key is the exact HF architecture class name from `config.json`'s `architectures` field: ```python MODULE_MAP.update({ 'MyModelForCausalLM': f'{LMDEPLOY_PYTORCH_MODEL_PATH}.my_model.MyModelForCausalLM', }) ``` ______________________________________________________________________ ### Step 3 — Add config builder (if needed) **File:** `lmdeploy/pytorch/configurations/<model_name>.py` **Skip this step** for models with a standard flat HF config — `DefaultModelConfigBuilder` handles them automatically. Only create this file when the HF config is non-standard, e.g.: - Nested config (e.g., Qwen3-Omni has `hf_config.thinker_config.text_config`) - Unusual `model_type` that needs special field remapping ```python from .builder import AutoModelConfigBuilder, DefaultModelConfigBuilder class MyModelConfigBuilder(AutoModelConfigBuilder): @classmethod def condition(cls, hf_config): # Must match model_type from config.json exactly return hf_config.model_type == 'my_model' @classmethod def build(cls, hf_config, model_path=None, **kwargs): # Extract the text config if nested; patch fields if needed cfg = DefaultModelConfigBuilder.build(hf_config, model_path, **kwargs) cfg.hf_config = hf_config # keep full config for VLM layers return cfg ``` Auto-discovery: subclasses of `AutoModelConfigBuilder` register themselves automatically via `__init_subclass__()` — no import needed elsewhere. ______________________________________________________________________ ### Step 4 — Add quantization mappings (optional) Only needed to support AWQ/SmoothQuant calibration for this model family. **`lmdeploy/lite/apis/calibrate.py`** — add layer name, norm name, and head name mappings for the new model type. **`lmdeploy/lite/quantization/awq.py`** — add entries to `NORM_FCS_MAP` (norm → downstream FC layers) and `FC_FCS_MAP` (FC → downstream FC layers) following the existing patterns.