在PyTorch中用3090 GPU训练Transformer模型的显存与计算问题

问题1：训练随机冻结、内核无响应及WebSocketClosedError/BufferError报错 #

可能原因 #

• SSH远程连接不稳定：WebSocketClosedError直接关联远程通信链路，比如网络波动、连接超时，或是Jupyter Notebook的WebSocket服务异常。
• 显存接近饱和：3090的24G显存用了18-19G，剩余空间不足，训练中反向传播的临时张量可能触发GPU驱动阻塞，导致进程挂起。
• PyTorch 2.0编译优化冲突：torch.compile可能在某些序列建模场景下存在兼容性bug，引发训练进程冻结。

解决步骤 #

• 优化远程训练方式：
  • 修改本地SSH配置（~/.ssh/config），添加ServerAliveInterval 60和ServerAliveCountMax 3，避免连接超时断开。
  • 放弃Jupyter Notebook，直接在终端后台运行训练脚本：nohup python train.py > train.log 2>&1 &，彻底规避WebSocket通信问题。
• 显存冗余优化：
  • 启用梯度检查点，用torch.utils.checkpoint.checkpoint包装Transformer的编码/解码层，减少反向传播时的显存占用。
  • 适当减小batch size，给显存留足2-3G的冗余空间，避免峰值显存耗尽。
  • 每个epoch结束后，手动删除临时变量（如del outputs, loss），再执行import gc; gc.collect()触发Python垃圾回收，让PyTorch自动释放无用显存。
• 排查PyTorch编译问题：
  • 暂时关闭torch.compile，用普通模式训练，若冻结问题消失，说明是编译优化的bug，升级PyTorch到2.0.1及以上稳定版本。

问题2：调用torch.cuda.empty_cache()引发梯度爆炸、NaN问题 #

原因 #

torch.cuda.empty_cache()会强制释放CUDA冗余显存，但会打乱CUDA内存分配策略，导致后续张量分配、梯度计算出现数值异常；加上原本显存接近饱和，该操作可能引发GPU资源竞争，进一步破坏梯度计算的稳定性。

解决步骤 #

• 彻底移除torch.cuda.empty_cache()调用，改用安全的显存管理：
  • epoch结束后先删除临时变量，再执行gc.collect()，依赖PyTorch的自动显存回收机制即可。
  • 启用自动混合精度训练：用torch.cuda.amp.autocast()包裹前向传播，配合GradScaler处理梯度，既省显存又提升数值稳定性。
• 强化梯度稳定策略：
  • 固定梯度裁剪时机：在反向传播后、优化器更新前执行torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)（max_norm可根据任务调整）。
  • 检查模型初始化：确保Transformer的注意力层、全连接层用Xavier/He初始化，避免参数初始值过大导致梯度爆炸。
  • 调整学习率：降低初始学习率（比如从1e-4调到5e-5），或用ReduceLROnPlateau调度器，在loss异常时自动降学习率。

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