31B 音视频扩散压到 14B：INT8 量化 + 6 步蒸馏

这篇能做出什么

先说明结果，让你判断是否值得继续往下读。

跟着这篇走完，你会得到一条可复用的压缩流水线，把一个 31B 量级的音视频联合扩散模型（视频 DiT 主干 + 音频 DiT 支路 + 跨模态注意力融合层）压到约 14B 参数、权重以 INT8 存储、采样步数从几十步降到 6 步，并且有一套固定的对比方法来判断“质量掉多少、显存和时延赚多少”。

需要先明确三个前提，避免后续预期错位：

1. 参数量不等于显存。31B 降到 14B 是参数口径；再叠加 INT8 权重存储，权重占用大约再减半。但总显存还包括激活、VAE、文本编码器、KV 缓存，长视频高分辨率时激活可能比权重还大。

2. 质量掉点主要来自剪枝和步数蒸馏，量化通常只是次要因素。实践里量化带来的损失常常可以压到很小，而删掉 40% 的层、把采样压到 6 步，才是真正需要反复调的部分。所以顺序上先剪枝、再量化、最后蒸馏，每一步都单独评测。

3. 时延收益不等于显存收益。INT8 权重能不能真的跑快，取决于你的推理后端有没有对应 kernel。没有 INT8 kernel 时，权重量化只省显存不省时间，这点要提前在目标硬件上验证。

下面所有代码都是通用骨架，函数名和超参需要按你的模型结构替换；涉及具体库版本、算子支持范围的地方，以官方文档当前版本为准。

前置条件清单

评测集要提前冻结，并且存下原始模型生成的样片。压缩过程中你会反复回看这批样片，没有它就只能凭感觉。

第 0 步：先把“31B”拆开，确定压哪里

不要一上来就动刀。先统计参数分布，找出真正的大头。

```python

inventory.py

import collections

import torch

def summarize(model, name="model"):

total = 0

buckets = collections.defaultdict(int)

for n, p in model.named_parameters():

c = p.numel()

total += c

buckets[n.split(".")[0]] += c

print(f"{name}: {total / 1e9:.2f}B")

for k, v in sorted(buckets.items(), key=lambda x: -x[1]):

print(f" {k:28s} {v / 1e9:7.2f}B {v / total:6.1%}")

return total

if __name__ == "__main__":

model = build_model() # 替换成你的加载逻辑

summarize(model, "full")

```

典型输出会长这样（结构名示意）：视频主干占大头，音频支路占比小但承担对齐职责，文本编码器和 VAE 加起来可能只有几个百分点。

结论通常是：压缩预算几乎全部要花在视频主干上。音频支路参数少，压它省不下多少显存，却最容易把音画同步搞崩，建议最后再考虑。

同时记录基线：

```python

import torch, time

@torch.no_grad()

def profile(fn, warmup=3, iters=10):

for _ in range(warmup):

fn()

torch.cuda.synchronize()

torch.cuda.reset_peak_memory_stats()

t0 = time.perf_counter()

for _ in range(iters):

fn()

torch.cuda.synchronize()

dt = (time.perf_counter() - t0) / iters

peak = torch.cuda.max_memory_allocated() / 1024 ** 3

return dt, peak

```

torch.cuda.max_memory_allocated() 只统计张量，不含 CUDA 上下文和显存碎片，报数字时说明清楚口径。

第 1 步：剪枝选层——用敏感度排序，别凭感觉删

目标是把 31B 砍到 14B 量级。这么大的比例，只能做结构化剪枝：整块删除 transformer block，加上 FFN 中间维收窄。非结构化稀疏在这个比例下既难恢复质量，也难拿到实际加速。

1.1 给每个 block 打敏感度分

思路很朴素：把某个 block 换成恒等映射（直接透传输入），看模型输出偏离多少。偏离越小，说明这层越可删。

