使用WandB.watch记录梯度触发CUDA显存溢出的原因及优化方法

• 梯度临时副本开销：wandb.watch会为所有模型参数注册反向传播钩子，每次触发日志时会读取全量梯度数据，在GPU上执行展平、NaN/Inf过滤、统计量计算等操作，布尔索引、逻辑判断步骤都会产生新的GPU临时张量副本。从报错栈可以看出，flat = flat[~torch.isinf(flat)]这步尝试分配4.65GB的临时张量，刚好超出剩余显存容量。
• 全量日志带来双倍开销：你使用了log='all'配置，会同时采集参数值和梯度两类数据做统计，处理开销是仅采集梯度的两倍。
• 高频日志放大显存压力：log_freq=3意味着每3个step就执行一次全量梯度/参数的统计处理，临时张量产生频率过高，PyTorch缓存分配器来不及释放闲置显存，容易触发峰值溢出。

• 调整wandb.watch配置
  • 把log='all'改为log='gradients'，不需要监控参数值分布的情况下直接砍掉一半数据处理量；如果不需要梯度分布直方图，仅需要监控梯度范数的话可以加log_graph=False, log_weights=False进一步裁剪功能。
  • 调大log_freq数值，比如改为10、50甚至每1个epoch记录一次，降低统计操作的触发频率，给显存回收留足时间。
• 把梯度统计转移到CPU执行：可以给WandB的张量统计逻辑打补丁，处理梯度前先把张量移动到CPU，所有临时张量都占用CPU内存，不会挤占GPU训练显存，不需要修改WandB源码也能实现。
• 优化训练侧显存冗余：你现有训练代码中comb = torch.zeros(1,1,100,1632).to(device)可以提到训练循环外复用，避免每次迭代都重新申请释放显存；不需要参与反向传播的张量及时调用.detach()，减少不必要的显存占用，为WandB操作预留空间。
• 降精度处理统计数据：如果需要保留高频日志，可以将梯度转换为float16精度后再传入统计逻辑，临时张量的显存占用直接降低一半。

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