如何管理PyTorch显存分配？微调BERT遇CUDA Out of Memory问题

一、直接管理PyTorch显存预留量 #

• 显式限制进程显存占比：PyTorch默认会尽可能占用全部可用显存，你可以通过torch.cuda.set_per_process_memory_fraction(0.8, device=0)手动限制进程最多使用80%的GPU显存（可根据实际情况调整比例，比如0.7或0.9），避免显存被占满触发OOM。注意要在模型加载前调用这个函数。
• 主动释放显存的正确姿势：gc.collect()和torch.cuda.empty_cache()要放在每个epoch结束后或者验证环节结束后调用，不要零散插在训练循环里——频繁调用反而会影响效率，且要确保此时没有未释放的张量引用。

二、排查隐性显存占用 #

• 先查GPU进程：用nvidia-smi命令查看当前GPU的进程占用情况，确认是不是有Paperspace后台服务或者其他无关进程占了显存。如果有，联系平台客服清理或者手动杀掉无关进程。
• 模型加载避坑：不要重复执行model.to('cuda')，也别在循环里多次移动张量到GPU。建议用上下文管理器指定设备：

  with torch.cuda.device('cuda:0'):
      model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased')
  

• 数据集加载优化：别把整个数据集预加载到GPU，而是用DataLoader的pin_memory=True参数，让数据分批从CPU转移到GPU，减少一次性显存占用。

三、进阶显存优化（针对HuggingFace微调） #

• 开启梯度检查点：HuggingFace的BERT模型支持梯度检查点，通过model.gradient_checkpointing_enable()开启，会牺牲少量训练速度来大幅降低显存占用，适合大模型微调场景。
• 验证环节无梯度+分批计算：验证时一定要用torch.no_grad()包裹计算逻辑，同时在Trainer配置里设置eval_accumulation_steps=8（可根据显存调整数值），避免验证时一次性处理太多样本。
• 4bit量化微调：用bitsandbytes库做4bit量化，加载模型时设置load_in_4bit=True，这比fp16更省显存，且A100硬件支持高效计算，几乎不影响微调精度。

四、其他细节排查 #

• 检查张量泄漏：训练循环里别创建不必要的大张量，比如不要在循环内定义torch.zeros()这类大张量，尽量移到循环外复用。
• 确认A100的显存是否被共享：有些云平台的GPU可能是多用户共享的，用nvidia-smi看显存总占用量，如果实际可用远低于80GiB，就联系平台确认资源分配。

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