潞晨科技尤洋：如何用低成本做出类Sora？成功复现有四个要素

作者|GenAICon2024

2024中国生成式AI大会于4月18-19日在北京举行，在大会第二天的主会场AIInfra专场上，新加坡国立大学校长青年教授、潞晨科技创始人兼董事长尤洋以《技术共享：类Sora开源架构模型与训练细节》为题发表演讲。

值得一提的是，这个模型在低成本下训练而成，相比Sora耗费数千张H100GPU、花费数千万美元乃至数亿美元，Open-Sora的训练成本仅不到1万美元。

因此，Open-Sora模型也进行了全面的训练流程开源，开源了包括模型架构、模型权重、训练细节、数据处理在内的多项技术细节，让更多的开发者可以尝试Open-Sora模型，共同迭代与升级。

以下为尤洋的演讲实录：

本次演讲将分为几个部分。首先，我会简单介绍Open-Sora模型，以及与之相关的OpenAISora。需要明确的是，尽管我们称之为Open-Sora，也确实采用了类似于OpenAI技术报告中类似的技术，但它实际上是一个不同的模型。然后我将介绍Open-Sora的技术要点、性能表现，以及我们对未来发展的规划。

在介绍了Sora及其影响力后，第二部分介绍一下Open-Sora。

首先，需要了解模型的架构，比如我到底用的是Diffusion、Llama、GPT还是BERT，不同的架构决定我模型基本的骨架。

其次，一旦模型训练完成，分享训练得到的权重也是非常重要的。这意味着其他人可以直接拿来用，而不需要从头开始训练模型。例如Meta刚刚开放了Llama3，尽管4000亿参数版本还没有完全训练完成，但已经可以从中看到很好的效果。通过分享这些权重，社区可以快速地将模型部署到各种应用中。

第三点非常关键，它涉及到开源模型的透明度和可控性。

虽然现有的一些开源模型，如Llama1和Llama2，已经公开了模型参数和使用方式，但它们并没有公开训练过程的具体细节，包括超参数的设置。这导致了我们无法完全复现其预训练过程，也就是说，模型的预训练并不是百分百自主可控的。

此外，数据处理也是决定模型性能的一个关键因素。

我将展示一些我们的demo和教程，这将涵盖开源模型的几个重要组成部分。

从技术角度来看，Open-Sora模型采用了STDiT架构。我们选择STDiT的主要原因是考虑到成本效益。我们的目标是将Open-Sora的成本控制在1万美金或者更少。

STDiT架构的核心思想在于它包含时间维度的SelfAttention和空间维度的SelfAttention，这两个方面是分开处理的，而不是合并计算，这样的设计可以显著降低模型的训练和推理成本。相比于DiT模型，STDiT在成本上有着显著的优势，而且在相同的硬件条件下，其吞吐量也更高，这对于提升模型效率来说是非常有利的。

至于如何实现这一技术，考虑到成本的压力，我们可以将其分为三个阶段进行。

在成本受限的现实条件下，我们的策略是，首先找到一个不错的文本到图像生成模型，这类模型目前比较容易获取，市场上也有很多选择，当然我们也可以自己训练一个。

理论上我们有无限算力的话，我们应该直接拿最好的数据让它去训练。包括Llama、GPT，它们训练时也都参考了类似的思路，先在短序列上大规模训练，之后再在长序列或者更好的数据上去做微调，提升最终的模型生成质量。

具体而言，我们可以看一下三个阶段究竟是怎么操作的。

第一个阶段还比较简单，现在有很多文生图的模型，即便不自己训练，也可能找一些不错的文生图模型，它其实就是我的基准，我从起点开始去构造我的整个方案。我们改造StableDiffusion，可以快速把这件事完成。

我们的创新点有两部分，用了STDiT，有时间信息和空间信息，我们新加了时间上的Attention模块，因为本身已经有空间上的Attention模块。比如空间上就是S，时间上就是T，现在有一个S和T，S是已经训得差不多了，T刚刚开始。S相当于是一个初中生，T相当于是一个婴儿，但现在我们希望S和T都能达到大学生的水平。

有时我们会采用一种混合训练的方法，即将成熟度不同的模型一起训练。这种方法听起来可能有些冒险，因为S可能已经相当于一个初中生，而T可能还只是一个刚刚起步的婴儿。如果将它们放在一起训练，可能会担心它们无法跟上对方学习的节奏。

然而，现代的大型模型拥有庞大的参数量，这使得它们能够通过适当的调整迅速自适应不同的学习速度。在这种情况下，尽管S模型最初学习速度较慢，但T模型可以快速增长，最终两者都能迅速达到相同的水平，最终都能达到相当于大学生的能力水平。

即使我们没有自己的S模型，也可以利用一些现有的资源。当我们引入T模型后，通过适当的整合和调整，可以显著提升整体系统的性能。

这种策略体现了深度学习模型训练的灵活性和适应性，通过合理的设计和调整，即使是成熟度不同的模型也能够协同工作，最终实现性能的共同提升。

介绍完整体的模型、算法、设计流程之后，接下来考虑如何把成本降到最低。

要想把成本控制在1万美金左右，显然我们不能用太多的GPU，我们可以简单地做一笔数学计算。现在H800一台月租8万-10万人民币，假设有8台H800，每月的租金就要80万，如果用20台，每月的租金大概需要200万。要想一次性试验成本控制在10万以下，只能用8台H80064个H800GPU，就需要把速度、效率破到最高。

之前我们打造了Colossal-AI系统，从三个角度，高效的内存优化、N维并行系统、低延迟推理，通过Colossal-AI进一步实现2-4倍的加速。

训练过程中，一个关键因素是它们需要处理的序列长度通常非常长。无论是国内还是美国的大模型，研究者们都在努力扩展模型的序列长度，以期获得更高的预测精度。以GPT模型为例，其损失函数依赖于一个窗口的信息来预测下一个词的概率，窗口越大，即包含的信息越多，预测的准确性也就越高。

因此，将时间信息和空间信息进行分割处理是非常关键的一步。通过时空分割，我们可以显著降低计算成本和内存压力。具体来说，这意味着我们不是同时计算时间信息和空间信息，而是分步骤进行，先处理时间维度，再处理空间维度，这样可以大幅提升处理效率。

经过我们的优化之后，训练策略提升了很大。从右图可以看出，即使在8个GPU上训练速度也提升了16%，尤其在Encoder部分，计算密集型任务也实现了显著加速。

最后展示下我们的demo。我们的demo远差于OpenAI，主要有两个原因：

首先，我们的demo是在低成本条件下完成的，OpenAI使用了2000到4000个H100GPU，花费了五千万美元到两亿美元，而我们仅用了不到1万美金进行试验。在如此有限的预算下，我们取得的效果是可接受的。

以上是尤洋演讲内容的完整整理。