Wan2.2-S2V-14B 24GB GPU 推理优化报告

1. 项目背景

目标: 在单张 RTX 4090 D (24GB VRAM) 上运行 Wan2.2-S2V-14B (14B 参数) 语音驱动视频生成模型，使用 68 秒音频生成完整视频。

挑战: 模型 bfloat16 精度下需要 ~28GB 显存，超出 24GB 硬件限制。官方推荐最低 80GB VRAM。

硬件环境:

GPU: NVIDIA GeForce RTX 4090 D, 24GB VRAM
CPU: 服务器级 CPU, 充足内存
PCIe: 4.0 x16 (~25GB/s 带宽)
OS: Windows Server, Python 3.12, PyTorch 2.11+cu128
无 flash_attn (Windows 不支持)

模型组件:

组件	参数量	bfloat16 大小	用途
DiT (去噪网络)	14B (40 blocks)	~28GB	核心扩散去噪
T5 (文本编码器)	~4.7B	~8GB	文本理解
VAE (视频编解码)	~350M	~1.2GB	像素↔潜空间
wav2vec2 (音频编码)	~315M	~1.2GB	音频特征提取

2. 优化迭代过程

2.1 v1: accelerate.cpu_offload (基线)

方案 : 使用 accelerate.cpu_offload() 将整个 DiT 模型放在 CPU，每个 transformer block 在前向传播时临时搬到 GPU 计算，完后搬回。

实现:

python 复制代码

from accelerate import cpu_offload
self.noise_model = cpu_offload(self.noise_model, execution_device=self.device)

结果:

每步耗时: 48s/step
每 clip: ~16 min
GPU 利用率: < 50% (每次仅 1 个 block 在 GPU, ~703MB)
68s 音频: 8 小时超时未完成 (完成 20/23 clips)

问题: 所有 40 个 block 都从 CPU 搬运，GPU 大量时间空闲等待 PCIe 传输。显存严重浪费。

2.2 v2: device_map="auto" (失败)

方案 : 使用 diffusers 的 device_map="auto" 自动分配模型到 GPU/CPU。

实现:

python 复制代码

self.noise_model = WanModel_S2V.from_pretrained(
    checkpoint_dir, torch_dtype=self.param_dtype, device_map="auto")

结果:

29 个模块在 GPU, 21 个在 CPU
每步耗时: >10 min/step (极慢，不可用)
2 小时仅完成不到 1 clip

失败原因 : dispatch_model 将 CPU 模块的计算也放在 CPU 上执行 (CPU 矩阵乘法比 GPU 慢数十倍)，而非搬到 GPU 计算。这是与 cpu_offload 的本质区别。

教训 : device_map="auto" 适用于 LLM 文本生成 (小 batch, 序列化推理)，不适用于扩散模型 (大 batch attention, 需要 GPU 并行)。

2.3 v3: 混合常驻 + CPU Offload

方案 : 将尽可能多的 block 常驻 GPU (零传输开销)，剩余 block 用 AlignDevicesHook(offload=True) 做 CPU↔GPU 穿梭 (计算始终在 GPU)。

实现:

python 复制代码

# 计算 GPU 可容纳的 block 数
gpu_avail = gpu_total - gpu_used - 12 * 1024**3  # 12GB 预留
n_gpu = int((gpu_avail - non_block_size) / block_size)

# 前 N 个 block 常驻 GPU
for i in range(n_gpu):
    model.blocks[i].to(self.device)

# 剩余 block 挂 offload hook
for i in range(n_gpu, len(model.blocks)):
    hook = AlignDevicesHook(execution_device=self.device, 
                            offload=True, place_submodules=True)
    add_hook_to_module(model.blocks[i], hook)

关键发现 - place_submodules=True 必须设置:

默认 place_submodules=False 只搬运 block 的直接参数
子模块 (self_attn.q, self_attn.k 等) 的权重不会被搬运
导致 RuntimeError: mat1 is on cuda:0, different from other tensors on cpu
设置 True 后递归搬运所有子模块参数

预留空间的探索:

