【Bug已解决】allegro model/pipeline review 解决方案

【Bug已解决】allegro model/pipeline review 解决方案

一、现象长什么样

对 diffusers 的 Allegro 视频生成模型（allegro model/pipeline review，即 Rhymes 的文生视频 transformer）做审查时，发现一个时间注意力掩码 bug：Allegro 的视频 transformer 在时空联合注意力里，对时间维（帧间）使用因果 mask（第 t 帧只能看 ≤ t 的帧），但共享注意力层在构造时间因果 mask 时，把帧索引的轴向搞错了——按“空间 token 位置”而不是“帧序号”做 causal，导致同一帧内的 token 互相看不到未来帧、却错误地让不同帧的同位置token 单向可见，破坏了视频的时间一致性。现象：

# 现象 A：生成视频帧间闪烁/跳变，物体在第 3 帧突然“瞬移” # 时间因果被破坏，帧间依赖关系错乱 # 现象 B：和官方实现对拍，运动轨迹不一致 # 权重、结构都对，唯独帧序列运动不连贯 —— 定位到 temporal mask # 现象 C：不报错，但长视频（帧数多）更明显 # 因为帧越多，轴向错误累积的可见性错位越严重

最隐蔽的是现象 B：能跑、不报错、单帧看着还行，但帧序列的运动逻辑是错的，只能靠和官方逐帧对拍发现。

二、背景

Allegro 把视频当成“帧 × 空间 patch”的 3D token 序列送入 transformer。注意力分两种：① 空间注意力（每帧内 patch 互相看）；② 时间注意力（跨帧同位置 patch 看）。时间注意力必须按帧序号做因果：第 t 帧的时间 query 只能 attend 第 0..t 帧的对应 patch。

审查发现：共享时间注意力层在算 causal mask 时，输入的 token 布局是[frames, patches_per_frame, ...]展平后的 1D 序列，但 mask 构造代码按“展平后的绝对位置”直接做上三角 causal，没先把绝对位置映射回(frame_idx, patch_idx)再只对frame_idx维度 causal。于是它实际上是对“展平位置”做了 causal——这等价于让第 1 帧的第 100 个 patch 看不到第 0 帧的第 50 个 patch，但能看到第 0 帧的第 99 个 patch，完全不是“按帧 causal”的语义。

这是视频 transformer 审查里极典型的坑：3D 布局展平后，轴向语义丢失，mask 按错维度施加。

三、根因

1. 时间 causal 按展平位置而非帧序号：mask 构造对 1D 展平序列做上三角，丢失了frame_idx维度，导致可见性不是“按帧”而是“按绝对位置”。

2. 帧/ patch 布局假设不一致：代码假设[patches, frames]布局，实际是[frames, patches]，轴向假设错导致 mask 整体错位。

3. 缺少与参考实现逐帧对拍：没有断言“相同输入下 temporal-masked attention 输出与官方一致”，轴向错误长期存在。

本质：是视频 transformer 时间因果 mask 在 3D 布局展平后丢失了帧维度语义，按错轴施加，且缺少参考对拍。

四、最小可运行复现

下面复现“时间 causal 按展平位置而非帧序号，导致可见性错位”：

import torch def temporal_mask_buggy(num_frames, patches_per_frame): """buggy: 对展平后的绝对位置做上三角 causal。""" total = num_frames * patches_per_frame # 上三角：pos j > pos i 不可见 —— 这是“按绝对位置”causal，错！ mask = torch.triu(torch.ones(total, total), diagonal=1) * float("-inf") return mask def temporal_mask_fixed(num_frames, patches_per_frame): """fixed: 只对 frame_idx 维度 causal，patch 维度内全可见。""" total = num_frames * patches_per_frame mask = torch.zeros(total, total) for q in range(total): q_frame = q // patches_per_frame for k in range(total): k_frame = k // patches_per_frame if k_frame > q_frame: # 只看过去和当前帧 mask[q, k] = float("-inf") return mask mb = temporal_mask_buggy(2, 2) # 帧0:[0,1] 帧1:[2,3] mf = temporal_mask_fixed(2, 2) print("buggy: frame0-patch0 看 frame1-patch0 (pos2)?", mb[0, 2].item() == 0.0) # True → 错误地可见（跨帧未来） print("fixed: frame0-patch0 看 frame1-patch0 (pos2)?", mf[0, 2].item() == float("-inf")) # True → 正确不可见

buggy里位置 0（帧0）能看位置 2（帧1），违反时间因果；fixed正确禁止。

五、解决方案（第一层：最小直接修复）

最小修复：时间 causal mask 必须先把展平位置映射回frame_idx，只对帧维度做因果，patch 维度内部保持全可见：

import torch def build_temporal_mask(num_frames, patches_per_frame): total = num_frames * patches_per_frame mask = torch.zeros(total, total) for q in range(total): q_frame = q // patches_per_frame # 还原帧维度 for k in range(total): k_frame = k // patches_per_frame if k_frame > q_frame: # 仅按帧因果 mask[q, k] = float("-inf") return mask

