Open-Sora 旨在为扩散模型提供一个高速训练框架。在 64 帧 512x512 视频上训练时，我们可以实现 55% 的训练速度加速。我们的框架支持训练 1分钟1080p视频。
Open-Sora 通过以下方式提高训练速度：
- 内核优化，包括 flash attention, 融合 layernorm 内核以及由 colossalAI 编译的内核。
- 混合并行性，包括 ZeRO。
- 用于更大批量的梯度检查点。
我们在图像上的训练速度可与 OpenDiT 相媲美，这是一个加速 DiT 训练的项目。训练速度是在批处理大小为 128、图像大小为 256x256 的 8 个 H800 GPU 上测量的。
| 模型 | 吞吐量 (img/s/GPU) | 吞吐量 (tokens/s/GPU) |
|---|---|---|
| DiT | 100 | 26k |
| OpenDiT | 175 | 45k |
| OpenSora | 175 | 45k |
我们的 STDiT 采用时空注意力对视频数据进行建模。与直接全神贯注在 Dit 相比，我们的 STDiT 随着帧数的增加而更有效率。我们当前的框架仅支持序列超长序列的并行性。
训练速度是在 8 个 H800 GPU 上测量的，应用了加速技术，GC 表示梯度检查点。 两者都具有像 PixArt 一样的 T5 调节。
| 模型 | 设置 | 吞吐量 (sample/s/GPU) | 吞吐量 (tokens/s/GPU) |
|---|---|---|---|
| DiT | 16x256 (4k) | 7.20 | 29k |
| STDiT | 16x256 (4k) | 7.00 | 28k |
| DiT | 16x512 (16k) | 0.85 | 14k |
| STDiT | 16x512 (16k) | 1.45 | 23k |
| DiT (GC) | 64x512 (65k) | 0.08 | 5k |
| STDiT (GC) | 64x512 (65k) | 0.40 | 25k |
| STDiT (GC, sp=2) | 360x512 (370k) | 0.10 | 18k |
使用 Video-VAE 在时间维度上进行 4 倍下采样时，24fps 视频有 450 帧。STDiT(28k tokens/s) 和 DiT 对图像 (高达 45k tokens/s) 两者之间的速度差距主要来自 T5 和 VAE 编码，以及时间注意力。
在训练过程中，文本由 T5 编码，视频由 VAE 编码。通常有两种方法可以加速训练：
- 提前预处理文本和视频数据并保存到磁盘。
- 在训练过程中对文本和视频数据进行编码，并加快编码过程。
对于选项 1，一个样本的 120 个令牌需要 1M 磁盘空间，而 64x64x64 的潜在可能需要 4M。考虑训练 包含 10M 视频剪辑的数据集，所需的总磁盘空间为 50TB。我们的存储系统目前还没有准备好 这种数据规模。
对于选项 2，我们提高了 T5 速度和内存要求。根据在OpenDiT，我们发现 VAE 消耗了大量的 GPU 内存。因此，我们 将批大小拆分为较小的批大小，以便进行 VAE 编码。使用这两种技术，我们可以大大加快训练速度。
训练速度是在 8 个带有 STDiT 的 H800 GPU 上测量的。
| 加速模式 | 设置 | 吞吐量 (img/s/GPU) | 吞吐量 (tokens/s/GPU) |
|---|---|---|---|
| Baseline | 16x256 (4k) | 6.16 | 25k |
| w. faster T5 | 16x256 (4k) | 7.00 | 29k |
| Baseline | 64x512 (65k) | 0.94 | 15k |
| w. both | 64x512 (65k) | 1.45 | 23k |