预留	GPU blocks	结果
4GB	23/40	OOM, GPU 无响应 (激活值空间不足)
6GB	21/40	OOM, VAE 编码崩溃
8GB	17/40	VAE 编码崩溃 (VAE + DiT 共享 GPU)
12GB	11/40	稳定运行, 26s/step

结果:

11/40 blocks 常驻 GPU, 29 个用 offload hook
每步耗时: 26s/step (比 v1 快 46%)
GPU 利用率: ~70%+
1 clip: ~10 min

教训: 480×832 分辨率下, attention 激活值需要 ~11GB 显存。预留不足会导致显存抖动, 速度反而更慢。

2.4 v4: 顺序独占 GPU (架构优化)

动机: v3 中 VAE (~1.2GB) 常驻 GPU, 挤占了 DiT 的空间。理想情况下每个阶段应独占全部 GPU。

方案:

VAE 加载到 CPU, 按需搬入 GPU
DiT 每个 clip 去噪前 _setup_dit(), 去噪后 _teardown_dit()
每个阶段独占 GPU, 用完释放

VAE 搬运辅助方法 (Wan2_1_VAE 不是 nn.Module):

python 复制代码

def _move_vae(self, device):
    self.vae.model.to(device)
    self.vae.mean = self.vae.mean.to(device)
    self.vae.std = self.vae.std.to(device)
    self.vae.scale = [self.vae.mean, 1.0 / self.vae.std]
    self.vae.device = device

结果:

13/40 blocks on GPU (比 v3 多 2 个, 因为 VAE 不占 GPU 了)
每步耗时: 26s/step
但每个 clip 有 ~3min 的 setup/teardown 开销

问题: 每个 clip 都要拆建一次 DiT (搬入搬出 ~9GB), 23 clips 增加了大量开销。

2.5 v5: 加载一次 + VAE 穿插 (最终方案)

动机: DiT 应该只加载一次, 处理所有 clips 后才释放。VAE 足够小 (1.2GB), 可以和 DiT 共存在 GPU 上做快速穿插。

方案:

复制代码

_setup_dit() ← 一次
for each clip:
    DiT 去噪 (一直在 GPU, 不拆不建)
    VAE CPU→GPU (1.2GB, 与 DiT 共存) → 解码+编码motion → GPU→CPU
_teardown_dit() ← 一次

实现 : 将 _setup_dit() / _teardown_dit() 移到 clip 循环外。

结果:

11/40 blocks 常驻 GPU (12GB 预留, 兼容 VAE 共存)
每步耗时: 29s/step (稍慢于 v3/v4, 因为预留多了 VAE 的空间)
clip 间无间隔, 直接开始下一个 clip
模型加载次数: DiT 1次, VAE 按需穿插 (~0.5s/次)

3. flash_attn 回退方案

问题 : Windows 无预编译 flash_attn, 模型硬性调用 flash_attention() 函数。

原始代码 (wan/modules/attention.py):

python 复制代码

if FLASH_ATTN_3_AVAILABLE:
    x = flash_attn_interface.flash_attn_varlen_func(...)
else:
    assert FLASH_ATTN_2_AVAILABLE  # ← 直接崩溃
    x = flash_attn.flash_attn_varlen_func(...)

修复 : 添加 scaled_dot_product_attention 回退分支:

python 复制代码

elif FLASH_ATTN_2_AVAILABLE:
    x = flash_attn.flash_attn_varlen_func(...)
else:
    # SDP fallback - 无需 flash_attn
    q = q.unflatten(0, (b, lq)).transpose(1, 2)
    k = k.unflatten(0, (b, lk)).transpose(1, 2)
    v = v.unflatten(0, (b, lk)).transpose(1, 2)
    x = torch.nn.functional.scaled_dot_product_attention(
        q, k, v, attn_mask=None, is_causal=causal, 
        dropout_p=dropout_p, scale=softmax_scale)
    x = x.transpose(1, 2).contiguous()

影响: SDP attention 比 flash_attn 慢约 10-20%, 但在 Windows 上是唯一可用方案。

4. 最终性能对比

单步速度

方案	GPU blocks	每步耗时	相比基线
v1 cpu_offload (基线)	0/40	48s	1.0x
v2 device_map="auto"	N/A	>600s	0.08x (不可用)
v3 混合常驻	11/40	26s	1.85x
v4 顺序独占	13/40	26s	1.85x
v5 加载一次	11/40	29s	1.66x

