跳到內容

記憶體最佳化

大型模型可能會導致您的機器記憶體不足 (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 的額外記憶體。

您可以調整 compilation_config 以在推理速度和記憶體使用之間取得更好的平衡。

程式碼
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,
        # 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 RAM 不足,請嘗試以下選項。

  • (僅限多模態模型) 您可以透過設定 mm_processor_cache_gb 引擎引數(預設為 4 GiB)來設定多模態快取的大小。
  • (僅限 CPU 後端) 您可以使用 VLLM_CPU_KVCACHE_SPACE 環境變數(預設為 4 GiB)來設定 KV 快取的大小。

多模態輸入限制

您可以允許每個提示中的多模態專案數量較少,以減少模型的記憶體佔用。

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},
)

可配置選項

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

# Up to 5 images per prompt, profile with 512x512.
# Up to 1 video per prompt, profile with 32 frames at 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}}

注意

  • 大小提示僅影響記憶體剖析。它們會塑造用於計算預留啟用大小的虛擬輸入。它們不會改變實際輸入在推理時的處理方式。
  • 如果提示的大小超過了模型可以接受的最大值,vLLM 會將其限制為模型的有效最大值,並可能記錄警告。

警告

這些大小提示目前僅影響啟用記憶體剖析。編碼器快取的大小由執行時實際輸入決定,不受這些提示的限制。

多模態處理器引數

對於某些模型,您可以調整多模態處理器引數以減小已處理多模態輸入的大小,從而節省記憶體。

以下是一些示例。

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
)