Colab微调DialoGPT-medium模型触发CUDA OOM报错问题求解

DialoGPT-medium 参数量约355M，当前FP32训练、512序列长度的配置本身就超出16GB显存承载上限，按优先级依次做以下调整即可解决问题：

• 开启FP16混合精度训练
  当前配置中self.fp16 = False是显存占用过高的核心原因之一，P100支持混合精度计算，直接将该参数改为self.fp16 = True，精度等级保持O1即可。该调整可直接将模型参数显存占用降低50%，同时训练速度提升30%以上，是16GB显卡跑medium规模模型的必开配置。
• 下调训练序列长度
  当前self.block_size = 512的配置远高于日常对话训练的需求，绝大多数多轮对话拼接后的token长度在128-256区间，直接将该值调整为self.block_size = 256，如果数据集对话偏短可进一步下调到192。显存占用和序列长度线性正相关，该调整可降低30%-50%的激活值显存占用。
• 调整批次与梯度累积配置
  如果完成前两步调整仍有OOM，将self.per_gpu_train_batch_size = 2改为self.per_gpu_train_batch_size = 1，同时将self.gradient_accumulation_steps = 1改为self.gradient_accumulation_steps = 4。调整后等效全局批次大小仍为4，和初始配置的训练效果完全一致，单步显存占用会进一步下降。
• 开启梯度检查点
  模型加载完成后新增一行代码model.gradient_checkpointing_enable()，该功能通过重新计算中间激活值的方式换显存，可额外降低20%-30%的显存占用，仅会带来15%-20%的训练速度损耗，适合显存缺口不大的场景。
• 其他显存优化细节
  • 模型初始化时传入参数use_cache=False，训练阶段不需要缓存KV注意力矩阵，可避免无意义的显存占用
  • 启动训练前重启Colab内核清空显存残留，不要在训练进程中保留多余的张量、日志变量
  • 若以上调整后仍存在显存缺口，可接入8bit量化加载（依赖bitsandbytes库），将模型参数转为8bit精度存储，可再降低近50%的参数显存占用，精度损失可忽略，16GB显存下甚至可以跑到batch size=2+256序列长度的配置。

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