跳至内容

内存优化

Large models might cause your machine to run out of memory (OOM). Here are some options that help alleviate this problem.

张量并行 (TP)

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

以下代码将模型分布在2个GPU上运行。

from vllm import LLM

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

警告

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

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

注意

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

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

量化

量化模型以降低精度为代价占用更少内存。

静态量化模型可直接从HF Hub下载(部分热门模型可在Red Hat AI获取),无需额外配置即可使用。

动态量化也通过quantization选项支持——详见此处获取更多信息。

上下文长度与批量大小

您可以通过限制模型的上下文长度(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的额外内存。

警告

在V1版本中,CUDA图捕获比V0版本占用更多内存。

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

from vllm import LLM
from vllm.config import CompilationConfig, CompilationLevel

llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    compilation_config=CompilationConfig(
        level=CompilationLevel.PIECEWISE,
        # By default, it goes up to 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)

调整缓存大小

If you run out of CPU RAM, try the following options:

  • (仅限多模态模型) 您可以使用环境变量VLLM_MM_INPUT_CACHE_GIB来设置多模态输入缓存的大小(默认为4 GiB)。
  • (仅限CPU后端) 您可以通过VLLM_CPU_KVCACHE_SPACE环境变量设置KV缓存大小(默认为4 GiB)。

多模态输入限制

您可以限制每个提示中的多模态项目数量以减少模型的内存占用:

from vllm import LLM

# Accept up to 3 images and 1 video per prompt
llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct",
          limit_mm_per_prompt={"image": 3, "video": 1})

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

from vllm import LLM

# Accept any number of images but no videos
llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct",
          limit_mm_per_prompt={"video": 0})

你甚至可以运行一个多模态模型仅用于文本推理:

from vllm import LLM

# Don't accept images. Just text.
llm = LLM(model="google/gemma-3-27b-it",
          limit_mm_per_prompt={"image": 0})

多模态处理器参数

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

以下是一些示例:

from vllm import LLM

# Available for Qwen2-VL series models
llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct",
          mm_processor_kwargs={
              "max_pixels": 768 * 768,  # Default is 1280 * 28 * 28
          })

# Available for InternVL series models
llm = LLM(model="OpenGVLab/InternVL2-2B",
          mm_processor_kwargs={
              "max_dynamic_patch": 4,  # Default is 12
          })