作者 | Yi Tay 编译 | 苏宓
出品 | CSDN（ID：CSDNnews）
一年前的 3 月，在 Google Brain 和 AI 部门工作了 3.3 年的 Yi Tay 做了一个重要的决定：辞职！然后开启下一次冒险之旅。
在他看来，辞职的确有种“毕业”的感觉。“毕竟，我从 Google、从我出色的同事、导师、经理等那里学到了很多东西”，Yi Tay 说道。
回看在 Google 任职期间的工作量，Yi Tay 总共写了约 45 篇论文，其中约 16 篇是第一作者（第一作者或共同第一作者等）。就实际影响而言，他大概促成了约 20 项产品的发布。
同时，他一直深耕于 Transformers 架构的研究与应用，为 PaLM、UL2、Flan-2 和 Bard 等许多业界领先的 LLM 以及 PaLI-X 和 ViT-22B 等多模态模型做出了贡献。值得一提的是，Yi Tay 还是 PaLM-2 和 PaLM-2 API 建模的共同负责人。
在从 Google Brain 离开的一年间，Yi Tay 和很多前同事一样，选择了一家创业公司，担任联合创始人兼首席科学家。
其所在的初创公司名为 Reka，除了有 Yi Tay 的加入，还有来自 Google、百度和 Meta、微软等大厂出来的研究员共同参与。
在不到一年的时间里，Yi Tay 的工作量可谓是翻了好几番，其团队不仅让 Reka 拿到了 5800 万美元的首轮融资，也发布了 21B 的多模态语言模型 Reka Flash，其性能与 Gemini Pro 和 GPT-3.5 相媲美。与此同时，团队还推出了一个参数更少、效率更高的 7B 模型 Reka Edge，适用于资源受限的场景。
整体而言，从资源丰富、人才济济的大厂到人少事多的创业公司，Yi Tay 在这一年趟过了很多「坑」，也有太多的话想说了。
于是，他将其总结为一篇《作为一家初创公司，完全从零开始训练优秀的 LLM》的文章，围绕在不同计算提供商处采购计算资源带来的差异和经验、讨论“野蛮生长”的基础设施/代码以及过渡到其在 Google 时的做法、训练模型时的新思维等维度进行了深刻的思考与分享，其直言：自己训练模型真的太难了！
以下是 Yi Tay 的博客内容：
在 Reka 成功训练出一些强大的多模态语言模型（如 Reka Flash、Reka Edge）后，许多人对从零开始构建基础设施，以及训练大型语言模型和多模态模型的经历表现出浓厚的兴趣。
我在社交媒体上经常抱怨（除了 Google 之外的）基础设施和代码，这引发了人们对我在这片“荒野”中错过了什么、我对这个环境的喜好和厌恶的真实好奇。因此，这里终于有了一篇博文。这篇博文将阐述从零训练大模型过程中面临的挑战以及所学到的经验教训。希望这篇博文对许多人来说既有趣又有教育意义。
LLM 时代的硬件采购就像买“彩票”，一切存在未知
训练模型的首要条件是获取计算资源。 这一点没有什么可质疑的，也相对容易。
然而，整个过程中存在的最大意外在于计算资源提供商的不稳定性，以及集群、加速器及其连接的质量在不同来源之间存在的巨大差异。
人们总是认为这只是事关选择加速器（TPUs 与 GPU）的问题或争论，而所有 GPU 集群都是平等的。对我们而言，这一观点很快就被证明是错误的。
在对不同的服务提供商取样时，我们发现即使对于相同的硬件，例如 GPU（H100s），硬件质量也存在巨大的差异。 请注意，这里的硬件指的是整体集群质量，而不一定是芯片或加速器本身，这感觉就像买“彩票”一样。基本上：
并非所有硬件都是相同的。不同硬件提供商带来的集群质量可能天差地别，因此在训练良好模型时要经历多少痛苦就像是在刮彩票一样，一切未可知。一句话概括，这是在 LLM 时代的硬件彩票。
更具体地说，我们从多家计算资源提供商那里租赁了一些集群，每个集群有数百到数千个芯片。我们看到的集群有的还能将就说得过去（一些可以通过较小的软件工程师解决的烦人问题），也有的完全无法使用，由于各种原因每隔几小时就会出故障。详细来说，一些集群的节点每隔 N 小时就会出现故障，问题涉及到布线问题（其中 N 值非常小）、GPU 硬件错误等。
更令人惊讶的是，即使是同一提供商的每个集群在稳健性方面也可能大不相同。
与此同时，即使其他一些集群的节点可能更加稳定，它们也可能受到 I/O 和文件系统不佳的影响，甚至保存检查点都可能导致超时，或耗费大量时间来降低集群利用率。还有其他一些计算资源需要完全不同的软件层才能运行，这对于自带代码库的团队而言非常不友好，他们往往需要额外的迁移成本来运行实验或大型任务。
没有什么是完美的！但有些提供商的服务质量确实比其他的要糟糕得多。
最令人沮丧的部分是什么？几乎不可能提前知道，尤其是在什么都准备好的情况下，自己将会得到什么样的硬件，以及体验将会有多么健壮/容错，都不知道。
此外，你也无法判断供应商是否会按时交付，有时候他们还会推迟几个月交付，让人在数周或数月内无法从其他渠道获取计算资源。甚至有些供应商还会意外删除你的检查点。
再来看看，对于不同集群，你还会得到不同的模型算力利用率（Model FLOPs Utilization， MFU）吗？
当然，如果运气不好找到一个布线混乱或存在其他问题的供应商，这将是一笔不可忽视的计算资源浪费。当团队成员开始跨集群传输大量数据时，如果系统的文件系统非常不理想，训练运行的 MFU 会瞬间下降。
每个服务提供商还提供不同级别的支持。这些支持从礼貌到漫不经心，从“ChatGPT风格”的预设回复到将每一件出错的事情归咎于用户。
总的来说，我们尝试的每个集群都感觉有它们自己的氛围、斗争和失败模式。几乎每个集群似乎都需要对其自身的一系列问题进行修复。话虽如此，我们已经了解到，设置故障保护很重要，并且为任何集群找到快速修复的方法可能是关键。
在过去的几个月里，我们已经做了很多工作，只是为了确保事情能够使用，例如监控周围的工具、高效的检查点，以及各种其他优化，甚至包括安装我们的定制文件系统以实现可伸缩的数据存储——这只是实际所需的冰山一角。
GPU vs. TPU
在 Reka，我们主要在 GPU 上训练我们的模型。
就我个人而言，我在 Reka 之前的 Google 工作历程里，大型语言模型的训练一直在使用 TPU。CUDA 和 nccl 对我来说是很陌生的东西。（我只是从一个曾在 Nvidia 工作的同事那里学到它的发音是“Nickel”，哈哈）
与我在 Google 使用 TPU 的经验相比，我对 GPU 的故障率感到完全震惊。
实际上，我不太记得 TPU 在大规模运行中经常失败，尽管我不确定我是否因为当时有极好的基础设施和专门的硬件团队而被保护免于了解这一点。
在 Google，UL2 20B 模型（https://blog.research.google/2022/10/ul2-20b-open-source-unified-language.html）是通过意外将作业运行一个月而训练的。它从未失败过。如果这是在 GPU 领域，它肯定会在开始的前几天内就会失败。
也就是说，我认为这更多地与管理加速器的硬件团队的能力有关，而不仅仅是底层芯片的问题。
拥有良好的硬件支持（来自计算资源提供商）很重要。因此，很多事情都取决于它们是否真正有能力，这加强了“硬件彩票”的概念。
GPU 领域感觉奇怪。与 TPU 架构上将多节点训练作为一等公民的分布式训练相比，在 GPU 领域，它似乎更像是事后考虑的。在 GPU 领域，感觉就像不同的提供商以不同的方式布线，以启用多节点训练，导致在不同地方的做事方式存在很大的差异。
虽然我对硬件不是专家，但这是我的真实印象。
多集群设置的痛苦
我在职业生涯的大部分时间里都在 Google 基础设施上工作，主要运行在 Borg、Xmanager 和 Colossus 上，几乎可以从任何地方访问所有资源。因此，对我来说，在不同集群中实际设置新环境的概念是陌生的。
在当前的世界中，除非专门在单一位置为大量集群构建，否则拥有多个加速器池集群似乎是不可避免的。更具体地说，GPU 供应（或缺乏供应）自然也导致了这种集群采购模式，其中事务本质上是碎片化的。
训练大型模型还需要大量 TB 的数据，即使只是将它们移动一下也可能会带来很多不便。同时，在极大规模的情况下复制数据通常也不是简单的事情。
显然，最理想的情况是建立某种协调层，专门将任务发送到不同的服务器。我相信，许多以 AI 为主的大公司一般都有某种基础设施，来提高人工智能研究人员的生活质量。
但是，对于一家精益求精且刚起步的初创公司来说，在开始阶段建立这种复杂而高级的 ML 训练基础设施确实不太可能。
目前，我们开发了许多内部工作流程来缓解这些问题，并继续朝着世界一流的实验基础设施的黄金标准迈进。
(有人告诉我，对于非顶级/大型公司来说，这种简陋的设置或多或少是一种常态）。
野生代码
我的最爱的代码库是 T5X 和 Mesh Tensorflow，但由于以下原因，这些选项很快就变得不可行：
它们在谷歌之外得不到足够的支持；
它有点过时了；
它们对我们团队中不是 xooglers 的人不友好。
最终，我们选择了看起来更稳定、流行（如 Pytorch）的东西，这样大多数团队成员都能够更容易地使用（除了我）。
在最初的几个月里，我被 pip、git、docker 以及所有这些狂野的东西弄得晕头转向。不过，我也不能百分之百确认在外部使用 Google 代码库会有多稳定或者对用户有多友好（我猜应该会相当糟糕）。
坦率地说，我必须说外部代码库的质量明显落后于我在 Google 习惯的代码库。主要因为 Google 内部的代码库通常由 ML 行业的大佬们亲自编写（例如 Noam Shazeer、Barret Zoph、Adam Roberts、Hyung Won Chung 等），与我在外部尝试过的代码库相比，感觉更好（更卓越的氛围）。特别是，我发现自己在接触其他公司构建的东西时，发现自己对代码质量感到非常恼火（有些产品比其他产品差很多）。
而且，我以前从来不知道改变模型并行性的功能不是自动的（免费的），直到有些代码库要求我编写一个转换器来更改模型的并行性。对我来说确实是一个令人费解的时刻。
另一个令人震惊的事情是，这些代码库对大规模编码器-解码器训练，甚至是prefixLM 训练的支持是多么的有限。为此，尽管在 Github 问题上提出了合理的要求，但由于种种原因，即使是 flash attention 也一直拒绝提供对 prefixLM 训练（即自定义掩码）的支持。
我知道我应该使用 Jax。一个朋友因我使用 pytorch 而责备了我，但这是一家初创公司，我们决定要快速行动。
很抱歉，我们下次会做得更好。我并不以此为荣。
少一点原则，多一点 Yolo
要系统地扩展模型，一般需要有原则地从小到大进行扩展，即分多个阶段（1B->8B->64B->300B 等）进行实验，然后挑选优胜者并不断扩展其规模。
在初创公司中，我们的计算能力远不足以执行这些大规模的扫描来检查 hparams。最后，我们不得不使用许多 Yolo run（幸运的是结果还不错）。
最终，我们只用了极少量的较小规模和较短时间的消融运行，就获得了强大的 21B Reka Flash 和 7B Reka Edge 模型（以及我们即将推出的 Reka Core 模型）。
用非常有限的运行次数找到可靠的方法是一项挑战，鉴于搜索空间大得离谱，需要同时改变许多变量。为此，我们必须放弃大型科技公司的系统性，大量依靠 “Yolo”、直觉和本能。
值得庆幸的是，我（以及团队中的许多人）在我们的 ML 职业生涯中积累了大量的“直觉”经验，从而在很短的时间内就能获得正确的结果。
虽然我们在以前的工作中已经训练出了非常好的模型，但训练基础架构、数据、新想法的融入以及其他环境问题的差异仍会导致结果出现非同小可的差异。也就是说，强大的先验有助于大大缩小搜索空间，这可能也是我们为什么能在如此少的试验、资源和实验中训练出真正强大的模型的最简单的解释之一。
一言以蔽之
在荒野中摸索是一次有趣的经历。不幸的是，它并非毫无痛苦。计算资源的稀缺和计算提供商的不可靠，让事情变得比想象的要困难得多，但我们很高兴我们凭借强大的技术力量挺了过来。
总而言之，这只是我们如何创办公司、筹集资金、购买芯片、匹配 Gemini pro/GPT-3.5 并在不到一年的时间内超越其他公司、从零开始构建一切故事的一小部分。
未来还有更多内容想和大家分享，如数据管道、人工评估等，但这已经是一篇很长的文章了，期待下一次的分享。
原文链接：https://www.yitay.net/blog/training-great-llms-entirely-from-ground-zero-in-the-wilderness