```python

sensitivity.py

import torch

def make_skip_handles(block):

"""把 block 的输出替换成它的输入，实现'跳过该层'。"""

def pre_hook(module, args, kwargs):

module._skip_cache = args[0] if args else kwargs["hidden_states"]

def post_hook(module, args, kwargs, output):

return module._skip_cache

with_kwargs 的 hook 签名以 PyTorch 官方文档当前版本为准

h1 = block.register_forward_pre_hook(pre_hook, with_kwargs=True)

h2 = block.register_forward_hook(post_hook, with_kwargs=True)

return [h1, h2]

@torch.no_grad()

def block_sensitivity(model, batches, blocks, metric_fn):

base = [metric_fn(model, b) for b in batches]

scores = []

for idx, blk in enumerate(blocks):

handles = make_skip_handles(blk)

try:

cur = [metric_fn(model, b) for b in batches]

finally:

for h in handles:

h.remove()

delta = sum((c - b_).pow(2).mean().item() for c, b_ in zip(cur, base)) / len(base)

scores.append((idx, delta))

return sorted(scores, key=lambda x: x[1])

```

metric_fn 建议直接用最终输出（解码后的视频帧 + 音频波形的组合损失），而不是中间层特征。中间层特征相似的层，对最终画质的影响可能完全不同。

1.2 加两道保护

纯敏感度排序会误删关键层，务必手工保护：

• 首块与末块：几乎在所有扩散模型里都高度敏感，直接加入白名单。
• 含跨模态注意力的块：把每个 block 内 cross-attention 的梯度范数或注意力熵算出来，融合层通常熵更低、更“专一”，删掉容易出音画不同步。

```python

def cross_modal_score(block):

"""用 cross-attn 输出的梯度范数做代理指标。"""

for name, p in block.named_parameters():

if "cross" in name and p.grad is not None:

return p.grad.norm().item()

return 0.0

```

1.3 定删除比例并做宽度收窄

第一轮建议只删 block，先不要动 FFN 宽度。删到参数量降 30% 左右时停手，重测一遍敏感度（删完后排序会变，要迭代），再进入第二阶段：对剩下的 block 做 FFN 中间维裁剪，按神经元权重 L2 排序，成组裁剪到 8 的倍数，保证 kernel 友好。

31B → 14B 通常是 2～3 轮“删层 + 重评 + 再删层”迭代出来的，一次到位很难。

第 2 步：剪枝后的短恢复训练

剪完不要直接量化，先做一次轻量恢复。用原始模型做 teacher，让学生（剪枝后的模型）回归 teacher 的 velocity 预测：

```python

def restore_step(student, teacher, batch, optimizer, t_scale=0.5):

x, cond = batch["latent"], batch["cond"]

t = torch.rand(x.shape[0], device=x.device) * t_scale # 只用低噪声段，收敛更快

x_t = add_noise(x, t)

with torch.no_grad():

v_t = teacher(x_t, t, cond)

v_s = student(x_t, t, cond)

loss = torch.nn.functional.mse_loss(v_s, v_t)

optimizer.zero_grad(set_to_none=True)

loss.backward()

torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)

optimizer.step()

return loss.item()

```

学习率取原始预训练时的一个小比例，步数不用多，重点是让输出分布回到 teacher 附近。这一步做扎实，后面量化和蒸馏都会好过很多。

第 3 步：INT8 量化与校准

3.1 先做层融合，再量化

把 LayerNorm + Linear、QKV 的三路投影、卷积 + 归一化这些能合的合掉。融合会改变数值分布，必须在量化前完成，否则校准出来的 scale 全部作废。

3.2 量化策略

• 权重：per-channel（按输出通道）对称 INT8，这是最稳的选择。
• 激活：per-token 动态量化，即每次推理实时算 scale，避免校准集覆盖不到的长尾。
• 不量化：所有 LayerNorm/RMSNorm、softmax、时间步嵌入、调制参数（scale/shift 那一路）。这些张量小但敏感。
• 谨慎量化：音频支路的第一层和最后一层、跨模态注意力的 value 投影。

3.3 用假量化先评估，别急着上真 kernel

