vLLM显存优化

最新推荐文章于 2025-03-24 14:17:23 发布

xnuscd

最新推荐文章于 2025-03-24 14:17:23 发布

阅读量4.6k

点赞数 13

文章标签：人工智能机器学习算法

本文链接：https://blog.csdn.net/xnuscd/article/details/143232080

版权

在使用 vLLM 框架进行大模型推理时，为了最大程度地减少 GPU 显存的占用，可以从以下几个方面调整参数和配置：

1. 调整 `max_batch_size` 参数

max_batch_size：这是批处理的最大大小。较大的批处理尺寸会占用更多的显存。如果显存有限，可以降低 max_batch_size。
建议：从小的 batch_size 开始，逐步调高，找到合适的平衡点。

2. 使用 `offload` 功能

vLLM 支持将部分张量或模型参数从 GPU 显存转移到 CPU 内存中，以减少显存的占用。这可以通过 offload 功能实现。
offload_weights：如果模型的参数很大，可以启用权重的 offloading，将模型参数存储在 CPU 内存中，只有在需要时才将部分参数加载到 GPU。
offload_activations：同样地，激活值（即推理过程中产生的中间结果）也可以进行 offloading。

建议：尝试启用 offload_weights 和 offload_activations，并观察显存使用情况的变化。

3. 使用梯度累积（Gradient Checkpointing）

虽然梯度累积一般用于训练，但也可以在推理过程中应用类似的技术，通过重新计算部分中间结果来减少显存占用。在 vLLM 中，使用梯度检查点的策略会大幅减少中间激活存储在显存中的需求。

建议：如果模型支持 checkpointing，启用该功能以减少内存使用。

4. 减小模型规模或使用量化模型

如果可能，考虑使用量化后的模型，如 8-bit 或 4-bit 量化，这样可以显著降低显存的占用。

vLLM 支持加载和推理不同格式的模型，量化模型会用更少的显存来存储权重。

建议：尝试加载量化后的模型（8-bit 或 4-bit），这可以在减少显存使用的同时维持模型的性能。

5. 降低生成序列长度

生成序列的长度越长，模型需要存储和计算的中间结果越多，从而增加显存占用。如果可能的话，可以减少推理时生成的 token 数量。
例如，减少 max_length 或 max_new_tokens 的值。

建议：根据需求，适当调整 max_new_tokens 参数，减少生成长度以降低显存占用。

6. 优化显存分配（Memory Allocation）

在某些情况下，可以使用 torch.cuda.empty_cache() 来手动释放不再需要的显存，防止显存碎片化问题导致的额外占用。
你可以定期调用该方法来确保没有多余的显存占用。

7. 使用 `fp16` 模式

如果硬件支持，可以使用混合精度（fp16）模式来减少显存占用。vLLM 支持在推理过程中启用混合精度。

建议：尝试将模型权重和激活切换为 fp16 格式，这样每个浮点数只占用一半的显存。

8. 分布式推理

如果显存仍然不足，可以考虑使用多个 GPU 来分布式推理。vLLM 支持多 GPU 模式，可以根据模型规模和可用资源配置多个 GPU。

建议：如果有多块 GPU，使用分布式推理将模型切分到不同的设备上，以减少单卡的显存压力。

配置示例

假设你使用的是 Hugging Face 的模型，可以通过以下代码启动 vLLM 并减少显存使用：

from vllm import LLM

# 配置 LLM 模型
llm = LLM(model="huggingface/your-model-name",
          offload_weights=True,        # 启用权重的 offload
          offload_activations=True,    # 启用激活值的 offload
          dtype="fp16",                # 启用混合精度
          max_batch_size=1,            # 控制 batch size
          max_new_tokens=50)           # 控制生成序列长度