运行LLaVA-1.5-13B预训练模型时CUDA内存不足求助

你的RTX4070Ti（12G显存）显存容量不足以完整加载float16精度的LLaVA-1.5-13B模型——13B参数的LLM仅float16权重就需要约26G显存，再加上多模态模型的视觉编码器、推理时的临时缓存，显存占用远超12G上限。CPU能运行是因为它依赖系统内存，不受GPU显存限制，但速度较慢。以下是可行的解决方案：

1. 4bit量化加载模型（最推荐，显存占用减半） #

借助bitsandbytes库对模型做4bit量化，能将显存占用降到13G以内（刚好适配你的显卡），且几乎不损失推理效果。

修改模型加载代码：

from PIL import Image
import requests
import torch
from transformers import AutoProcessor, LlavaForConditionalGeneration

model_id = "llava-hf/llava-1.5-13b-hf"

prompt = "USER: <image>\nWhat are these?\nASSISTANT:"
image_file = "http://images.cocodataset.org/val2017/000000039769.jpg"

model = LlavaForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    load_in_4bit=True,  # 启用4bit量化
    bnb_4bit_use_double_quant=True,  # 双重量化进一步压缩显存占用
    bnb_4bit_quant_type="nf4",  # 针对LLM优化的量化类型
    bnb_4bit_compute_dtype=torch.float16  # 计算时保留float16精度保证速度
)

processor = AutoProcessor.from_pretrained(model_id)

raw_image = Image.open(requests.get(image_file, stream=True).raw)
inputs = processor(prompt, raw_image, return_tensors='pt').to("cuda", torch.float16)

output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
print(processor.decode(output[0][2:], skip_special_tokens=True))

2. 切换到更小的模型版本 #

如果不想用量化，可以换成llava-hf/llava-1.5-7b-hf（7B参数），float16精度下权重约14G，配合梯度检查点优化刚好能在12G显存运行：

model = LlavaForConditionalGeneration.from_pretrained(
    "llava-hf/llava-1.5-7b-hf",
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
).to(0)
model.gradient_checkpointing_enable()  # 开启梯度检查点减少推理缓存

3. 模型分片（自动分配到CPU/GPU） #

用device_map="auto"让transformers自动将模型部分层放到CPU、部分放到GPU，降低显存占用，代价是推理速度比全GPU慢，但比纯CPU快：

model = LlavaForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    device_map="auto"  # 自动分配设备
)

4. 推理时的额外显存优化 #

• 推理前手动清理GPU缓存：

torch.cuda.empty_cache()

• 生成时指定pad_token_id（避免自动生成冗余token占用显存）：

output = model.generate(
    **inputs,
    max_new_tokens=200,
    do_sample=False,
    pad_token_id=processor.tokenizer.pad_token_id
)

内容的提问来源于stack exchange，提问作者Jakob Hürlemann