Skip to content

内存优化

大型模型可能导致您的机器内存不足(OOM)。以下是一些有助于缓解此问题的选项。

张量并行(Tensor Parallelism)

张量并行(tensor_parallel_size 选项)可用于将模型分布在多个 GPU 上。

以下代码将模型分布在 2 个 GPU 上:

from vllm import LLM

llm = LLM(model="ibm-granite/granite-3.1-8b-instruct", tensor_parallel_size=2)

Warning

为确保 vLLM 正确初始化 CUDA,您应避免在初始化 vLLM 之前调用相关函数(例如 torch.cuda.set_device)。 否则,您可能会遇到类似 RuntimeError: Cannot re-initialize CUDA in forked subprocess 的错误。

要控制使用哪些设备,请改用设置 CUDA_VISIBLE_DEVICES 环境变量。

Note

启用张量并行后,每个进程都会读取完整模型并将其分割为块,这会进一步延长磁盘读取时间(与张量并行大小成正比)。

您可以使用 examples/offline_inference/save_sharded_state.py 将模型检查点转换为分片检查点。转换过程可能需要一些时间,但之后您可以更快地加载分片检查点。无论张量并行大小如何,模型加载时间应保持不变。

量化

量化模型占用更少内存,但精度较低。

静态量化模型可以从 HF Hub 下载(一些流行模型可在 Red Hat AI 获得),无需额外配置即可直接使用。

还支持动态量化,通过 quantization 选项实现——更多详情请参见 here

上下文长度和批大小

您可以通过限制模型的上下文长度(max_model_len 选项)和最大批大小(max_num_seqs 选项)进一步减少内存使用。

from vllm import LLM

llm = LLM(model="adept/fuyu-8b", max_model_len=2048, max_num_seqs=2)

减少 CUDA Graphs

默认情况下,我们使用 CUDA 图优化模型推理,这会占用 GPU 的额外内存。

您可以调整 compilation_config 来在推理速度和内存使用之间取得更好的平衡:

Code
from vllm import LLM
from vllm.config import CompilationConfig, CompilationMode

llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    compilation_config=CompilationConfig(
        mode=CompilationMode.VLLM_COMPILE,
        # 默认情况下,它会递增至 max_num_seqs
        cudagraph_capture_sizes=[1, 2, 4, 8, 16],
    ),
)

您可以通过 enforce_eager 标志完全禁用图捕获:

from vllm import LLM

llm = LLM(model="meta-llama/Llama-3.1-8B-Instruct", enforce_eager=True)

调整缓存大小

如果您内存不足,可以尝试以下选项:

  • (仅多模态模型)您可以通过设置引擎参数 mm_processor_cache_gb 来设置多模态缓存的大小(默认 4 GiB)。
  • (仅 CPU 后端)您可以通过环境变量 VLLM_CPU_KVCACHE_SPACE 设置 KV 缓存的大小(默认 4 GiB)。

多模态输入限制

您可以通过限制每提示允许的多模态项目数量来减少模型的内存占用:

from vllm import LLM

# 每个提示最多接受 3 张图像和 1 个视频
llm = LLM(
    model="Qwen/Qwen2.5-VL-3B-Instruct",
    limit_mm_per_prompt={"image": 3, "video": 1},
)

您还可以进一步完全禁用未使用的模态,通过将其限制设为零。例如,如果您的应用仅接受图像输入,则无需为视频分配任何内存。

from vllm import LLM

# 接受任意数量的图像,但不接受视频
llm = LLM(
    model="Qwen/Qwen2.5-VL-3B-Instruct",
    limit_mm_per_prompt={"video": 0},
)

您甚至可以仅对文本进行推理运行多模态模型:

from vllm import LLM

# 不接受图像。仅文本。
llm = LLM(
    model="google/gemma-3-27b-it",
    limit_mm_per_prompt={"image": 0},
)

可配置选项

limit_mm_per_prompt 还接受每模态的可配置选项。在可配置形式中,您仍需指定 count,并且可以选择提供大小提示,这些提示控制 vLLM 如何对您的多模态输入进行分析和预留内存。这有助于您为预期的实际媒体内容调优内存,而不是模型的绝对最大值。

每模态的可配置选项:

  • image: {"count": int, "width": int, "height": int}
  • video: {"count": int, "num_frames": int, "width": int, "height": int}
  • audio: {"count": int, "length": int}

详情请参阅 ImageDummyOptionsVideoDummyOptionsAudioDummyOptions

示例:

from vllm import LLM

# 每个提示最多 5 张图像,使用 512x512 进行分析。
# 每个提示最多 1 个视频,使用 32 帧 640x640 进行分析。
llm = LLM(
    model="Qwen/Qwen2.5-VL-3B-Instruct",
    limit_mm_per_prompt={
        "image": {"count": 5, "width": 512, "height": 512},
        "video": {"count": 1, "num_frames": 32, "width": 640, "height": 640},
    },
)

为了向后兼容,传递整数仍然按原方式工作,被解释为 {"count": <int>}。例如:

  • limit_mm_per_prompt={"image": 5} 等价于 limit_mm_per_prompt={"image": {"count": 5}}
  • 您可以混合格式:limit_mm_per_prompt={"image": 5, "video": {"count": 1, "num_frames": 32, "width": 640, "height": 640}}

Note

  • 大小提示仅影响内存分析。它们用于塑造用于计算预留激活大小的虚拟输入。它们不会改变推理时实际处理输入的方式。
  • 如果提示值超过模型可接受的范围,vLLM 会将其限制为模型的有效最大值,并可能记录警告。

Warning

这些大小提示目前仅影响激活内存分析。编码器缓存大小由运行时的实际输入决定,不受这些提示限制。

多模态处理器参数

对于某些模型,您可以调整多模态处理器参数以减小处理后的多模态输入大小,从而节省内存。

以下是一些示例:

from vllm import LLM

# 适用于 Qwen2-VL 系列模型
llm = LLM(
    model="Qwen/Qwen2.5-VL-3B-Instruct",
    mm_processor_kwargs={"max_pixels": 768 * 768},  # 默认为 1280 * 28 * 28
)

# 适用于 InternVL 系列模型
llm = LLM(
    model="OpenGVLab/InternVL2-2B",
    mm_processor_kwargs={"max_dynamic_patch": 4},  # 默认为 12
)