当前位置：首页 > news >正文

LMDeploy终极贡献指南：如何快速提交新模型支持的完整教程

news 2026/8/23 10:18:07

LMDeploy终极贡献指南：如何快速提交新模型支持的完整教程

【免费下载链接】lmdeployLMDeploy is a toolkit for compressing, deploying, and serving LLMs.项目地址: https://gitcode.com/gh_mirrors/lm/lmdeploy

LMDeploy是一个功能强大的大语言模型压缩、部署和服务工具包，支持众多主流LLM和VLM模型。如果你想让自己的模型在LMDeploy中获得支持，这篇完整教程将指导你完成从零开始到成功提交PR的整个过程。🚀

📋 准备工作与开发环境搭建

在开始贡献之前，首先需要搭建开发环境。克隆LMDeploy仓库并设置开发分支：

git clone https://gitcode.com/gh_mirrors/lm/lmdeploy cd lmdeploy git checkout -b your_username/add_new_model_support

配置pre-commit钩子以确保代码风格一致：

pip install -U pre-commit pre-commit install pre-commit run --all-files

🏗️ 理解LMDeploy模型支持架构

LMDeploy支持两种推理引擎：TurboMind（追求极致性能）和PyTorch（开发友好）。添加新模型主要涉及PyTorch引擎的实现。

核心文件结构

添加新模型需要关注以下几个关键目录：

模型配置: lmdeploy/pytorch/configurations/ - 存放模型配置解析器
模型实现: lmdeploy/pytorch/models/ - 存放具体的模型实现
模型注册: lmdeploy/pytorch/models/module_map.py - 模型注册映射表
官方文档: docs/en/advance/pytorch_new_model.md - 官方添加新模型指南

🔧 第一步：创建模型配置解析器

如果你的模型参数命名与常见transformers模型不同，需要创建自定义的ConfigBuilder。以Gemma模型为例：

# lmdeploy/pytorch/configurations/gemma.py from lmdeploy.pytorch.config import ModelConfig from .builder import AutoModelConfigBuilder class GemmaModelConfigBuilder(AutoModelConfigBuilder): @classmethod def condition(cls, hf_config): # 检查hf_config是否适合此构建器 return hf_config.model_type in ['gemma', 'gemma2'] @classmethod def build(cls, hf_config, model_path: str = None): # 使用transformers加载的hf_config # 为pytorch引擎构造ModelConfig return ModelConfig( hidden_size=hf_config.hidden_size, num_layers=hf_config.num_hidden_layers, num_attention_heads=hf_config.num_attention_heads, num_key_value_heads=hf_config.num_key_value_heads, bos_token_id=hf_config.bos_token_id, eos_token_id=hf_config.eos_token_id, head_dim=hf_config.head_dim, vocab_size=hf_config.vocab_size )

🧩 第二步：实现模型核心逻辑

参考现有的Llama模型实现，创建你的模型类。关键组件包括：

1. 模型结构定义

from typing import Any, Dict, Iterable, List, Optional, Tuple import torch from torch import nn from transformers import YourModelConfig from lmdeploy.pytorch.model_inputs import StepContext, StepContextManager from lmdeploy.pytorch.nn import ApplyRotaryEmb, Attention, RMSNorm, SiluAndMul from lmdeploy.pytorch.nn.linear import ( build_down_linear, build_gateup_linear, build_o_proj, build_qkv_proj, build_rowwise_linear ) class YourModelForCausalLM(nn.Module): def __init__(self, config: YourModelConfig, ctx_mgr: StepContextManager, dtype: torch.dtype = None, device: torch.device = None): super().__init__() self.config = config self.ctx_mgr = ctx_mgr # 构建模型组件 self.model = YourModelModel(config, dtype=dtype, device=device) self.lm_head = build_rowwise_linear( config.hidden_size, config.vocab_size, bias=False, dtype=dtype, device=device ) def forward(self, input_ids, position_ids, past_key_values, attn_metadata=None, inputs_embeds=None, **kwargs): # 前向传播逻辑 hidden_states = self.model( input_ids=input_ids, position_ids=position_ids, past_key_values=past_key_values, attn_metadata=attn_metadata, inputs_embeds=inputs_embeds, ) logits = self.lm_head(hidden_states) return logits.float()

