鉴于我们在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等等），并且与我在外部尝试过的代码库相比，感觉更好（例如，优越的氛围）。特别是，当我尝试使用其他公司构建的东西时，我发现自己对代码质量感到非常恼火（有些比其他更糟糕