跳到內容

記憶體最佳化

大型模型可能會導致您的機器記憶體不足 (OOM)。以下是一些有助於緩解此問題的選項。

張量並行 (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 圖

預設情況下,我們使用 CUDA 圖最佳化模型推理,這會佔用 GPU 中的額外記憶體。

警告

CUDA 圖捕獲在 V1 中比在 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)

調整快取大小

如果您的 CPU 記憶體不足,請嘗試以下選項

  • (僅限多模態模型)您可以使用 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
        })