```python

fake_quant.py

import torch

def quant_weight_per_channel(w, n_bits=8, dim=0):

qmax = 2 ** (n_bits - 1) - 1

qmin = -qmax - 1

w32 = w.float()

reduce_dims = [d for d in range(w32.dim()) if d != dim]

amax = w32.abs().amax(dim=reduce_dims, keepdim=True).clamp(min=1e-8)

scale = amax / qmax

q = torch.round(w32 / scale).clamp(qmin, qmax)

return (q * scale).to(w.dtype)

def quant_act_per_token(x, n_bits=8, dim=-1):

qmax = 2 ** (n_bits - 1) - 1

amax = x.abs().amax(dim=dim, keepdim=True).clamp(min=1e-8)

scale = amax / qmax

return (torch.round(x / scale).clamp(-qmax - 1, qmax) * scale)

```

把这两个函数挂到 Linear 上做“模拟量化”，跑一遍评测集，看质量掉点。这一步掉了多少，就是真量化大致会掉的量，如果假量化阶段就崩了，说明激活有严重离群值，需要先处理。

3.4 处理激活离群：SmoothQuant 思路

扩散模型的激活里常见个别通道数值远高于其他通道。做法是把量化难度从激活“搬”一部分到权重上：

```python

def smooth_linear(linear, act_amax, alpha=0.5):

w_amax = linear.weight.abs().amax(dim=0).clamp(min=1e-8)

s = (act_amax.clamp(min=1e-8).pow(alpha) / w_amax.pow(1 - alpha)).clamp(1e-3, 1e3)

linear.weight.data = linear.weight.data * s.unsqueeze(0)

if linear.bias is not None:

linear.bias.data = linear.bias.data * s

return s # 推理时在前一层输出上除以 s

```

alpha 从 0.5 开始试，音频支路一般要比视频支路小一些。

3.5 校准集怎么选

校准集的质量直接决定量化效果。硬要求：

• 音视频成对，不能只有画面没声音。
• 覆盖你的真实分布：单人说话、多人对话、音乐、环境音、快速运动、静态特写，各来一些。
• 时长和分辨率要覆盖上线时的上限，不要全部用最短的样本。
• 200～500 条通常够用，但分布比数量重要。

校准完，把每一层的 scale 存下来，导出的权重文件要可复现。

第 4 步：6 步蒸馏

4.1 思路

原始模型走 N 步（N 通常几十）。你要一个 6 步就能出片的 student。在 flow matching / rectified flow 框架下，可以用轨迹重排（reflow）：让 teacher 走完整条轨迹，记录中间状态，构造新的“直线化”配对数据，让 student 学一条更直的路径。

```python

@torch.no_grad()

def build_trajectory(teacher, z0, cond, n_steps):

ts = torch.linspace(0.0, 1.0, n_steps + 1, device=z0.device)

x = z0

traj = [x]

for i in range(n_steps):

t = ts[i].expand(x.shape[0])

v = teacher(x, t, cond)

x = x + v * (ts[i + 1] - ts[i])

traj.append(x)

return ts, traj

```

一次轨迹生成会占不少显存，建议离线跑一遍，把 (x_i, x_j, t_i, t_j, cond) 落盘，训练时只读不生成。

4.2 训练目标

```python

def distill_step(student, ts, traj, cond, optimizer):

b = traj[0].shape[0]

i = torch.randint(0, len(traj) - 1, (b,), device=traj[0].device)

j = torch.randint_like(i, len(traj))

i, j = torch.minimum(i, j), torch.maximum(i, j)

keep = (j > i)

if keep.sum() == 0:

return None

xi = torch.stack([traj[k] for k in range(len(traj))])[i, torch.arange(b)]

xj = torch.stack([traj[k] for k in range(len(traj))])[j, torch.arange(b)]

ti, tj = ts[i], ts[j]

v_target = (xj - xi) / (tj - ti).clamp(min=1e-6).view(-1, *([1] * (xi.dim() - 1)))

v_pred = student(xi, ti, cond)

loss = torch.nn.functional.mse_loss(v_pred, v_target)

optimizer.zero_grad(set_to_none=True)

loss.backward()

torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)

optimizer.step()

return loss.item()

```

4.3 音视频要加同步约束

