6GB显存GPU遇PyTorch CUDA内存不足问题，急求解决

文本任务专属优化 #

• 极致压缩批量大小与序列长度：从batch_size=1或2开始测试，同时将tokenizer的max_length设为模型有效运行的最小长度（比如BERT类设为128而非512），长文本截断能大幅降低单样本显存占用。
• 梯度累积替代大批次：若需保证训练稳定性，用梯度累积模拟大批次。例如目标等效batch_size=16，可设batch_size=2，每8步执行一次优化，代码示例：

  accumulation_steps = 8
  for step, batch in enumerate(dataloader):
      outputs = model(**batch)
      loss = outputs.loss / accumulation_steps
      loss.backward()
      if (step + 1) % accumulation_steps == 0:
          optimizer.step()
          optimizer.zero_grad()
  

• 混合精度训练：启用FP16混合精度，直接砍半显存占用，代码示例：

  scaler = torch.cuda.amp.GradScaler()
  for batch in dataloader:
      with torch.cuda.amp.autocast():
          outputs = model(**batch)
          loss = outputs.loss
      scaler.scale(loss).backward()
      scaler.step(optimizer)
      scaler.update()
      optimizer.zero_grad()
  

• 换轻量级模型：用DistilBERT、TinyBERT等替代标准大模型，参数量仅为原模型的1/3左右，文本任务精度损失可控。

显存碎片与PyTorch配置调整 #

• 设置显存分配参数：按照错误提示配置PYTORCH_CUDA_ALLOC_CONF减少碎片，代码开头添加：

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

  可在64-256区间调整数值，找到最优值。
• 启用按需显存分配：关闭PyTorch默认的大显存预分配，改为按需扩展：

  import os
  os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
  

• 定时清理缓存：每个epoch结束后手动清理显存碎片：

  torch.cuda.empty_cache()
  

其他排查项 #

• 用nvidia-smi检查GPU进程，关闭其他占用显存的程序。
• 确认模型、数据均已正确移至GPU（model.to('cuda')、batch.to('cuda')），避免CPU/GPU数据混存。

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