OpenChat项目中Gemma模型在RTX 4090上的VRAM优化实践

背景介绍

OpenChat是一个开源的大型语言模型项目，近期发布了基于Gemma架构的新版本模型openchat-3.5-0106-gemma。与之前基于Mistral架构的模型相比，Gemma模型在性能上有所提升，但也带来了更高的显存需求。

问题现象

用户在配备24GB显存的RTX 4090显卡上尝试运行新Gemma模型时遇到了显存不足的问题。具体表现为系统报错提示"模型的max seq len(8192)大于KV缓存能存储的最大token数(5888)"。而同样的硬件环境下，旧版Mistral模型则可以正常运行。

技术分析

Gemma模型相比Mistral模型需要更多显存的主要原因包括：

1. 更大的嵌入层：Gemma模型的嵌入层(embedding)规模显著增大，这是导致显存需求上升的主要因素之一。

2. KV缓存需求：模型在推理过程中需要维护键值(KV)缓存，序列长度越长，KV缓存占用的显存就越多。Gemma模型默认支持的最大序列长度(8192)超过了RTX 4090显卡在默认配置下的显存容量。

3. 显存利用率参数：vLLM引擎默认的显存利用率设置可能不足以支持Gemma模型的高需求。

解决方案

针对这一问题，可以通过调整vLLM引擎的显存利用率参数来解决：

python -m ochat.serving.openai_api_server \
    --model ~/path/to/openchat-3.5-0106-gemma/ \
    --engine-use-ray \
    --worker-use-ray \
    --max-model-len 8192 \
    --tensor-parallel-size 1 \
    --gpu-memory-utilization 0.95 \
    --host 0.0.0.0 \
    --disable-log-requests \
    --disable-log-stats \
    --log-file openchat.log

关键参数说明：

• --gpu-memory-utilization 0.95：将GPU显存利用率提高到95%，允许引擎使用更多的显存资源
• --max-model-len 8192：保持模型支持的最大序列长度不变

实践建议

1. 显存监控：在调整显存利用率时，建议同时监控GPU显存使用情况，避免设置过高导致系统不稳定。

2. 性能权衡：提高显存利用率可能会影响系统其他部分的性能，需要根据实际应用场景进行权衡。

3. 模型选择：如果显存限制严格，可以考虑使用序列长度较短的配置，或者选择对硬件要求更低的模型版本。

4. 硬件适配：对于显存有限的设备，可能需要考虑使用量化版本或参数更少的模型变体。

总结

OpenChat项目的Gemma模型在性能提升的同时也带来了更高的硬件需求。通过合理调整vLLM引擎的显存利用率参数，可以在RTX 4090等高端消费级显卡上成功运行这一模型。这为研究者和开发者在有限硬件资源下使用最新大语言模型提供了实用参考。