论文
MegaSlide-DiT:以内存为中心的适应和可变形局部注意力,实现高效视频扩散
MegaSlide-DiT: Memory-Centric Adaptation and Deformable Local Attention for Efficient Video Diffusion
摘要
基于 Diffusion Transformer (DiTs) 构建的高分辨率视频扩散模型可提供强大的保真度,但很快就会耗尽单个工作站的内存预算。超过 1000 亿的参数 DiT 很容易需要超过 TB 的持久状态,而朴素的时空自注意力在序列长度上呈二次方增长。这两堵墙——参数内存和激活内存——阻止研究人员在没有大型 GPU 集群的情况下适应大规模生成模型。我们从系统角度重新审视这个问题,并介绍 MegaSlide-DiT,这是一个原型,演示了如何在具有 1.5 TB 主机 RAM 的单个 H200 GPU 上调整预训练的 105B DiT。我们的主要见解是 GPU 不需要拥有模型状态:所有持久权重、主权重和优化器时刻都保留在主机内存中,而只有瞬态分片按需流式传输到 GPU。同时,我们用 3D Deformable Slide Attention (3D-DSA) 取代二次全局注意力,3D-DSA 是一种运动自适应局部注意力算子,可将内存和计算复杂度降低到序列长度的线性。我们报告详细的内存统计、执行跟踪和评估结果来证实我们的设计。 MegaSlide-DiT并没有声称可以在单个GPU上从头开始训练105B模型,也没有神奇地解决带宽限制;相反,它为高端工作站上的海量视频扩散模型的全参数适配提供了一条实用的路径。