上图展示了LMDeploy通过量化技术在不同批处理大小下的内存优化效果，WeightInt4和kCacheKVInt8量化策略显著降低了内存消耗。

2. 使用预定义的融合算子

LMDeploy提供了许多融合算子来简化模型构建，这些算子更好地支持张量并行和量化等功能：

class YourModelMLP(nn.Module): def __init__(self, config: YourModelConfig, dtype: torch.dtype = None, device: torch.device = None): super().__init__() quantization_config = getattr(config, 'quantization_config', None) # 使用预定义的线性层构建器 self.gate_up_proj = build_merged_colwise_linear( config.hidden_size, [config.intermediate_size, config.intermediate_size], bias=config.mlp_bias, dtype=dtype, device=device, quant_config=quantization_config, is_tp=True, # 支持张量并行 ) self.act_fn = SiluAndMul(inplace=True) self.down_proj = build_rowwise_linear( config.intermediate_size, config.hidden_size, bias=config.mlp_bias, quant_config=quantization_config, dtype=dtype, device=device, is_tp=True ) def forward(self, x): gate_up = self.gate_up_proj(x) act = self.act_fn(gate_up) return self.down_proj(act)

📝 第三步：注册模型到模块映射

在 lmdeploy/pytorch/models/module_map.py 中添加你的模型映射：

MODULE_MAP.update({ 'YourModelForCausalLM': 'lmdeploy.pytorch.models.your_model.YourModelForCausalLM', })

🧪 第四步：编写单元测试

创建测试文件确保你的实现正确：

# tests/test_lmdeploy/test_your_model.py import torch from lmdeploy.pytorch.models.your_model import YourModelForCausalLM from transformers import YourModelConfig def test_your_model_forward(): config = YourModelConfig( hidden_size=4096, num_hidden_layers=32, num_attention_heads=32, vocab_size=32000 ) model = YourModelForCausalLM(config) # 测试前向传播 batch_size = 2 seq_len = 128 input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len)) with torch.no_grad(): outputs = model(input_ids) assert outputs.shape == (batch_size, seq_len, config.vocab_size)

运行测试确保一切正常：

pytest tests/test_lmdeploy/test_your_model.py

🔍 第五步：验证配置解析

使用LMDeploy的检查工具验证配置解析是否正确：

from lmdeploy.pytorch.check_env import check_model # 验证模型配置 check_model("your_model_path", model_name="your_model_name")

📦 第六步：创建Pull Request

1. 提交代码到远程仓库

git add . git commit -m "[Feature] Add support for YourModel" git push -u origin your_username/add_new_model_support

2. 创建PR描述模板

在GitHub上创建Pull Request时，使用以下模板：

## 新增功能说明 - 添加了对YourModel模型的支持 - 实现了完整的PyTorch引擎集成 - 支持所有标准功能（量化、张量并行等） ## 相关Issue - 关联Issue #1234（如有） ## 测试验证 - [x] 单元测试通过 - [x] 配置解析测试通过 - [x] 与现有模型兼容性测试 ## 代码变更 - 新增文件: lmdeploy/pytorch/models/your_model.py - 修改文件: lmdeploy/pytorch/models/module_map.py - 新增测试: tests/test_lmdeploy/test_your_model.py

🚀 高级功能集成

支持CUDA图优化

如果你的模型支持CUDA图，可以添加相应属性：

class YourModelForCausalLM(nn.Module): # 表示模型是否支持cudagraph # 可以是可调用对象，接收forward输入 # 动态确定是否支持cudagraph support_cuda_graph = True

支持权重加载

实现权重加载方法：

def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): # 模型输入是state dict的键值对 weight_map = dict(weights) # 加载权重逻辑 self.model.load_state_dict(weight_map, strict=False)