总耗时 (68 秒音频, ~23 clips, 480×832, 20 steps)

方案	模型加载次数	clip间开销	总耗时	相比基线
v1	1次	0	8+ 小时 (未完成)	基线
v4	23×2次	~3min/clip	~4.8 小时	1.7x
v5	1次	~0.5s	~3.9 小时	2.0x

GPU 显存分布 (v5 最终方案)

阶段	GPU 占用	空闲
VAE 编码 (步骤3)	VAE 1.2GB	22.8GB
DiT 去噪 (步骤4)	DiT 11.1GB + 激活 ~11GB	~2GB
VAE 穿插 (步骤5)	DiT 11.1GB + VAE 1.2GB	~11.7GB

5. 核心技术总结

5.1 关键认知

device_map="auto" ≠ cpu_offload
- device_map: CPU 模块在 CPU 上计算 → 极慢
- cpu_offload: 权重在 CPU, 但搬到 GPU 上计算 → 合理
- AlignDevicesHook(offload=True): 最灵活, 可混合使用
激活值显存不可忽略
- 14B 模型在 480×832 分辨率下, attention 激活值需要 ~11GB
- 模型权重 + 激活值必须同时在 GPU 上
- 预留不足 → 显存抖动 → 速度暴跌甚至卡死
模型加载次数是关键开销
- 搬运 9GB+ 权重需要 ~0.4s (PCIe 4.0)
- 但 hook 设置/销毁、CUDA 缓存清理等额外开销可达数秒
- 23 clips × 每 clip 拆建一次 = 大量无效开销
顺序独占 vs 并行共存
- 理论上每阶段独占 GPU 最优 (最大利用)
- 实际上 VAE 仅 1.2GB, 与 DiT 共存代价很小
- DiT 加载一次 + VAE 快速穿插 = 最佳平衡

5.2 不可突破的瓶颈

motion 链式依赖: clip N 必须等 clip N-1 完成, 无法并行多 clip
PCIe 带宽: 29 个 offload blocks × 703MB × 2(来回) ≈ 40.8GB/次前向, 限制传输时间 ~1.6s
显存物理限制: 24GB 下最多常驻 11 个 blocks (28GB 模型的 ~28%)

5.3 进一步优化方向 (未实施)

方向	预期效果	代价
模型量化 (INT8/INT4)	模型体积减半, 更多 blocks 常驻	质量下降
flash_attn 安装	注意力计算快 10-20%	需要 Linux 或编译
多卡并行 (FSDP+Ulysses)	线性加速	需要多张 GPU
砍断 motion 依赖	可并行所有 clips	画面不连贯
潜空间 motion 传递	省去 clip 间 VAE 编解码	可能影响质量

6. 修改文件清单

文件	修改内容
`wan/speech2video.py`	新增 `use_device_map` 参数; `_setup_dit()` / `_teardown_dit()` / `_move_vae()` 方法; VAE 按需加载; DiT 混合 offload; generate() 顺序流水线
`wan/modules/attention.py`	flash_attention() 添加 SDP 回退分支
`generate.py`	新增 `--use_device_map` CLI 参数

7. 使用方式

bash 复制代码

# 设置 HF 镜像 (中国网络)
set HF_ENDPOINT=https://hf-mirror.com

# 运行生成
E:\Wan2.2\.venv\Scripts\python.exe generate.py ^
  --task s2v-14B ^
  --size 480*832 ^
  --ckpt_dir "E:\Wan2.2-S2V-14B\" ^
  --use_device_map ^
  --t5_cpu ^
  --prompt "A cat wearing sunglasses is talking on the beach" ^
  --image examples/i2v_input.JPG ^
  --audio "E:\千问剪辑skill\podcast-generator\output_qianfang\drama_qianfang.wav" ^
  --sample_steps 20 ^
  --infer_frames 48

报告生成时间: 2026-06-21

硬件: RTX 4090 D 24GB | 软件: PyTorch 2.11+cu128, diffusers 0.38.0, accelerate 1.14.0