这一层改动最小：用// patches_per_frame还原帧索引再比较，时间因果恢复正确。但它依赖“每个时间注意力层都写对”，下看第二层。

六、解决方案（第二层：结构性改进）

把“Allegro 时间因果 mask 的构造规则”固化成单一事实来源。下面这个 dataclass 集中管理：从 token 布局推导帧维度、构造时间因果、并与参考对拍。

from dataclasses import dataclass, field from typing import Callable import torch @dataclass class AllegroTemporalMaskPolicy: """单一事实来源：Allegro 视频 transformer 时间因果 mask 规则。""" patches_per_frame: int def build(self, num_frames: int) -> torch.Tensor: total = num_frames * self.patches_per_frame mask = torch.zeros(total, total) for q in range(total): qf = q // self.patches_per_frame for k in range(total): kf = k // self.patches_per_frame if kf > qf: mask[q, k] = float("-inf") return mask def verify_against_reference(self, ref_fn: Callable[[int, int], torch.Tensor], num_frames: int) -> None: mine = self.build(num_frames) ref = ref_fn(num_frames, self.patches_per_frame) if not torch.equal(mine, ref): raise AssertionError("temporal mask differs from reference") # 用法 policy = AllegroTemporalMaskPolicy(patches_per_frame=256) mask = policy.build(num_frames=16)

这一层的关键收益：

• 布局即参数：patches_per_frame是显式参数，杜绝“假设布局”导致的轴向错；
• 参考对拍：verify_against_reference直接比对官方 mask，轴向错误立刻暴露；
• 单一事实来源：所有 Allegro 时间因果约定收口在AllegroTemporalMaskPolicy，审查只盯它。

七、解决方案（第三层：断言 / CI 守护）

把第二层钉成 pytest，挂进 CI，确保时间因果按帧、跨帧未来不可见、变长一致：

import torch import pytest from your_package.allegro_mask import AllegroTemporalMaskPolicy def _ref(nf, ppf): total = nf * ppf m = torch.zeros(total, total) for q in range(total): for k in range(total): if (k // ppf) > (q // ppf): m[q, k] = float("-inf") return m def test_future_frame_invisible(): # 断言 1：未来帧对当前帧不可见 policy = AllegroTemporalMaskPolicy(patches_per_frame=2) m = policy.build(2) assert m[0, 2].item() == float("-inf") # 帧0 看不了帧1 assert m[2, 0].item() == 0.0 # 帧1 能看帧0 def test_within_frame_visible(): # 断言 2：同帧内 patch 互相可见（不按绝对位置 causal） policy = AllegroTemporalMaskPolicy(patches_per_frame=2) m = policy.build(2) assert m[0, 1].item() == 0.0 # 帧0 内 patch0 看 patch1 def test_variable_frames(): # 断言 3：不同帧数都正确 policy = AllegroTemporalMaskPolicy(patches_per_frame=4) for nf in (1, 4, 8): m = policy.build(nf) assert m[0, nf*4-1].item() == float("-inf") # 首帧看不了末帧 def test_reference_match(): # 断言 4：与参考对拍 policy = AllegroTemporalMaskPolicy(patches_per_frame=4) policy.verify_against_reference(_ref, 6) # 不抛异常

四条断言从“未来帧不可见”“同帧可见”“变长正确”“参考对拍”四面把轴向错误钉死在 CI。

八、排查清单

审查allegro或任何视频 transformer 时：

1. 用官方权重跑视频，和官方 repo 逐帧对拍。运动不连贯但单帧对，就怀疑 temporal mask。
2. 时间因果是按“帧序号”还是“展平绝对位置”？按绝对位置就是轴向错（现象 A）。
3. patches_per_frame布局假设是否和实际一致？不一致 mask 整体错位。
4. 用第二层AllegroTemporalMaskPolicy：布局显式参数化 + 参考对拍。
5. 加第三层 pytest，断言“未来帧不可见、同帧可见、变长正确、参考对拍”。
6. 视频模型 mask 错也“能跑”，必须靠对拍和断言才能发现。

九、小结

allegro审查发现的核心 bug 是视频 transformer 的时间因果 mask 在 3D token 布局展平后丢失了帧维度语义，按“绝对位置”而非“帧序号”施加 causal，导致帧间依赖关系错乱、生成视频闪烁跳变；且因能跑不报错，只能靠与官方对拍发现。修复分三层——第一层用// patches_per_frame还原帧索引只对帧维度 causal；第二层用AllegroTemporalMaskPolicy这个 dataclass 把时间因果规则收口成单一事实来源，布局显式参数化并内置参考对拍；第三层用四条 pytest 把“未来帧不可见、同帧可见、变长正确、参考对拍”钉死在 CI。核心心法：视频 transformer 的时间因果必须按帧序号施加，3D 布局展平后务必先还原帧维度，否则轴向错误只会静默毁掉时间一致性。