Google Colab中Llama-3交叉验证GPU内存不足问题求助

针对你在Google Colab中进行Llama-3交叉验证时，第二轮迭代出现GPU内存耗尽的问题，以下是针对性的解决方案：

1. 将基座模型加载移至循环外，仅重置LoRA权重 #

每次循环重新加载完整的Llama-3模型是内存占用累积的核心原因。基座模型只需加载一次，每次交叉验证迭代仅需重置LoRA适配器即可：

# 把模型加载移到循环外面
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map='auto',
    quantization_config=bnb_config,
    token=access_token
)
model.config.use_cache = False
model.config.pretraining_tp = 1

for train_index, test_index in kf.split(data):
    training_data, test_data = data['text'][train_index], data['text'][test_index]
    training_data = Dataset.from_pandas(training_data.to_frame().reset_index())
    test_data = Dataset.from_pandas(test_data.to_frame().reset_index())

    # 重置LoRA权重（第一次训练会自动初始化，后续迭代重置）
    if 'lora' in model.named_modules():
        for name, module in model.named_modules():
            if 'lora' in name:
                module.reset_parameters()

    trainer = SFTTrainer(
        model=model,
        train_dataset=training_data,
        peft_config=lora_config,
        dataset_text_field='text',
        max_seq_length=30,
        tokenizer=tokenizer,
        args=training_arguments,
        packing=True
    )

    trainer.train()

    # 清理当前迭代的Trainer资源
    del trainer
    gc.collect()
    torch.cuda.empty_cache()

2. 完善内存清理逻辑，强制解除对象引用 #

你的free_gpu_cache函数需要补充对大对象的引用删除，否则GC无法彻底回收内存。修改后的清理逻辑应放在每次迭代结束后：

def free_gpu_cache():
    print('Initial GPU Usage')
    gpu_usage()

    # 强制删除可能的大对象引用
    try:
        del trainer
    except NameError:
        pass

    gc.collect()
    torch.cuda.empty_cache()
    
    # 移除numba cuda相关操作（Colab中易引发设备冲突，无必要）
    print("GPU Usage after emptying the cache")
    gpu_usage()

3. 优化模型加载的设备与量化配置 #

避免device_map='auto'在内存碎片下自动将模块分配到CPU/Disk，强制将模型加载到GPU，并优化4bit量化参数进一步节省内存：

# 调整4bit量化配置，开启双量化和嵌套量化
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# 强制所有模块加载到GPU 0，避免自动分配到CPU/Disk
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map={'': 0},
    quantization_config=bnb_config,
    token=access_token
)

4. 调整训练参数减少内存占用 #

在training_arguments中添加以下内存优化参数：

training_arguments = TrainingArguments(
    per_device_train_batch_size=1,  # 调小单设备批量大小
    gradient_accumulation_steps=4,  # 用梯度累积弥补小批量的训练效率
    gradient_checkpointing=True,  # 开启梯度检查点，大幅降低内存占用
    fp16=True,  # 使用混合精度训练
    # 保留你原有的其他参数...
)

关键注意事项 #

• Colab的GPU内存有限（通常16GB左右），Llama-3 7B 4bit量化后约4GB，但训练时的梯度、优化器状态等会额外占用内存，需严格控制批量和序列长度。
• 每次迭代结束后必须删除Trainer对象，否则其内部持有的模型、优化器引用会导致内存无法释放。

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