OSP-Next：基于稀疏序列并行、HiF8量化和强化学习的高效高质量视频生成 OSP-Next: Efficient High-Quality Video Generation with Sparse Sequence Parallelism, HiF8 Quantization, and Reinforcement Learning
通过稀疏注意力、序列并行、量化和强化学习，实现高质量视频生成的1.64倍加速
前置知识
Diffusion Transformer (DiT)
DiT将UNet架构中的卷积层替换为Transformer，通过注意力机制处理图像/视频的时空token序列。在视频生成中，token数量随帧数和分辨率增长，注意力复杂度为O(N^2)，其中N是序列长度。
本文的核心是优化DiT的计算效率，理解DiT的二次复杂度瓶颈是理解本文动机的关键。
稀疏注意力
稀疏注意力通过限制token之间的交互来降低计算复杂度。常见策略包括基于动态token选择（根据相似度打分选择部分token）、基于固定规则（如局部窗口、步进模式）等。
Skipparse-2D Attention是本文的核心创新，理解稀疏注意力的原理和挑战有助于理解该创新的价值。
序列并行
序列并行将长序列切分到多个设备上并行处理。Ulysses-style SP在每个设备上计算部分token的Q、K、V，通过All-to-All通信聚合完整注意力结果。通信量是序列并行的主要瓶颈。
SSP是本文的重要贡献之一，理解传统SP的通信瓶颈有助于理解SSP减少75%通信量的意义。
FlashAttention
FlashAttention是一种针对GPU优化的注意力计算内核，通过分块计算、重计算和融合kernel减少内存访问和显存占用。它依赖连续的内存布局和标准的注意力模式，不直接支持复杂的稀疏掩码。
Skipparse-2D Attention的原生FlashAttention兼容性是一个关键优势，这解释了为什么本文强调避免动态掩码。
GRPO强化学习
GRPO是一种策略梯度方法，通过在多个采样轨迹间比较优势函数来更新策略。在扩散模型中，将去噪轨迹视为马尔可夫决策过程，使用奖励信号优化生成策略。
本文使用Mix-GRPO恢复稀疏微调带来的质量损失，理解RL范式有助于理解这一后训练策略的目的。
研究动机
Diffusion Transformers在视频生成中展现了强大的生成能力，但全注意力的二次复杂度严重限制了推理速度和可实现的分辨率。例如，HunyuanVideo和Wan2.1在单张NVIDIA A100上生成5秒720p视频需要约30-60分钟。现有优化方法各有局限：训练无关的稀疏注意力依赖预训练全注意力模型，固定权重无法适应注意力模式变化引起的分布偏移；基于训练的稀疏方法需要动态构建注意力掩码，通常用FlexAttention或自定义稀疏内核实现，难以与高效的FlashAttention内核原生集成；不规则掩码还导致并行时负载不均衡。线性注意力虽然降低复杂度到O(N)，但简化限制了表达能力。量化方法如FP8和INT8的表示范围远窄于BF16和FP16，8位精度训练的模型性能往往受限。这些问题共同制约了视频DiT的实用化。
本文的目标是本文的目标是构建一个集成稀疏注意力、并行化、量化和强化学习的高效文本到视频生成模型，在保持生成质量的同时显著提升推理速度。具体而言，目标包括：设计一种可训练的稀疏注意力机制，既减少计算量又保持与FlashAttention的原生兼容；开发针对稀疏注意力的序列并行策略，降低通信开销；实现稳定的8位精度联合训练；通过强化学习后训练恢复稀疏微调带来的质量损失。最终目标是跨越不同硬件平台（NVIDIA GPU和华为昇腾NPU）提供一致的加速效果。
与已有工作不同的是，本文的独特切入点在于将稀疏注意力设计、序列并行策略、低精度训练和强化学习优化整合到一个统一的框架中，而非孤立地解决单个问题。现有稀疏视频生成方法要么是训练无关的推理加速（无法在预训练阶段获益），要么基于动态token选择（难以并行），要么需要复杂掩码（不兼容FlashAttention）。本文提出Skipparse-2D Attention，其固定规则模式天然支持FlashAttention，并基于其局部等价性提出SSP，实现高效并行。此外，本文首次探索了稀疏模型与8位量化的联合训练可行性，以及强化学习对稀疏模型的优化效果。这种系统性的优化思路是本文的核心创新。
核心方法
OSP-Next的整体思路是利用视频模态的时空局部性，通过固定规则的稀疏注意力减少计算量，同时设计配套的并行策略使稀疏模式在分布式环境下高效运行，再用量化进一步加速，最后用强化学习恢复质量损失。技术路线包括：Skipparse-2D Attention将视频token在高度和宽度维度分别应用skipparse模式，形成token-wise和group-wise稀疏注意力的交替；Sparse Sequence Parallelism (SSP)利用skipparse rearrange的局部等价性，通过单次All-to-All通信实现稀疏模式切换；HiF8量化采用华为HiF8格式，支持动态的指数-尾数位分配，实现稳定的8位训练；Mix-GRPO后训练结合SDE和ODE采样，在部分去噪步骤应用RL优化。模型采用混合注意力架构，前后若干层使用全注意力以捕获细粒度特征，中间层使用Skipparse Attention以最大化效率。
核心创新点是Skipparse-2D Attention的固定规则模式和由此衍生的SSP并行策略。与基于动态token选择的稀疏注意力不同，Skipparse-2D采用预定义的模式：token-wise稀疏注意力类似pixel unshuffle，将H乘W的图分成k乘k的子图，每个子图的token作为独立序列；group-wise稀疏注意力类似patch unshuffle，按k乘k的块将token分组。两者交替执行，确保任意两个token最多通过两次注意力操作即可交互。这种固定模式的优势在于：避免了动态构建2D掩码，原生支持FlashAttention；skipparse rearrange具有局部等价性，即对整个序列应用rearrange等价于对每个最小重复单元独立应用rearrange。基于此，SSP将子序列分配到不同rank，模式切换时只需本地rearrange加一次All-to-All通信，相比Ulysses SP的4次All-to-All，通信量减少75%。
方法步骤详情
方法步骤如下：Skipparse-2D Attention计算：给定形状为(B, T, H, W, C)的视频token，其中T是帧数，H和W是空间维度。Token-wise稀疏注意力执行rearrange操作，生成P的平方个子序列，每个子序列长度为原始的1除以P的平方。Group-wise稀疏注意力执行另一种rearrange操作。两种模式在层间交替，最终rearrange回原始布局。当H或W不是P的平方的倍数时，采用any resolution策略，填充到最近倍数并生成注意力掩码，确保不同分辨率视频的相同位置token遵循相同交互模式。SSP并行：在P的平方个rank间分配子序列。模式切换时，每个rank先执行本地rearrange，然后通过All-to-All通信交换数据。具体而言，reshape数据为(group_size, G, G乘b, base, C)，其中group_size是通信组大小，G等于P的平方除以group_size。All-to-All后permute并reshape回目标布局。整个过程的通信量是O(NS)，其中N是通信组大小，S是每rank数据量，远小于All-Gather-Rearrange-Reshard的O(N的平方乘S)。HiF8量化：缩放因子等于HiF8max除以(Amax加epsilon)，其中Amax是当前tensor的最大绝对值，HiF8max在前向传播设为15，反向传播设为224。采用per-tensor量化，无需per-block粒度。Mix-GRPO后训练：采样25步，前10步用SDE，后15步用ODE。使用VideoAlign作为奖励模型，提供视觉质量、运动质量、文本对齐和总分四个维度。LoRA rank为32，alpha为64，rollout group size为4，学习率为0.0001，KL正则化系数beta为0.004。
技术新颖性
技术新颖性体现在多个维度：Skipparse-2D Attention将1D序列的skipparse模式正交应用到2D空间维度，更好地匹配视频的空间局部性，避免了flatten序列导致的相邻token交互受限问题。这相比Skipparse-1D更符合视觉模态特性，实验也显示2D版本更接近3D全注意力的交互模式。SSP是首个专门为固定模式稀疏注意力设计的序列并行策略，利用局部等价性性质，将Ulysses SP的4次All-to-All减少到1次，通信量降低75%。HiF8与稀疏模型的联合训练是首次尝试，HiF8的锥形精度设计（中心指数范围保持3位尾数，边缘逐渐降到1位）比传统FP8（E4M3或E5M2）更适合神经网络数据的类高斯分布，使得per-tensor量化即可保持训练稳定性，损失曲线与BF16几乎重叠。强化学习应用于稀疏模型的后训练也是首次探索，Mix-GRPO的混合SDE-ODE策略降低了视频RL的昂贵成本，实验显示RL显著改善了空间关系、文本对齐和美学质量等维度。
实验结果
核心发现包括：生成质量：OSP-Next在VBench上获得83.73%的总分，超过Wan2.1基线的83.69%，表明稀疏注意力模型可以达到甚至超越全注意力模型的质量。具体指标上，OSP-Next在动态程度（DD, 70.83% vs 65.46%）、多物体（MO, 90.10% vs 86.28%）、人类动作（HA, 77.48% vs 69.58%）等语义相关维度有显著提升，但在背景一致性（BC, 96.95% vs 98.09%）和时间闪烁（TF, 98.49% vs 99.46%）等一致性维度略有下降。OSP-Next-HiF8的总分为83.29%，仅比基线低0.4%，证明HiF8量化几乎不损失质量。推理速度：在单张NVIDIA H200上，OSP-Next在720P（有padding）和768P（无padding）设置下分别达到1.53倍和1.64倍加速；在8张H200上分别达到1.42倍和1.52倍加速。在华为昇腾910C上，单卡加速1.73倍到1.75倍；在950PR上，单卡OSP-Next加速1.27倍到1.76倍，而OSP-Next-HiF8加速1.69倍到2.27倍，最高达到2.27倍。无padding设置下加速更大，因为避免了注意力掩码的开销。消融研究：初始化gap分析显示，当所有层都使用Skiparse Attention时，零样本推理无法生成有效视频。随着全注意力块数量从0增加到16，生成质量逐步接近基线。Skipparse-2D在所有设置下都比Skipparse-1D更接近基线，最终选择8个全注意力块（前后各4个）平衡性能和效率。强化学习消融显示，Mix-GRPO后训练显著改善了空间关系、文本对齐和美学质量，证明RL对恢复稀疏微调损失有效。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 文本到视频生成 | VBench总分 | 83.73% | Wan2.1-T2V-14B: 83.69% | 0.04个百分点 |
| 文本到视频生成（8位量化） | VBench总分 | 83.29% | Wan2.1-T2V-14B: 83.69% | -0.4个百分点（可接受） |
| 单GPU推理速度（NVIDIA H200, 720P） | 加速比 | 1.53倍 (620s vs 946s) | Wan2.1基线 | 加速53% |
| 单GPU推理速度（NVIDIA H200, 768P） | 加速比 | 1.64倍 (645s vs 1061s) | Wan2.1基线 | 加速64% |
| 8GPU推理速度（NVIDIA H200, 720P） | 加速比 | 1.42倍 (91s vs 129s) | Wan2.1基线（Ulysses SP） | 加速42% |
| 单NPU推理速度（华为昇腾950PR, 768P, HiF8） | 加速比 | 2.27倍 (590s vs 1338s) | Wan2.1基线 | 加速127% |
局限与改进
作者承认的局限性包括：Skipparse-2D Attention的固定模式无法适应不同视频的个性化注意力需求，某些场景可能需要更灵活的稀疏策略。需要从全注意力预训练模型初始化，无法从头训练稀疏模型，这限制了训练阶段的加速效果。SSP目前仅支持固定稀疏比率（本文P等于2），扩展到可变稀疏比率需要额外设计。HiF8是华为的专有格式，在其他硬件平台上的支持可能受限。观察到的局限性包括：混合架构中全注意力块的数量（8个）是经验选择，未进行系统搜索，可能存在更优配置。RL后训练仅使用VideoAlign奖励，未探索多奖励组合或人类反馈。实验仅在固定分辨率（720乘1280）和时长（81帧）下进行，未验证对更长序列和更高分辨率的可扩展性。推理速度评估基于torch.compile和FlashAttention3，这些优化在不同PyTorch版本或硬件环境下的效果可能不同。
独立分析的弱点
独立分析的弱点包括：固定稀疏模式的泛化能力有限：Skipparse-2D的固定规则可能不适应所有视频内容，例如某些场景需要长程依赖而skip模式限制了交互范围。改进方向可以是学习稀疏模式或根据内容自适应调整模式。初始化依赖全注意力模型：从全注意力预训练模型初始化意味着预训练阶段无法利用稀疏注意力的加速，训练成本仍然高昂。改进方向可以是设计可从头训练的稀疏架构或渐进式稀疏训练。RL优化的成本：虽然Mix-GRPO降低了视频RL成本，但仍需要昂贵的奖励评估和轨迹采样，限制了在更大模型上的应用。改进方向可以是使用更高效的奖励代理或蒸馏RL模型。稀疏比率的固定性：当前P等于2是硬编码的，无法根据硬件或性能需求动态调整。改进方向可以是支持可变稀疏比率或分层稀疏策略。评估覆盖不全：实验未涵盖更高分辨率（如1080P）、更长时长（如10秒）或更复杂场景（如多物体交互）的生成质量，这些是实际应用的关键场景。改进方向是扩展评估范围。
未来方向
作者提出的未来方向包括：探索可学习的稀疏注意力模式，在保持FlashAttention兼容性的同时提高灵活性。研究从头训练稀疏视频DiT的可能性，释放训练阶段的加速潜力。扩展SSP到更复杂的并行场景，如与tensor parallelism或pipeline parallelism的集成。探索其他低精度格式在稀疏模型上的应用，如NF4或自定义格式。基于本文成果可延伸的方向包括：将OSP-Next的优化思路应用到图像生成、音频生成等序列生成任务。研究稀疏注意力与其他效率技术（如知识蒸馏、模型剪枝）的协同效果。开发自动化工具，根据硬件配置和性能需求自动选择最优稀疏比率、全注意力块数量等超参数。探索更高效的RL范式，如离线RL或基于价值函数的方法，进一步降低后训练成本。在更大规模数据集和模型上验证OSP-Next的可扩展性，为视频生成的实用化奠定基础。
复现评估
复现评估：论文提供了GitHub仓库（https://github.com/PKU-YuanGroup/OSP-Next）和HuggingFace模型（https://huggingface.co/yunyangge/OSP-Next），代码和权重开源。实验使用内部收集的高质量视频数据集，采用与Open-Sora Plan相同的过滤策略，具体数据集规模未公开。训练配置详细：8卡H200或NPU，AdamW优化器，学习率为0.00002，权重衰减0.001，全局batch size 24，FSDP训练策略，EMA衰减0.999，SFT权重在6000步后获得。RL配置：LoRA rank为32，alpha为64，rollout group size为4，学习率为0.0001，KL正则化beta为0.004。硬件需求较高：8张高端GPU（H200）或华为昇腾910C或950PR，训练和推理都需要大量显存和时间。复现难度中等偏高，主要原因包括：内部数据集不可公开获取、硬件要求高、HiF8格式在非华为硬件上可能需要模拟实现（论文提到使用HiF8模拟算子）。然而，详细的配置说明和开源权重有助于复现主要实验结果。
论文图表
