初创公司如何从零开始训练出优秀的LLMs 精华

鉴于我们在Reka成功地培训了相当强大的多模态语言模型，许多人对从零开始建立基础设施并训练大型语言和多模态模型的经验特别感兴趣。

我在社交媒体上经常抱怨外部（Google之外）的基础设施和代码，这让人们对我在荒野中错过了什么，以及我对什么讨厌/喜欢非常好奇。所以终于有了这篇文章。这篇博客文章揭示了挑战和经验教训。

我希望这篇文章对许多人来说既有趣又有教育意义。

在荒野中训练LLMs（图片由 Dall-E 生成）

在LLMs时代的硬件抽奖

​训练模型的第一个必要条件是获取计算资源。这看起来似乎很简单，也足够容易。然而，最大的惊喜是计算资源提供商的不稳定性，以及集群、加速器及其连接的质量取决于来源的巨大差异。

人们总是认为这只是一个加速器选择（例如TPUs vs GPUs等）的问题/争论，而所有GPU集群都是平等的。但对我们来说，这很快就被证明是错误的。当我们尝试不同的服务提供商时，我们发现即使是相同的硬件，例如GPU（H100），硬件质量的差异也是巨大的。请注意，这里的硬件指的是整个集群的质量，而不一定是芯片或加速器本身。就像彩票一样。基本上：

并非所有硬件都是一样的。不同硬件提供商的集群质量差异如此之大，以至于如何训练出良好模型会成为一场真正的抽奖。简而言之，这是LLM时代的硬件抽奖。

更具体地说，我们从多家计算资源提供商租用了几个集群，每个集群都有数百至数千个芯片。我们看到的集群范围从尚可（只是一些烦人的问题，只需花费一些 工程师的时间就可以解决）到完全无法使用的集群，由于各种原因每隔几个小时就会失败。具体来说，一些集群的节点每隔N个小时就会出现故障，问题涉及到布线问题（其中N是不合理小的数字）、GPU硬件错误等。更令人惊讶的是，同一提供商的每个集群在稳健性方面也可能截然不同。

与此同时，即使其他一些集群可能拥有明显更稳定的节点，它们可能会遭受到I/O和文件系统糟糕的影响，甚至保存检查点都可能导致超时或花费大量时间降低集群利用率。其他一些计算资源可能需要完全不同的软件层才能运行，对于带有自己代码库的团队来说并不友好，需要额外的迁移成本来运行实验或大型作业。

没有完美的东西！但有些肯定比其他的要糟糕得多。

最令人沮丧的部分是什么呢？几乎不可能事先真正了解，尤其是在一切都在疯狂进行的情况下，人们将得到什么样的硬件以及体验的鲁棒性/容错性。

除此之外，你还无法知道供应商是否会延迟交货几个月，让你处于无法从其他来源获取资源的困境中数周甚至数月。有些提供商还会意外删除你的检查点¯\_(ツ)_/¯。

我提到了不同集群的模型浮点操作（Model Flop Utilisation，MFU）也会不同吗？如果不幸地找到布线不良或其他问题的提供商，这将是一笔不可忽略的计算资源浪费。具有非常次优文件系统的系统在团队成员开始在集群之间传输大量数据时，训练运行的MFU会瞬间下降。

每个服务提供商也提供不同级别的支持。这些支持从礼貌到漠不关心，从“ChatGPT风格”的套话回复再到责怪用户每一件事情都出错。

总的来说，我们尝试过的每个集群都感觉有着自己的氛围、挣扎和失败模式。几乎每个集群似乎都需要为自己的一系列问题提供自己的修复措施 - 有些问题比其他问题更容忍。也就是说，我们已经学到了备用方案的重要性，快速为任何集群找到修复措施可能是关键。

在过去的几个月里，我们已经做了很多工作，只是为了确保事物能够使用，例如，围绕监控工具、高效的检查点，以及各种其他优化，甚至安装我们的自定义文件系统以实现可扩展的数据存储 - 并且这只是实际需求的冰山一角。

这些工具组合带来了MFU 的显着改进，同时还最大限度地减少了糟糕硬件带来的停机时间。

GPU与TPU的对比

在Reka，我们大部分时间都是在GPU上训练我们的模型。就我个人而言，在Reka之前的Google生活中，我一直使用TPU进行大型语言模型训练。CUDA和nccl对我来说是最陌生的东西。（我只是从一个曾在Nvidia工作过的同事那里学到它的发音是“Nickel” 哈哈）

与我在Google使用TPU的经验相比，我完全惊讶于GPU的故障率。事实上，我实际上不太记得在Google使用TPU时出现过太多故障，即使是在大规模运行时，尽管我不确定我是否被这个简直太好的基础设施和专门的硬件团队的强大保护着。事实上，UL2 20B模型（在Google）是通过不小心让作业运行了一个月来训练的。它从未失败过。如果这是在GPU领域，它肯定会在最初的几天内就失败了。

话虽如此，我认为这更多地与管理加速器的硬件团队的能力有关，而不是与底层芯片有关。拥有良好的硬件支持（来自计算资源提供商）很重要。很多事情取决于他们是否真的很有能力，这加强了“硬件抽奖”的概念。

GPU领域感觉很奇怪。与在TPU架构中作为分布式训练的首要选择不同，在GPU领域，多节点训练似乎更像是事后想到的。在GPU领域，似乎不同的提供商以不同的方式将它们连接起来以实现多节点训练，这导致了在不同地方进行事情的方式存在很大差异。虽然我不是硬件方面的专家，但这是我得到的印象。

多集群设置的痛苦

