微调XLNet-Large-Cased时CUDA内存不足求助（batch size=1仍报错）

问题背景 #

我正在使用Hugging Face训练xlnet-large-cased模型，训练参数如下：

args = TrainingArguments( 
    f"xlnet-large-finetuned", 
    evaluation_strategy = "epoch", 
    save_strategy = "epoch", 
    learning_rate=2e-5, 
    per_device_train_batch_size=1, 
    per_device_eval_batch_size=1, 
    num_train_epochs=3, 
    gradient_accumulation_steps=16, 
    weight_decay=0.01, 
    load_best_model_at_end=True, 
    metric_for_best_model="accuracy" 
)

调用代码：

trainer = Trainer( 
    model, 
    args, 
    train_dataset=tokenized_train_dataset, 
    eval_dataset=tokenized_val_dataset, 
    data_collator=data_collator, 
    tokenizer=tokenizer, 
    compute_metrics=compute_metrics 
)
trainer.train()

已将batch size降至1，清空CUDA缓存并通过gc清理所有变量，但仍出现CUDA内存不足错误：

RuntimeError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 15.78 GiB total capacity; 14.31 GiB already allocated; 2.75 MiB free; 14.78 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

无需额外GPU资源的解决办法 #

• 启用梯度检查点
  XLNet支持梯度检查点，通过牺牲少量计算速度降低显存占用。初始化模型后添加：

  model.gradient_checkpointing_enable()
  

  该操作会在反向传播时重新计算部分激活值，大幅减少显存消耗。

• 调整梯度累积与训练评估策略
  当前gradient_accumulation_steps=16，可尝试降至8（保持总batch大小=单设备batch*累积步数的等效性）。同时将evaluation_strategy从"epoch"改为"steps"，并设置较大的eval_steps，减少评估阶段的显存叠加压力。

• 优化PyTorch内存分配
  根据错误提示设置环境变量减少内存碎片，训练脚本开头添加：

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

  或启动脚本时设置：

  export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128"
  

  限制内存块最大分割大小，缓解碎片化导致的显存浪费。

• 启用混合精度训练
  在TrainingArguments中添加fp16=True，将大部分参数和激活值以16位浮点数存储，显存占用直接减半，几乎不影响模型性能：

  args = TrainingArguments(
      # 保留原有参数
      fp16=True
  )
  

• 压缩输入序列长度
  如果当前输入的token序列过长，会显著增加XLNet的显存占用。可以在预处理阶段设置更小的max_length（比如从512降到256，根据任务需求调整），减少单条数据的显存开销。

• 清理GPU冗余进程
  用nvidia-smi命令检查GPU占用情况，关闭无关进程释放显存，确保训练过程独占GPU资源。

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