联合模型只回归 velocity 会出现音画“各走各的”。加一项跨模态一致性损失——比如用音频特征去预测视频帧级别的运动能量，或反之，约束两者在时间轴上对齐：

```python

loss = loss_v + loss_a + lam * sync_loss(video_feat, audio_feat)

```

lam 从很小的值开始加，加太大会牺牲画质换同步。

4.4 把 CFG 折进 student

如果推理时还要跑两遍（条件 + 无条件）做 classifier-free guidance，6 步也省不下多少。蒸馏时把 guidance scale 作为条件输入 student，让它在单次前向里就产出“已引导”的结果，端到端时延能再接近减半。

4.5 步数不是越少越好

6 步是个常见折中点。压到 4 步以下时，通常是高频细节先崩：毛发、文字、快速运动拖影，音频则是齿音和瞬态。建议把 6 步、5 步、4 步各训一版，用同一批 seed 出片，肉眼对比后再定。

第 5 步：出片对比与收益核算

固定 prompt、固定 seed、固定分辨率、固定 guidance。逐项记录：

时延测量务必用 CUDA Event 而不是 time.time()，并且要先 warmup：

```python

start, end = torch.cuda.Event(True), torch.cuda.Event(True)

start.record()

out = student.sample(cond, steps=6)

end.record()

torch.cuda.synchronize()

print(f"{start.elapsed_time(end):.1f} ms")

```

权衡判断的实操标准：先设一条不可退让的底线（比如“人声口型不能飘”“文字不能糊”），在守住底线的方案里选显存和时延最优的那个。不要用一个总分把所有维度揉在一起，那样等于没有决策依据。

常见坑与排错

删层后音画不同步、口型对不上

大概率删到了跨模态融合 block。回退上一轮，把 cross-attn 梯度范数排名靠前的几层加白名单，重新剪。

量化后画面还行，声音先崩

音频支路对激活离群更敏感。优先把音频支路的第一/最后一层、value 投影改成 FP16，或者只对音频做 weight-only INT8、激活保持 FP16。

校准完指标断崖

检查校准集是不是只有单一场景。补上音乐、多人、环境噪声明亮暗差异大的样本，重跑校准。

显存没降下来

权重降了但激活没降，或者优化器状态、teacher 模型、VAE 还占着显存。分离部署评测：只加载 student 和 VAE 再测一次。

时延没降

最可能是算子回退。用后端的 profiler 看有多少层实际走的是 FP16 回退路径。没有 INT8 kernel 的层，量化只省显存不省时间。

假量化精度正常，导出后精度崩了

检查层融合是不是在量化之后才做、LayerNorm 有没有被误量化、per-channel 的维度是不是选错了（应该是输出通道，不是输入通道）。

蒸馏训练 loss 一直降但出片变差

target 构造有偏。检查 tj - ti 是否出现过接近 0 的情况导致数值爆炸，以及轨迹是不是用带 CFG 的 teacher 生成的——两者混用会让 student 学到一个不存在的分布。

对比结果不可信

种子没锁、guidance 没锁、分辨率没锁，或者两次采样之间重启了进程导致随机状态不同。把评测脚本写死成一份，全流程只改模型路径。

下一步建议

1. 先量化后剪枝会更容易失败。如果第一轮剪枝效果不理想，可以试试先做 INT8（对质量影响小），在量化模型上做敏感度分析，剪枝信号会更干净。

2. 权重 4bit + 激活 FP16 是压缩权重显存的下一个台阶，实现难度比全 INT8 低，收益在权重占主导时明显。

3. 分块推理/CPU 卸载 是显存不够时的兜底方案，代价是时延，但能让你在消费级卡上跑起来做验证。

4. 端侧部署要额外考虑算子支持度和内存带宽，同一个 INT8 模型在服务器卡和移动端的收益可能完全不同，务必在目标硬件上实测。

5. 把整条流水线脚本化：清点 → 敏感度 → 剪枝 → 恢复 → 校准 → 量化 → 蒸馏 → 评测，每一步产出带版本号的中间产物。压缩是个反复试错的活，能复现的流水线比单次调参值钱得多。