我职业生涯的大部分时间都花在了谷歌的基础设施上，它主要运行在Borg、Xmanager和Colossus上，从任何地方都可以访问所有内容。因此，实际上需要在不同的集群中设置新环境的概念对我来说是陌生的。

在当前世界中，除非专门为单个位置的大量集群构建，否则拥有多个加速器池集群似乎是不可避免的。更具体地说，GPU供应（或缺乏）也自然导致了这种以集群采购的模式，其性质是分散的。训练大型模型还需要大量的数据，即使只是将它们移动一下也会带来很多不便。与此同时，通常复制数据也不是直接的，在极大规模上是难以承受的。

显然，理想情况下是有一种专门建立的编排层，可以将作业发送到不同的服务器。我相信许多大型AI优先公司通常都有某种基础设施来改善AI研究人员的生活质量。然而，对于一个刚起步的精益的新创企业来说，构建这种复杂和花哨的ML训练基础设施实际上是不太可能的。

目前，我们开发了许多内部工作流来减轻许多这些问题，并且正在继续朝着世界一流的实验基础设施的黄金标准迈进。

（有人告诉我，这种杂乱的设置或多或少是非顶级/大公司的常态）。

代码在野外

毫无疑问，我有史以来最喜欢的代码库是T5X和Mesh Tensorflow（命名张量太赞了），但这些选项很快就变得不可行，原因是：1）它们在谷歌之外得不到太多支持，2）它们已经有点被弃用了，3）对于我们团队中的非谷歌人员来说并不友好。

我们最终选择了一些常规的、看起来比较稳定且更受欢迎的选项（即，pytorch），这对大多数团队成员来说更易于接触（除了我哈哈）。在我最初的几个月里，我在pip、git、docker和所有这些野生的东西感到困惑。不过话说回来，我不太确定使用谷歌的代码库外部会有多稳定或用户友好（我猜这会非常令人讨厌）。

坦率地说，我不得不说外部代码库的质量明显落后于我在谷歌习惯的那些。主要是因为谷歌内部的代码库往往是由机器学习界的明星们亲自编写的（例如，Noam Shazeer、Barret Zoph、Adam Roberts、Hyung Won Chung等等），并且与我在外部尝试过的代码库相比，感觉更好（例如，优越的氛围）。特别是，当我尝试使用其他公司构建的东西时，我发现自己对代码质量感到非常恼火（有些比其他更糟糕🤗）。

另外，我从来不知道更改模型并行性的能力不是自动（免费）的，直到某些代码库要求我编写一个转换器来更改模型的并行性。对我来说这绝对是一个WTF时刻。

另一个引人注目的事情是这些代码库对大规模编码器-解码器训练甚至 prefixLM 训练的支持非常少。为此，即使 Flash Attention 也一直拒绝为 prefixLM 训练（即自定义掩码）提供支持，尽管出于某种原因对其 github 问题有合理的需求。

我知道我应该使用Jax。一个朋友刚刚因我使用pytorch而羞辱了我，但这是一个创业公司，我们决定快速行动。对此我很抱歉，下次我们会做得更酷一些。我对这个事实并不感到自豪。

少一点原则，多一点Yolo

系统地扩展模型通常需要以一种原则性的方式从小规模到大规模进行，即在多个阶段运行实验（从1B到8B到64B到300B等），然后选择出获胜者并持续扩展它们。在创业公司中，我们拥有的计算资源远远不足以进行大规模的参数扫描。最终，我们不得不依靠许多“Yolo”运行（幸运的是，结果还不错）。

最终，我们只需要进行了非常少量的小规模和较短的割裂实验，就能够获得强大的21B Reka Flash模型和7B边缘模型（以及即将推出的最大核心模型）。在有限数量的运行中找到一个可靠的配方是具有挑战性的，因为搜索空间非常巨大，必须同时改变许多变量。为了做到这一点，人们必须放弃大型科技公司的系统性，大量依赖于“Yolo”、直觉和本能。

值得庆幸的是，我（以及我们团队中的许多人）在我们的机器学习职业生涯中已经积累了相当多的直觉，以便在相当短的尝试次数内得到正确的结果。虽然我们之前在以往的工作中训练过非常好的模型，但是训练基础设施、数据、新思想的融入以及其他环境问题的差异仍然可能导致结果的巨大差异。尽管如此，强大的先验知识有助于显著减少搜索空间，并且可能是我们能够用如此少的尝试、资源和实验训练出非常强大的模型的最简单的解释之一。

简而言之

探索荒野中的事物是一次有趣的经历。不幸的是，这并不是一帆风顺的。计算资源的稀缺以及不可靠的计算资源提供商会使事情变得比预期的要困难得多，但我们很高兴用强大的技术实力度过了难关。

总的来说，这只是我们如何在不到一年的时间里创办一家公司、筹集一些资金、购买一些芯片并匹配 Gemini pro/GPT 3.5 并超越许多其他公司的故事的一小部分，而一切都从头开始。

还有更多要写的，比如数据管道、人工评估等，但这已经是一篇很长的文章了。下次再写。

关于作者：

Yi Tay是Reka的联合创始人兼首席科学家。Reka致力于构建先进的生成式 AI 模型。

此前，他曾是Google Brain的高级研究科学家，专注于大型语言模型和人工智能领域。在加入Brain团队之前，他是Google Research美国研究团队的技术负责人，致力于Transformer的扩展和架构研究，并为约20个产品发布做出了贡献。

他在Google任职期间，参与了许多产业级的大语言模型的开发，如PaLM、UL2、Flan-{PaLM/UL2/T5}、LaMDA/Bard、MUM等。他还为大型多模态模型，如ViT-22B和PaLI-X做出了贡献。他还是新PaLM-2和PaLM API的建模联合负责人。