使用PyTorch微调阿拉伯语Whisper模型时遇CUDA内存不足错误

一、设置max_split_size_mb解决内存碎片 #

错误提示里的max_split_size_mb是PyTorch CUDA内存分配器的核心参数，用来控制内存块的拆分阈值，从而减少内存碎片导致的无法分配问题。具体设置方法有两种：

1. 环境变量配置
   在终端运行脚本前直接设置环境变量：

   export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
   

   或者在Python代码开头加入：

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

   数值可根据场景调整（建议从64、128、256逐步测试）：较小值能降低碎片但会增加分配开销，较大值则相反，找到平衡即可。

2. 验证配置生效
   可在代码中加入以下代码确认设置是否生效：

   import torch
   print(torch._C._cuda_get_allocator_config())
   

   输出结果会显示当前的max_split_size_mb参数值。

二、额外内存优化技巧 #

既然已尝试过减小批量大小和基础代码优化，再补充几个实用的进阶方案：

• 启用梯度检查点
  对大模型牺牲计算速度换取内存，PyTorch可通过torch.utils.checkpoint.checkpoint或模型自带的enable_grad_checkpointing()方法实现（如Transformer系列模型）。

• 手动清理内存缓存
  在训练循环的间隙加入以下代码，清理未使用的张量和GPU缓存：

  import gc
  torch.cuda.empty_cache()
  gc.collect()
  

  注意不要过于频繁调用，避免影响训练效率。

• 开启混合精度训练
  使用torch.cuda.amp自动混合精度，大幅降低内存占用：

  from torch.cuda.amp import GradScaler, autocast
  
  scaler = GradScaler()
  for inputs, labels in dataloader:
      optimizer.zero_grad()
      with autocast():
          outputs = model(inputs)
          loss = loss_fn(outputs, labels)
      scaler.scale(loss).backward()
      scaler.step(optimizer)
      scaler.update()
  

• 模型轻量化改造
  若模型规模过大，可考虑使用模型蒸馏后的轻量化版本，或直接减少模型的层数、通道数等参数规模。

• CPU卸载非核心计算
  将部分非核心计算步骤（如某些预处理、辅助逻辑）临时转移到CPU执行，释放GPU内存给核心模型计算。

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