Llama2-7B模型CUDA内存不足问题排查（Colab A100环境）

核心问题根源 #

你遇到的CUDA显存不足问题，主要来自模型加载无显存优化、生成阶段未限制资源使用，以及代码逻辑的低效处理。

具体问题与修正方案 #

1. 模型加载未做显存优化

Llama2-7B全精度（float32）占用约28GB显存，即使A100显存充足，直接加载也会预留大量空间，导致后续生成操作溢出。必须启用显存优化策略：

import torch

# 方案1：用bf16半精度加载（显存占用约14GB）
model = LlamaForCausalLM.from_pretrained(path, torch_dtype=torch.bfloat16).to("cuda")

# 方案2：启用4位量化（需安装bitsandbytes库，显存占用约4GB）
model = LlamaForCausalLM.from_pretrained(path, load_in_4bit=True, device_map="auto")

2. 生成阶段未限制显存消耗

model.generate()默认参数会无限制生成token，且未关闭不必要的采样逻辑，进一步加剧显存占用。需添加显存友好的参数：

generate_ids = model.generate(
    **inputs,
    max_new_tokens=100,  # 限制生成的最大token数，根据需求调整
    do_sample=False,     # 关闭采样，减少显存占用
    pad_token_id=tokenizer.eos_token_id  # 避免pad token报错
)

每次生成后手动清理显存，释放临时张量：

torch.cuda.empty_cache()

3. 代码语法错误

生成的结果变量是generate_ids，但解码时误用了未定义的outputs，修正后：

output = tokenizer.batch_decode(generate_ids)[0]

4. 假批次处理逻辑

你拆分了50条的子批次，但process_data是逐条循环处理，没有实现真正的批量推理，反而因循环中未及时释放显存导致累积占用。改为批量处理逻辑：

def process_data(query, batch):
    responses = []
    valid_prompts = []
    # 先收集非NaN的文本
    for text in batch:
        if pd.isna(text):
            responses.append("NaN")
        else:
            valid_prompts.append(query + text)
    
    # 批量tokenize并推理
    if valid_prompts:
        inputs = tokenizer(
            valid_prompts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512  # 限制输入长度，避免显存溢出
        ).to("cuda")
        generate_ids = model.generate(
            **inputs,
            max_new_tokens=100,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id
        )
        valid_outputs = tokenizer.batch_decode(generate_ids, skip_special_tokens=True)
        
        # 将结果填充回responses列表
        idx = 0
        for i in range(len(responses)):
            if responses[i] == "NaN":
                continue
            responses[i] = valid_outputs[idx]
            idx += 1
    
    return responses

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