HuggingFace+FastAPI+Uvicorn多Worker部署时GPU进程异常排查

我使用FastAPI和Uvicorn部署HuggingFace模型进行推理，核心代码如下：

app = FastAPI() 
@app.post("/inference") 
async def func(text:str):     
    output = huggingfacepipeline(text)          
    return ...

启动服务器的命令为：

uvicorn app:app --host 0.0.0.0 --port 8080 --workers 4

服务器配备80GB显存的GPU，预期每个Worker对应一个CPU进程和独立的GPU内存空间，通过nvidia-smi可查看GPU内存分配，应存在4个CPU进程和4个GPU进程。使用小模型（如GPT Neo 125m）时该预期稳定实现，但使用大模型（如16位GPT-J）时行为不可预测：有时4个CPU进程对应仅3个GPU进程（仍有充足显存剩余），有时仅1个GPU进程却对应4个CPU进程。

可能原因 #

• 模型加载时序与Worker启动机制：大模型加载耗时远长于小模型，Uvicorn的Worker进程异步启动。若首个Worker未完成模型加载，后续Worker可能触发Hugging Face Transformers的显存复用机制，共享已加载的模型显存空间，导致多个CPU进程对应单个GPU进程。
• CUDA上下文共享特性：CUDA支持进程间共享上下文，当大模型占用显存较高时，后续Worker进程会尝试依附已存在的CUDA上下文，而非创建新的，从而减少GPU进程数量。
• 模型加载位置错误：若模型在主进程（而非Worker进程）中加载，所有Worker会共享主进程的模型实例，直接导致仅1个GPU进程对应所有CPU Worker。
• Hugging Face Pipeline默认策略：Pipeline默认会优化显存使用，当检测到GPU上已有同模型实例时，会自动复用，而非为每个Worker加载独立副本。

排查与解决步骤 #

• 确认模型加载时机：将模型加载逻辑从全局作用域移至路由函数内部或Worker初始化回调中，确保每个Worker进程独立加载模型。示例：

  from fastapi import FastAPI
  from transformers import pipeline
  
  app = FastAPI()
  pipe = None
  
  @app.on_event("startup")
  async def startup_event():
      global pipe
      # 每个Worker启动时加载独立模型
      pipe = pipeline("text-generation", model="EleutherAI/gpt-j-6B", device=0)
  
  @app.post("/inference")
  async def func(text: str):
      output = pipe(text)
      return output
  

• 添加进程与GPU日志：在模型加载时记录进程PID和GPU设备ID，确认每个Worker是否独立加载模型：

  import os
  import logging
  import torch
  
  logging.basicConfig(level=logging.INFO)
  logger = logging.getLogger(__name__)
  
  @app.on_event("startup")
  async def startup_event():
      global pipe
      logger.info(f"Worker PID: {os.getpid()} - Starting model load")
      pipe = pipeline(...)
      logger.info(f"Worker PID: {os.getpid()} - Model loaded on GPU {torch.cuda.current_device()}")
  

• 强制禁用CUDA上下文共享：在代码开头添加以下配置，强制每个Worker使用独立的CUDA上下文：

  import torch
  torch.multiprocessing.set_start_method('spawn', force=True)
  

• 实时监控显存变化：使用watch -n 1 nvidia-smi实时观察Worker启动过程中的显存分配，确认是否存在模型复用或加载失败的情况。
• 调整Uvicorn启动参数：移除--preload参数（若使用），避免主进程预加载模型导致Worker共享；尝试降低Worker数量至2个，验证是否能稳定对应GPU进程数，逐步排查资源竞争问题。

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