在多卡环境下基于 Wan2.1（文生视频模型）对外提供 API 接口时，核心诉求是启动服务时一次性加载模型到多卡，调用接口仅执行推理（避免重复加载）。torchrun本身是分布式启动工具，需结合「服务化框架 + 分布式模型预加载」实现该目标，以下是完整的落地方案（兼容 PyTorch 分布式、支持多卡推理）：

核心思路

用 torchrun 启动分布式服务进程：在服务启动阶段，通过torch.distributed完成多卡模型初始化，将模型权重加载到所有 GPU 并保持显存常驻；
封装推理逻辑为 API 接口：基于 FastAPI/Flask 搭建 HTTP 服务，接口调用时仅传递输入数据、触发推理计算，不重新加载模型；
分布式推理适配：确保多卡环境下的输入数据分发、推理结果聚合逻辑与预加载的模型兼容。

方案一：FastAPI + torch.distributed（推荐，高性能）

步骤 1：环境准备

bash

运行

# 安装依赖
pip install fastapi uvicorn torch transformers accelerate pydantic

步骤 2：编写服务代码（wan2.1_api.py）

python

运行

import os
import torch
import uvicorn
from fastapi import FastAPI, UploadFile, File
from pydantic import BaseModel
from typing import Optional
import torch.distributed as dist
from transformers import AutoTokenizer, AutoModelForCausalLM  # 替换为Wan2.1实际加载逻辑

