PyTorch预留显存远超实际需求，微调SentenceTransformer遇OOM问题

问题本质 #

PyTorch默认的显存分配策略会预分配几乎全部可用GPU显存，这就导致你看到torch.cuda.memory_reserved()远大于实际模型参数占用的情况——1.3GB的模型却占了13GB预留空间，剩余显存不足以加载批次数据，最终触发OOM。

实用解决方法 #

1. 开启显存按需分配模式 #

在代码最开头添加以下配置，让PyTorch仅在需要时分配显存，而非预分配大段空间：

import os
import torch

# 启用可扩展显存分段（PyTorch 1.10+支持）
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

# 限制进程显存使用率（可选，比如限制到80%）
torch.cuda.set_per_process_memory_fraction(0.8)
torch.cuda.empty_cache()

这个设置会彻底改变显存分配逻辑，避免一开始就占满显存。

2. 优化SentenceTransformer模型初始化 #

给Transformer模块添加显存优化参数，同时启用梯度检查点以牺牲少量计算速度换显存：

word_embedding_model = models.Transformer(
    model_name, 
    max_seq_length=512,
    low_cpu_mem_usage=True,  # 减少CPU到GPU加载时的显存占用
    gradient_checkpointing=True  # 不存储全部中间激活值，大幅降低训练显存占用
)

3. 调整训练参数减少显存压力 #

• 缩小批次大小：直接降低训练的batch_size，比如从32降到16或8，根据显存剩余空间调整；
• 启用混合精度训练：自动将部分张量转为FP16精度，节省近一半显存：
  如果用SentenceTransformer自带的Trainer，只需在训练参数中设置use_amp=True；
  手动训练则可以用以下代码：

  from torch.cuda.amp import autocast, GradScaler
  
  scaler = GradScaler()
  optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
  
  for batch in train_dataloader:
      with autocast():
          embeddings = model(batch["text"])
          loss = your_loss_function(embeddings, batch["labels"])
      
      scaler.scale(loss).backward()
      scaler.step(optimizer)
      scaler.update()
      optimizer.zero_grad()
  

4. 清理无关显存占用 #

训练前执行nvidia-smi命令检查GPU显存，确认没有其他后台进程占用资源；每次模型初始化前调用torch.cuda.empty_cache()清理残留的未释放显存。

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