48GB GPU上微调SBERT（TSDAE）反复训练后遇CUDA内存不足

租用了48GB显存的RTX A6000 GPU，用TSDAE微调Hugging Face的SBERT。采用分块读取训练文本、每次取1000句训练的方式，训练后保存模型，下一轮加载模型继续训练。前5-6次训练正常，但后续即使调用torch.cuda.empty_cache()清理缓存，仍出现CUDA内存不足错误。当前batch size设为16，尝试过多种设置仍无法解决。

相关代码 #

def train_denoising(train_sentences,modelName):

  torch.cuda.empty_cache()
  # os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:22"
  os.environ["CUDA_LAUNCH_BLOCKING"] = "1"


  word_embedding_model = models.Transformer(modelName)
  # Apply **cls** pooling to get one fixed sized sentence vector
  pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), 'cls')
  model = SentenceTransformer(modules=[word_embedding_model, pooling_model])

  #model = SentenceTransformer(modelName)

  
  train_dataset = datasets.DenoisingAutoEncoderDataset(train_sentences)
  train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
  train_loss = losses.DenoisingAutoEncoderLoss(model, decoder_name_or_path=modelName, tie_encoder_decoder=True)

  model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=num_epochs,
    weight_decay=0,
    scheduler='constantlr',
    optimizer_params={'lr': 3e-5},
    show_progress_bar=True,
    # checkpoint_path=model_output_path,
    use_amp=False,  # Set to True, if your GPU supports FP16 cores
    output_path='./yobb_model')
  return model

num_sentences = 1000

for path in paths:
 with gzip.open(path, 'rt', encoding='utf8') if path.endswith('.gz') else open(path, encoding='utf8') as f:
 
    for piece in read_in_chunks(f, chunk_size=500*1024): 
        aux = [line.lower() for line in splitter.split(piece) if len(line) > 10]

        count = len(aux)//num_sentences
        index = 0
        # iterate over the sentences getting <num_sentences> each time
        for i in range(count):                            
            train_sentences.extend(aux[index:index+num_sentences])  
            index += num_sentences
            if(len(train_sentences) <= 0):
                continue
            print("Numero de sentencas {}".format(len(aux)))
            logging.info("{} train sentences".format(len(train_sentences)))
            train_denoising(train_sentences,"./yobb_model")
            train_sentences.clear()

        count = len(aux) % num_sentences
        if(count > 0):                
            print("Numero de sentencas {}".format(len(aux[-count:])))
            logging.info("{} train sentences".format(len(aux[-count:])))
            train_denoising(aux[-count:],"./yobb_model")

报错信息 #

OutOfMemoryError: CUDA out of memory. Tried to allocate 920.00 MiB (GPU 0; 47.54 GiB total capacity; 43.35 GiB already allocated; 517.88 MiB free; 46.65 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

核心问题原因 #

1. 旧模型实例未被彻底释放：每次调用train_denoising都会重新创建SentenceTransformer实例，旧模型的GPU张量因Python垃圾回收延迟清理，导致显存持续占用。torch.cuda.empty_cache()只能释放PyTorch已放弃引用的显存，对未被回收的模型实例无效。
2. TSDAE组件显存累积：DenoisingAutoEncoderLoss会加载Decoder模型，训练产生的中间张量、优化器状态等未完全清理，多次迭代后显存碎片和占用量逐步上升。
3. 循环变量无显式销毁：外层循环中，train_denoising返回的模型未被处理，GPU上的参数一直被占用，没有主动释放。

具体解决步骤 #

• 显式清理模型并触发垃圾回收
  训练完成后，主动删除模型实例，触发Python垃圾回收，再清理显存。修改调用代码：

  model = train_denoising(train_sentences,"./yobb_model")
  # 清理模型显存
  del model
  import gc
  gc.collect()
  torch.cuda.empty_cache()
  train_sentences.clear()
  

  同时在train_denoising函数返回前，把模型移到CPU，减少GPU占用：

  model.to('cpu')
  return model
  

• 优化模型加载逻辑
  不要每次训练都重新初始化模型，直接加载本地已保存的模型，避免重复初始化的显存浪费。修改train_denoising函数：

  def train_denoising(train_sentences, model_path):
      torch.cuda.empty_cache()
      os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
      
      # 直接加载本地保存的模型
      model = SentenceTransformer(model_path)
      
      train_dataset = datasets.DenoisingAutoEncoderDataset(train_sentences)
      train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
      train_loss = losses.DenoisingAutoEncoderLoss(model, decoder_name_or_path=model_path, tie_encoder_decoder=True)
      
      model.fit(
          train_objectives=[(train_dataloader, train_loss)],
          epochs=num_epochs,
          weight_decay=0,
          scheduler='constantlr',
          optimizer_params={'lr': 3e-5},
          show_progress_bar=True,
          use_amp=False,
          output_path='./yobb_model')
      return model
  

• 启用混合精度训练
  RTX A6000支持FP16，将model.fit中的use_amp=True，可大幅降低显存占用，且不影响训练精度。

• 缓解显存碎片化
  取消注释os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:22"，让PyTorch更高效分配显存，减少碎片堆积。

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