# ====================== 1. 分布式初始化（torchrun启动时自动注入环境变量） ======================
def init_distributed():
    # 从torchrun的环境变量中读取分布式配置
    rank = int(os.environ.get("RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    
    # 初始化分布式进程组（NCCL后端，多卡必须用NCCL）
    torch.cuda.set_device(local_rank)
    dist.init_process_group(
        backend="nccl",
        init_method="env://",  # 依赖torchrun的环境变量
        rank=rank,
        world_size=world_size
    )
    return rank, local_rank, world_size

# ====================== 2. 预加载模型（服务启动时仅执行一次） ======================
class Wan21InferenceService:
    def __init__(self, local_rank: int):
        self.local_rank = local_rank
        self.device = torch.device(f"cuda:{local_rank}")
        # ========== 替换为Wan2.1实际的模型加载逻辑 ==========
        # 示例：加载Wan2.1模型（多卡分布式加载，需适配模型的parallel方式）
        self.model_path = "/path/to/wan2.1-model"  # 内网模型路径
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
        # 多卡加载：使用device_map或分布式并行（根据Wan2.1的实现调整）
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_path,
            torch_dtype=torch.float16,
            device_map={"": self.local_rank},  # 每个进程加载到对应GPU
            trust_remote_code=True
        )
        self.model.eval()  # 推理模式，禁用梯度
        
    @torch.no_grad()
    def infer(self, prompt: str, video_length: int = 16) -> str:
        """核心推理逻辑：仅执行前向计算，不重新加载模型"""
        # 1. 输入预处理（分发到当前GPU）
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
        
        # 2. 多卡推理（根据Wan2.1的分布式逻辑调整，示例为基础前向）
        outputs = self.model.generate(
            **inputs,
            max_new_tokens=video_length,
            do_sample=True,
            temperature=0.7
        )
        
        # 3. 结果聚合（若多卡输出分片，需跨进程聚合，示例简化）
        result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return result

# ====================== 3. FastAPI服务封装 ======================
app = FastAPI(title="Wan2.1 文生视频API")
service: Optional[Wan21InferenceService] = None  # 全局服务实例（预加载模型）

# 输入参数模型
class InferenceRequest(BaseModel):
    prompt: str
    video_length: Optional[int] = 16

# 仅主进程（rank=0）启动API服务，避免多进程端口冲突
@app.post("/generate_video")
async def generate_video(request: InferenceRequest):
    global service
    # 主进程接收请求，分发到各卡推理（简化版，实际可优化为分布式数据并行）
    result = service.infer(request.prompt, request.video_length)
    return {
        "code": 200,
        "message": "success",
        "data": {
            "prompt": request.prompt,
            "video_text_desc": result,  # 替换为实际视频生成结果（如帧数据/视频路径）
            "gpu_rank": service.local_rank
        }
    }

# ====================== 4. 服务启动入口 ======================
if __name__ == "__main__":
    # 初始化分布式
    rank, local_rank, world_size = init_distributed()
    
    # 所有进程预加载模型（每个进程加载到对应GPU）
    service = Wan21InferenceService(local_rank)
    print(f"✅ Rank {rank} (GPU {local_rank}) 模型加载完成，显存常驻")
    
    # 仅主进程（rank=0）启动HTTP服务，避免多进程监听同一端口
    if rank == 0:
        uvicorn.run(
            app="__main__:app",
            host="0.0.0.0",  # 对外提供接口
            port=8000,
            workers=1,  # 分布式下仅1个worker，避免模型重复加载
            reload=False  # 禁用热重载，防止模型重新加载
        )
    else:
        # 非主进程保持存活，等待主进程分发任务（阻塞）
        while True:
            torch.cuda._sleep(1000000)

步骤 3：用 torchrun 启动多卡服务

bash

运行

# 假设用4卡（--nproc_per_node指定卡数），启动时自动加载模型到所有卡
torchrun --nproc_per_node=4 --master_port=29500 wan2.1_api.py

--nproc_per_node=4：指定使用 4 张 GPU；
--master_port=29500：分布式通信端口，避免冲突；
启动后日志会显示「Rank 0/1/2/3 模型加载完成」，模型权重常驻所有 GPU 显存，无重复加载。

步骤 4：调用 API 接口（仅推理，不加载模型）

bash

运行

# POST请求示例（curl）
curl -X POST "http://服务器IP:8000/generate_video" \
     -H "Content-Type: application/json" \
     -d '{"prompt": "生成一段海边日落的视频，时长10秒", "video_length": 16}'

返回示例：

json

{
  "code": 200,
  "message": "success",
  "data": {
    "prompt": "生成一段海边日落的视频，时长10秒",
    "video_text_desc": "（视频生成结果）",
    "gpu_rank": 0
  }
}

方案二：Flask + 分布式进程守护（兼容传统 API 框架）

若偏好 Flask，核心逻辑与 FastAPI 一致，仅替换 API 层封装：

服务代码（wan2.1_flask_api.py）

python

运行

import os
import torch
from flask import Flask, request, jsonify
import torch.distributed as dist

# 1. 分布式初始化（同方案一）
def init_distributed():
    rank = int(os.environ.get("RANK", 0))
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl", init_method="env://", rank=rank, world_size=world_size)
    return rank, local_rank, world_size

# 2. 预加载模型（同方案一）
class Wan21InferenceService:
    def __init__(self, local_rank):
        self.local_rank = local_rank
        self.device = torch.device(f"cuda:{local_rank}")
        # 加载Wan2.1模型（省略，同方案一）
        self.model = ... 
        self.tokenizer = ...

    def infer(self, prompt, video_length):
        # 推理逻辑（同方案一）
        ...

# 3. Flask服务
app = Flask(__name__)
service = None

@app.route("/generate_video", methods=["POST"])
def generate_video():
    data = request.get_json()
    prompt = data.get("prompt")
    video_length = data.get("video_length", 16)
    result = service.infer(prompt, video_length)
    return jsonify({
        "code": 200,
        "data": {"result": result}
    })

if __name__ == "__main__":
    rank, local_rank, world_size = init_distributed()
    service = Wan21InferenceService(local_rank)
    print(f"✅ Rank {rank} 模型加载完成")
    if rank == 0:
        app.run(host="0.0.0.0", port=8000, debug=False)  # 仅主进程启动服务
    else:
        while True:
            torch.cuda._sleep(1000000)

启动命令

bash

运行

torchrun --nproc_per_node=4 --master_port=29501 wan2.1_flask_api.py

关键优化点（解决多卡推理核心问题）

避免模型重复加载：
- 模型初始化逻辑放在__main__中，仅在torchrun启动进程时执行一次；
- 非主进程（rank≠0）不启动 HTTP 服务，仅保持进程存活，等待主进程分发推理任务。

多卡推理数据分发：若 Wan2.1 需多卡并行推理（如张量并行 / 流水线并行），需在infer方法中添加分布式数据分发逻辑：

python

运行

@torch.no_grad()
def infer(self, prompt: str, video_length: int = 16) -> str:
    # 1. 主进程（rank=0）预处理输入，广播到所有卡
    inputs = None
    if dist.get_rank() == 0:
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
    dist.broadcast(inputs, src=0)  # 主进程将输入广播到所有卡
    
    # 2. 各卡独立推理
    outputs = self.model.generate(**inputs, max_new_tokens=video_length)
    
    # 3. 聚合结果（仅主进程返回）
    result = None
    if dist.get_rank() == 0:
        result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result

显存优化：

模型加载时使用torch.float16/bfloat16降低显存占用；
禁用梯度计算（torch.no_grad()）、关闭模型的train()模式；

若显存不足，使用accelerate的init_empty_weights + load_checkpoint_and_dispatch实现分片加载：

python

运行

from accelerate import init_empty_weights, load_checkpoint_and_dispatch
with init_empty_weights():
    self.model = AutoModelForCausalLM.from_pretrained(self.model_path, trust_remote_code=True)
self.model = load_checkpoint_and_dispatch(
    self.model,
    checkpoint=self.model_path,
    device_map={"": self.local_rank},
    dtype=torch.float16
)

服务运维与监控

进程守护：使用supervisor/systemd管理torchrun进程，避免服务意外退出：

ini

# /etc/supervisor/conf.d/wan21_api.conf
[program:wan21_api]
command=torchrun --nproc_per_node=4 --master_port=29500 /path/to/wan2.1_api.py
directory=/path/to/project
user=root
autostart=true
autorestart=true
stderr_logfile=/var/log/wan21_api.err.log
stdout_logfile=/var/log/wan21_api.out.log

显存监控：启动后通过nvidia-smi确认模型已加载且显存常驻：

bash

运行
```
watch -n 1 nvidia-smi
```
正常情况下，所有指定 GPU 的显存占用会稳定在模型加载后的数值，接口调用时仅小幅波动（推理计算占用）。
接口性能优化：
- 使用异步接口（FastAPI 的async）避免请求阻塞；
- 开启请求队列（如uvicorn的--workers=1 --limit-concurrency=10），控制并发推理数，防止多卡过载。

常见问题解决

端口冲突：torchrun的--master_port需唯一，避免与其他分布式服务冲突；
多卡通信失败：确保服务器开启 NCCL 通信（需 GPU 支持、关闭防火墙）；
模型加载 OOM：降低--nproc_per_node卡数、使用模型量化（INT8/INT4）、开启显存复用（export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True）；
接口调用无响应：检查非主进程是否阻塞（需保持while True循环）、分布式进程组是否初始化成功。

该方案可确保服务启动时一次性加载模型到多卡，接口调用仅执行推理，完全解决torchrun每次加载模型的问题，且兼容 Wan2.1 的多卡分布式推理需求。