Dolphin模型基于TensorRT-LLM的部署与推理指南

模型架构概述

Dolphin模型采用了一种创新的Swin编码器+MBart解码器混合架构，这种架构在HuggingFace Transformers配置中被明确标注为"VisionEncoderDecoderModel"。这种架构设计使得Dolphin能够同时处理视觉和文本信息，非常适合文档理解和多模态任务。

与Nougat和Donut等模型类似，Dolphin也采用了视觉编码器-文本解码器的结构。TensorRT-LLM已经支持Nougat模型，基于此，我们成功地将Dolphin模型移植到了TensorRT-LLM平台上。

重要提示：在使用Dolphin模型时，prompt_ids必须使用int32类型，否则TensorRT-LLM会产生错误的推理结果。

环境准备

系统要求

• 操作系统：Linux（已在Ubuntu 20.04/22.04上测试）
• Python版本：3.8或更高
• GPU：NVIDIA GPU（建议使用具有Tensor Core的架构，如Ampere或Hopper）

安装TensorRT-LLM

建议使用TensorRT-LLM 0.18.1版本，可以通过以下步骤安装：

1. 创建并激活Python虚拟环境
2. 安装PyTorch（与您的CUDA版本匹配）
3. 安装TensorRT-LLM核心包
4. 安装其他依赖项

离线推理实践

基本命令结构

离线推理的核心命令遵循以下模式：

python run_dolphin.py \
    --batch_size 1 \
    --hf_model_dir [模型目录] \
    --visual_engine_dir [视觉引擎目录] \
    --llm_engine_dir [LLM引擎目录] \
    --max_new_tokens 4096 \
    --repetition_penalty 1.0 \
    --input_text [提示文本] \
    --image_path [图片路径]

典型应用场景

1. 文档阅读顺序解析

python run_dolphin.py \
    --input_text "Parse the reading order of this document." \
    --image_path "page_1.jpeg"

2. 文本/LaTeX识别

# 识别普通文本
python run_dolphin.py \
    --input_text "Read text in the image." \
    --image_path "para_1.jpg"

# 识别数学公式
python run_dolphin.py \
    --input_text "Read text in the image." \
    --image_path "block_formula.jpeg"

3. 表格解析

python run_dolphin.py \
    --input_text "Parse the table in the image." \
    --image_path "table_1.jpeg"

在线服务部署

1. 启动API服务器

python api_server.py \
    --hf_model_dir [模型目录] \
    --visual_engine_dir [视觉引擎目录] \
    --llm_engine_dir [LLM引擎目录] \
    --max_batch_size 16

2. 客户端调用示例

文档结构分析

python api_client.py --image_path page_1.jpeg --prompt "Parse the reading order of this document."

内容识别

# 文本识别
python api_client.py --image_path para_1.jpg --prompt "Read text in the image."

# 公式识别
python api_client.py --image_path block_formula.jpeg --prompt "Read text in the image."

# 表格解析
python api_client.py --image_path table_1.jpeg --prompt "Parse the table in the image."

性能优化建议

1. 批处理大小：根据GPU内存调整batch_size参数，较大的批处理可以提高吞吐量
2. 精度选择：bfloat16在保持精度的同时能提供较好的性能
3. 令牌限制：合理设置max_new_tokens以避免不必要的计算
4. 重复惩罚：调整repetition_penalty可以控制生成文本的多样性

常见问题解答

Q: 为什么需要特别注意prompt_ids的数据类型？ A: TensorRT-LLM对输入数据类型有严格要求，使用int32以外的类型可能导致模型行为异常。

Q: 如何选择合适的max_new_tokens值？ A: 根据任务需求设置，对于短文本识别可以设置较小值，复杂文档解析则需要较大值。

Q: 视觉引擎和LLM引擎是否可以分开部署？ A: 可以，这种分离设计提供了部署灵活性，可以根据需求独立扩展视觉或语言处理能力。

通过本指南，您应该能够成功部署和运行Dolphin模型的TensorRT-LLM版本，充分利用其强大的文档理解和多模态处理能力。