在Google Colab T4实例运行Llama-2-7b时遇CUDA内存不足错误

问题详情 #

• 运行环境：Google Colab T4实例（15GB GPU显存）
• 触发场景：执行trainer.train()训练Llama-2-7b模型时，出现CUDA内存溢出错误
• 显存变化：崩溃前显存占用约6.5/15GB，崩溃时瞬间飙升至14.7GB
• 已尝试方案：将batch size降至1，测试Stack Overflow常见内存优化方法，问题仍存在

错误日志 #

OutOfMemoryError Traceback (most recent call last)
in <cell line: 2>()

1 # Train model

----> 2 trainer.train()

-31 frames-

/usr/local/lib/python3.10/dist-packages/bitsandbytes/autograd/_functions.py in forward(ctx, A, B, out, bias, quant_state)

514 # 1. Dequantize
515 # 2. MatmulnN

--> 516 output = torch.nn.functional.linear(A, F.dequantize_4bit(B, quant_state).to(A.dtype).t(), bias)

517
518 # 3. Save state

OutOfMemoryError: CUDA out of memory. Tried to allocate 172.00 MiB. GPU 0 has a total capacty of 14.75 GiB of which 72.81 MiB is free. Process 38354 has 14.67 GiB memory in use. Of the allocated memory 13.37 GiB is allocated by PyTorch, and 303.44 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CON

针对性解决方案 #

1. 优化4-bit量化参数 #

错误发生在bitsandbytes的4-bit反量化步骤，调整量化配置可减少临时内存开销：

• 启用双重量化和NF4量化类型，代码示例：

from transformers import BitsAndBytesConfig
import torch

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

2. 缓解内存碎片化 #

根据错误提示的内存碎片问题，设置PyTorch内存分配参数：

• 在代码开头添加环境变量配置：

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

3. 启用梯度检查点 #

牺牲少量计算速度，大幅降低显存占用：

• 加载模型时开启梯度检查点并禁用缓存：

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config,
    gradient_checkpointing=True
)
model.config.use_cache = False

4. 使用LoRA进行参数高效微调 #

避免全量训练，仅训练模型小部分参数：

• 配置并应用LoRA：

from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # 查看可训练参数占比

5. 清理显存与重启Runtime #

• 执行代码清理未使用的显存：

import torch
torch.cuda.empty_cache()

• 若仍有残留占用，可通过Colab「Runtime」→「Restart Runtime」重启环境后重新运行代码

内容的提问来源于stack exchange，提问作者sri_s