多线程调用PyTorch XTTS模型触发CUDA设备断言错误求助

核心原因分析 #

该错误仅在多线程快速请求时触发，说明问题出在并发场景下的GPU资源竞争、CUDA上下文管理或模型状态共享冲突，单线程下这些问题不会暴露。

针对性解决方案 #

1. 用线程锁保护CUDA操作 #

PyTorch的CUDA张量操作并非完全线程安全，多线程同时调用model.synthesize会导致GPU指令队列混乱，触发设备端断言。

• 实现方式：用threading.Lock包裹计算密集型的合成步骤，确保同一时间只有一个线程访问GPU模型：

  import threading
  
  # 全局锁，保护所有CUDA相关操作
  synthesize_lock = threading.Lock()
  
  def tts_process(text, output_path):
      with synthesize_lock:
          # 执行合成及保存逻辑
          wav = model.synthesize(text, speaker_wav="reference.wav", language="zh")
          torchaudio.save(output_path, wav.unsqueeze(0), sample_rate=24000)
  

2. 显式绑定CUDA上下文到线程 #

快速创建的线程可能未正确初始化CUDA上下文，导致设备端操作异常。

• 实现方式：在每个线程任务内部显式指定GPU设备并初始化上下文：

  def tts_process(text, output_path):
      # 绑定当前线程到指定GPU设备
      with torch.cuda.device("cuda:0"):
          # 确保上下文初始化完成
          if not torch.cuda.is_initialized():
              torch.cuda.init()
          wav = model.synthesize(text, speaker_wav="reference.wav", language="zh")
          torchaudio.save(output_path, wav.unsqueeze(0), sample_rate=24000)
  

3. 统一输入预处理避免形状断言失败 #

即使单请求输入格式正确，并发场景下可能因输入长度差异导致模型内部张量形状检查触发断言。

• 实现方式：在合成前统一预处理文本，确保分词后的张量形状符合模型要求：

  def preprocess_input(text):
      tokens = model.tokenizer(text, return_tensors="pt").input_ids.to("cuda")
      # 动态padding到当前批次最大长度（单请求时可固定为模型支持的最大长度）
      max_len = model.config.max_text_len
      tokens = torch.nn.functional.pad(tokens, (0, max_len - tokens.shape[1]), value=0)
      return tokens
  
  def tts_process(text, output_path):
      with synthesize_lock:
          processed_tokens = preprocess_input(text)
          wav = model.synthesize(processed_tokens, speaker_wav="reference.wav", language="zh")
          torchaudio.save(output_path, wav.unsqueeze(0), sample_rate=24000)
  

4. 主动清理GPU内存避免碎片化 #

快速并发请求会产生大量临时张量，导致GPU内存碎片化，隐性触发内存分配断言。

• 实现方式：在每个任务结束后显式释放内存：

  def tts_process(text, output_path):
      with synthesize_lock:
          try:
              wav = model.synthesize(text, speaker_wav="reference.wav", language="zh")
              torchaudio.save(output_path, wav.unsqueeze(0), sample_rate=24000)
          finally:
              # 删除大张量并清理缓存
              del wav
              torch.cuda.empty_cache()
  

5. 改用进程池替代线程池 #

线程共享同一进程的CUDA上下文，容易引发状态冲突；进程拥有独立的CUDA上下文，兼容性更好。

• 实现方式：用multiprocessing.Pool创建独立进程处理请求（注意GPU内存容量，避免进程过多）：

  from multiprocessing import Pool
  import torch
  import torchaudio
  
  def init_process():
      # 每个进程独立加载模型
      global model
      model = load_xtts_model()  # 替换为你的模型加载逻辑
      model.to("cuda")
  
  def tts_process(text_output_pair):
      text, output_path = text_output_pair
      wav = model.synthesize(text, speaker_wav="reference.wav", language="zh")
      torchaudio.save(output_path, wav.unsqueeze(0), sample_rate=24000)
      return output_path
  
  if __name__ == "__main__":
      # 待处理的请求列表：(文本, 输出路径)
      tasks = [("文本1", "output1.wav"), ("文本2", "output2.wav")]
      # 创建进程池，进程数不超过GPU核心数
      with Pool(processes=2, initializer=init_process) as pool:
          results = pool.map(tts_process, tasks)
  

内容的提问来源于stack exchange，提问作者Özgürcan Karakurt