RuntimeError: Worker failed with error ‘CUDA driver error: out of memory’, please check the stack trace above for the root cause

17次阅读

日志问题完整分析与解决方案

一、根因定位

日志末尾核心报错：

plaintext

RuntimeError: CUDA driver error: out of memory

问题场景：Wan2.2 视频生成流程，执行 VAE 解码（vae.decode）阶段显存溢出，堆栈卡在自定义 Norm 层 F.normalize 计算。

业务链路：API Server → RPC 分布式 Worker → Diffusion Pipeline → Wan2.2 VAE Decoder 显存打爆，请求直接终止。

关键特征

溢出点：VAE 解码器逐帧解码 x[:, :, i : i + 1, :, :] 上采样 ResNet Norm 层；
架构：vllm-omni 分布式多进程 RPC 推理，feat_cache 缓存特征占用大量显存；
模型：Wan2.2 原生 VAE 本身显存开销极大，视频分辨率、帧数、批量是主要膨胀点。

二、分级修复方案（从快速临时根治到长期架构优化）

方案 1：紧急临时规避（立刻可用，无需改代码）

1）降低视频生成参数（最有效）

下发请求时强制缩小显存占用配置：

json

{
  "height": 480,
  "width": 720,
  "num_frames": 16,
  "num_inference_steps": 20,
  "batch_size": 1,
  "fp8": true
}

禁止 720P+、32 帧以上、batch>1 同时使用；关闭高清增强、高分插值。

2）释放环境残留显存

WSL 环境容易显存缓存不释放，重启前执行：

bash

运行

# 清空cuda缓存
export PYTHONFAULTHANDLER=1
export CUDA_VISIBLE_DEVICES=0
# 启动前强制垃圾回收
python -c "import torch; torch.cuda.empty_cache()"

彻底 kill 残留进程：

bash

运行

fuser -v /dev/nvidia*
kill -9 对应PID

3）启动参数加显存限制开关

启动 api_server 追加参数：

bash

运行

--diffusion-max-batch-size 1 \
--cuda-malloc-conf rounding_up_power2 \
--pipeline-fp8 \
--disable-feat-cache

--disable-feat-cache 关闭日志中 feat_cache=self._feat_map 特征缓存，直接削减 VAE 显存占用。

方案 2：代码层显存优化（稳定根治，修改配置 / 层逻辑）

1）关闭 Wan VAE 特征缓存

文件：vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

找到调用 self.decoder 传参 feat_cache=self._feat_map 位置，全局禁用缓存：

python

运行

# 修改前
out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
# 修改后
out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=None, feat_idx=self._conv_idx)

缓存会保存所有上采样中间特征，是显存溢出核心元凶之一。

2）自定义 Norm 层强制混合精度 + 分块计算

溢出文件：vllm_omni/diffusion/layers/norm.py

forward_native 中 F.normalize 增加精度截断，使用 torch.float16/float8，超大张量分块归一化，避免一次性加载完整张量进显存。

示例改造：

python

运行

# forward_native 内部
x = x.half()
chunk = 128
for i in range(0, x.shape[0], chunk):
    x[i:i+chunk] = F.normalize(x[i:i+chunk], dim=1)

3）VAE 解码逐帧及时释放张量

autoencoder_kl_wan.py 的 _decode 循环内，每帧解码完成后手动释放中间变量：

python

运行

out_ = self.decoder(...)
torch.cuda.empty_cache()
del x_chunk

方案 3：分布式 Worker 调度优化（多请求并发场景）

日志出现 RPC call failed 分布式报错，多 Worker 争抢显存：

单卡仅启动 1 个 diffusion worker，禁止多进程分摊；
增加请求队列限流，同一时间只处理 1 条视频任务；
开启 Worker 显存隔离，任务结束自动重启子进程回收显存；
开启梯度检查点 gradient_checkpointing=True 加载 Wan2.2 VAE & Transformer。

方案 4：量化降级（硬件显存不足终极方案）

全局启用 FP8 量化加载 Wan2.2 全套模型（UNet+VAE）；
关闭 VAE float32 高精度解码，强制 fp16；
更换轻量版 Wan2.1 替代 Wan2.2，同等分辨率显存降低 40%+。

三、排查定位命令（确认显存瓶颈）

启动服务时埋显存监控，复现前观测占用：

bash

运行

# 实时查看显存
watch -n 0.5 nvidia-smi
# 代码内埋点，打印VAE解码前后显存
print("VAE decode pre mem:", torch.cuda.memory_allocated() / 1024**2)
# 解码逻辑执行后
print("VAE decode post mem:", torch.cuda.memory_allocated() / 1024**2)

若解码后显存暴涨 3~6GB，确认是 VAE 上采样 + 特征缓存导致。