LMDeploy中获取模型输出Logits的技术解析

背景介绍

LMDeploy作为InternLM团队推出的高效推理工具链，在部署大语言模型方面表现出色。在实际应用中，开发者经常需要获取模型输出的原始logits而非经过softmax处理后的概率值，特别是在构建奖励模型或进行模型分析等场景下。

Logits获取需求分析

在LMDeploy中，用户最初只能通过pipe.get_logits方法获取输入序列的logits，但该方法存在线程安全问题且无法直接获取生成序列的logits。这限制了在以下场景的应用：

1. 奖励模型构建：当需要扩展v_head权重时，直接获取logits更为方便
2. 模型分析：需要分析模型生成过程中的置信度分布
3. 高级采样策略：基于原始logits实现自定义采样逻辑

技术演进过程

LMDeploy团队在v0.7.0版本中对此功能进行了重要改进：

1. 初始方案：提供基础的get_logits方法，返回seq_len×vocab_size形状的张量
2. 线程安全改进：解决了原始方案的线程安全问题
3. 生成序列支持：新增对输出序列logits的获取能力
4. 多模态扩展：支持包括InternVL2在内的视觉语言模型

实际应用示例

对于InternVL2-1B这样的多模态模型，现在可以通过以下方式获取生成文本的logits：

from lmdeploy import pipeline, TurbomindEngineConfig
from lmdeploy.vl import load_image

# 初始化管道
model = 'OpenGVLab/InternVL2-1B'
pipe = pipeline(model, backend_config=TurbomindEngineConfig(session_len=8192))

# 处理图像并生成描述
image = load_image('tiger.jpeg')
response = pipe(('describe this image', image))

# 获取生成文本的logits
generation_logits = pipe.get_logits(response)

技术细节说明

1. logits形状：返回的张量维度为[生成序列长度, 词表大小]
2. 温度参数影响：当temperature=0时，实际采样退化为贪心搜索，此时top-1概率为1.0
3. 与logprobs区别：logits是softmax前的原始输出，而logprobs是经过log-softmax处理后的结果

最佳实践建议

1. 版本要求：确保使用v0.7.0或更高版本
2. 性能考虑：批量获取logits时注意内存占用
3. 应用场景：
   • 模型微调监控
   • 生成质量评估
   • 自定义解码策略实现

未来展望

随着LMDeploy的持续发展，预计将在以下方面进一步优化：

1. 更灵活的logits获取接口
2. 对更多模型架构的支持
3. 分布式场景下的优化
4. 与训练框架的深度集成

通过本文介绍的技术方案，开发者可以更高效地在LMDeploy中实现基于logits的高级应用，充分发挥大语言模型的潜力。