使用Fairseq训练MNMT模型时遭遇CUDA内存不足错误求助

核心参数调整 #

• 压缩有效batch规模：当前--max-tokens 4096搭配--update-freq 8的有效batch为32768 tokens，远超16GB GPU的承载上限。建议先将--max-tokens降至1024_{2048，`--update-freq`调整为2}4；同时删除--batch-size 64（Fairseq中--max-tokens优先级更高，该参数会造成配置冲突）。
• 更换轻量模型架构：transformer_wmt_en_de是针对英德双语设计的大模型（6层编解码、512嵌入维度），多语言场景下参数膨胀明显。可替换为transformer_iwslt_de_en轻量架构，或手动指定更小的模型尺寸：

  --encoder-layers 4 --decoder-layers 4 --encoder-embed-dim 256 --decoder-embed-dim 256 --encoder-attention-heads 4 --decoder-attention-heads 4
  

• 优化FP16训练配置：添加--fp16-no-flatten-grads参数，避免梯度扁平化带来的额外内存开销；确保使用PyTorch 1.8+版本，旧版本AMP实现存在内存泄漏风险。

内存碎片与系统优化 #

• 配置CUDA内存分配策略：训练前设置环境变量缓解内存碎片：

  export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
  

• 降低验证阶段内存占用：暂时关闭验证（添加--no-validate），待训练稳定后再开启；或单独设置验证集的token上限：--max-tokens-valid 1024。
• 减少Checkpoint留存数量：--keep-interval-updates 20会保留20个完整模型Checkpoint，建议改为--keep-interval-updates 5或--keep-last-epochs 3，减少保存时的临时内存占用。

多语言任务专属优化 #

• 共享嵌入层：添加--share-all-embeddings参数，共享源/目标语言的嵌入层参数，可大幅降低多语言模型的参数总量。
• 简化语言Token策略：当前同时启用--encoder-langtok src和--decoder-langtok，可尝试仅保留--decoder-langtok，减少编码器侧的额外内存开销。
• 截断长序列：预处理阶段添加--max-source-positions 512 --max-target-positions 512，截断超长句子，避免单个batch因长序列导致内存暴增。

其他实用技巧 #

• 启用梯度裁剪：添加--grad-clip-norm 1.0，限制梯度规模，减少梯度缓存的内存占用。
• 清理残留CUDA进程：用nvidia-smi查看GPU占用，杀掉无关进程（kill -9 <PID>），避免内存被残留进程占用。
• 降低保存频率：将--save-interval-updates 5000调整为--save-interval-updates 10000，减少Checkpoint保存次数。

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