如何在不到一年的时间里创办一家公司、筹集资金、购买芯片,并搭建出追赶 Gemini pro/GPT 3.5 的 LLM?
很多人都对构建基础架构和训练大语言模型和多模态模型感到好奇,但真正走完「从零开始」这一流程的人很少。我们普遍认为,储备技术人才是前提,掌握核心算法是关键,但实际上,工程实践中冒出来的挑战,也实在令人头疼。
一年前,乘着大模型的热潮,Yi Tay 离开了工作 3 年多的谷歌,参与创办了一家名为 Reka 的公司并担任首席科学家,主攻大型语言模型。
在谷歌时,Yi Tay 参与过许多知名的大型语言模型和多模态模型工作,包括 PaLM、UL2、Flan-U-PaLM、LaMDA/Bard、ViT-22B、PaLI、MUM 等。即使经验如此深厚,他还是遇到了以往无法想象的困难。为了帮助更多创业者避雷,Yi Tay 在一篇博客中分享了自己踩过的那些「坑」。
「计算稀缺和不可靠的计算提供商使事情比预期困难得多,但我们凭借强大的技术实力渡过了难关。终于,我写了这篇博文,揭示了其中的一些挑战和经验教训。我希望这篇文章对很多人来说都是有趣或有教育意义的。」
文章发出后,得到了众多技术创业者的议论和转发。
连 Andrej Karpathy 也深有同感:
成熟的公司有专门的团队维护集群。随着规模的扩大,集群已经脱离了工程学的范畴,变得更加生物化,因此需要专门负责「硬件健康」的团队。
「照看」训练运行是一项令人沮丧的训练大型模型日常生活体验。你需要仔细监控运行的生命体征:损失峰值、数值问题、吞吐量、梯度规范、策略熵等。每当运行性能下降或趋于平稳时(可能经常发生),你都要快速查找堆栈跟踪,看看发生了什么。你必须快速完成这项工作,否则可能会有 10000 个 GPU 闲置。通常,这是你从未见过的新的、奇特的、可怕的错误,所以你需要寻求帮助,看看是否有人能发现问题所在。
最严重的错误发生在凌晨 4 点。通常没人能看到,所以你只能禁止一些看起来有点可疑的节点,并尝试重新启动运行。有时,运行失败只是因为你当天没有得到神的眷顾,所以你在启动命令中加入了 while True: 循环。潜在的问题可能多种多样,从某些 GPU 发热过高、偶尔突然做错乘法运算到某些路由器宕机导致网络文件系统 I/O 减少,再到数据中心的某个人在未沟通的维护过程中物理断开电线连接。有的问题你甚至永远不会知道。
也有人发现了亮点:Yi Tay 所说的「荒野」(Wild)意思是「谷歌之外的公司」。
要是从基础设施和硬件的角度来说,能媲美谷歌的团队还真是不多。
现在,让我们一起看看博客内容:
LLM 时代的硬件彩票
训练模型的首要条件是获得计算能力。这看似简单易行,然而,最大的惊喜却是计算提供商的不稳定性,以及集群、加速器及其连接质量因来源不同而存在的巨大差异。
人们总以为这只是一个加速器选择的问题 / 争论(TPU 与 GPU 等),所有 GPU 集群都是一样的。我们的体验是,这很快就被证明是错误的。
我们对不同的服务提供商进行了抽样调查,发现即使是相同的硬件,即 GPU(H100),硬件质量的差异也非常大。请注意,这里的硬件指的是集群的整体质量,而不一定是芯片或加速器本身。整体感觉就像购买彩票一样。
更具体地说,我们从几家计算提供商那里租用了几个集群,每个集群都有数百到数千个芯片。我们所见过的集群有的还过得去(只存在一些小问题,但只需花几个小时的时间就能解决),有的则完全无法使用,每隔几个小时就会因各种原因出现故障。具体来说,有些集群的节点每隔 N 个小时就会出现故障,问题包括布线问题(N 小得不合理)、GPU 硬件错误等。更令人惊讶的是,同一家提供商的每个集群在鲁棒性方面也可能存在巨大差异。
同时,即使其他一些集群的节点明显更稳定,它们也可能存在 I/O 和文件系统不佳的问题,甚至连保存检查点都可能导致超时,或耗费大量时间来降低集群利用率。其他一些计算资源甚至需要完全不同的软件层才能运行,而且对自带代码库的团队不友好 — 运行实验或大型工作需要额外的迁移成本。
凡事都不会尽善尽美,但可以确定的是,提供商的服务质量是参差不齐的。
最令人沮丧的是什么?几乎不可能真正提前知道,尤其是在万事俱备的情况下,会得到什么样的硬件,以及体验会有多么强大 / 容错性如何。
此外,如果供应商不能按时交货,将装备时间推迟几个月,导致用户在数周或数月内无法从其他来源采购,你更无从得知。有些供应商还会不小心删除你的检查点。
我有没有说过,不同的集群会有不同的模型翻转利用率(MFU)?如果你不幸找到了一个节点布线不良或存在其他问题的提供商,那么浪费的计算量是不可忽视的。如果系统的文件系统非常不理想,那么当团队成员开始跨集群传输大量数据时,训练运行的 MFU 就会下降。
每个服务提供商的售后水平也各不相同。从礼貌客气到不冷不热,从「对话式」的预制回复到将所有问题都归咎于用户,不一而足。
总之,我们尝试过的每一个集群都有自己的风格、斗争和失败模式。而且,几乎每个集群都需要自己的热修复程序来解决一系列问题。尽管如此,我们还是认识到故障安全是非常重要的,为任何集群找到快速的热修复方案都是关键所在。
在过去的几个月里,我们构建了许多工具,以确保一切都可用,例如,围绕监控、高效检查点和其他各种优化的工具,甚至安装了我们的自定义文件系统,以实现可扩展的数据存储,而这只是实际需求的冰山一角。
这些工具的组合为 MFU 带来了非同小可的改进,同时也最大限度地减少了在硬件条件恶劣的情况下的停机时间。
GPU vs TPU
就我自己的公司来说,大部分时间都在使用 GPU 训练模型。不过在加入 Reka 之前,我在谷歌的大型语言模型训练中一直使用 TPU。CUDA 和 nccl 对我来说是最陌生的东西 (我是从一位曾在 Nvidia 工作的同事那里才知道它的发音是 Nickel 的)。
与我在谷歌使用 TPU 的经历相比,GPU 的故障率让我完全大吃一惊。事实上,我并不记得 TPU 发生过很多故障,即使是在大型运行中也是如此,不过我不确定,自己是否只是因为拥有出色的基础架构和专门的硬件团队才不知道这一点。事实上,谷歌的 UL2 20B 模型是通过意外运行一个月来训练的。如果是在 GPU 领域,它肯定会在最初几天内就失败。
话虽如此,我认为这可能更多与管理加速器的硬件团队的能力有关,而不是底层芯片。拥有良好的硬件支持(来自计算提供商)非常重要。而这在很大程度上取决于他们是否真的有能力,于是,又印证了「硬件彩票」的概念。
GPU 领域给人的感觉很奇怪。与分布式训练在 TPU pods 上的一等公民地位相比,多节点训练更像是一种事后考虑。在 GPU 领域,感觉就像不同的提供商以不同的方式将它们连接起来,以实现多节点训练,这导致不同地方的做法差异很大。
虽然我不是硬件专家,但这就是我的真实印象。
多集群设置的痛苦
我职业生涯的大部分时间都是在谷歌基础架构上度过的,这些基础架构主要运行在 Borg、Xmanager 和 Colossus 上。因此,必须在不同的集群中建立新环境的概念对我来说非常陌生。
在当今时代,拥有多个加速器池集群似乎是不可避免的,除非专门在一个地点建立大量的加速器池。更具体地说,GPU 的供应(或供应不足)自然而然地造成了这种集群采购模式,在这种模式下,事物的性质是支离破碎的。训练大型模型还需要大量的 TB 级数据,即使只是移动数据也会带来诸多不便。同时,复制数据通常也不是一件简单的事情,而且在超大规模的情况下,复制数据的成本也很高。
显然,最理想的情况是建立某种编排层,专门将作业发送到不同的服务器。我相信,许多注重人工智能的大公司一般都有某种基础设施,以提高研究人员的生活质量。但是,对于一家初创公司来说,在开始阶段建立这种复杂而花哨的 ML 训练基础设施其实是不可能的。
目前,我们公司开发了许多内部工作流程来缓解这些问题,并继续朝着世界级实验基础设施的黄金标准迈进。(有人告诉我,对于非顶级 / 大型公司来说,这种简陋的设置或多或少是一种常态)。
野生代码
众所周知,一直以来我最喜欢的代码库是 T5X 和 Mesh Tensorflow,但它们有一些缺点:
1)它们在 Google 之外没有得到那么多的支持;
2)它们有点被弃用了;
3)它们对我们团队中的非 xoogler 不友好。
我们最终选择了一些普通的、看似稳定且更流行的东西,即 pytorch。pytorch 对团队中的大多数人(除了我)来说更容易使用。
在最初的几个月里,我对 pip、git、docker 和所有「野生(wild)」的东西感到困惑。话又说回来,我不能 100% 确定在外部使用 google 代码库有多稳定或多用户友好。
坦率地说,外部代码库的质量明显落后于我在谷歌习惯使用的代码库。主要是因为谷歌内部的代码库往往是由 ML 大牛自己编写的(例如 Noam Shazeer、Barret Zoph、Adam Roberts、Hyung Won Chung 等人),并且与我在外部尝试过的代码库相比感觉更好。
另外,我从来不知道更改模型并行性的能力不是自动(免费)的,直到某些代码库要求我编写一个转换器来更改模型的并行性。对我来说,这绝对是一个「WTF 时刻」。
令人惊讶的是,这些代码库对大规模编码器 - 解码器训练甚至 prefixLM 训练的支持非常少。
少一点原则,多一点 Yolo
系统地扩展模型通常需要有原则地从小到大,即分多个阶段(1B→8B→64B→300B 等)进行实验,然后选出获胜者并不断扩展它们。在初创公司中,我们执行大规模扫描来检查超参数的计算量要少得多。最后,我们不得不多次运行 Yolo,幸运的是结果很好。
最终,我们只需要极少数的较小规模和较短的烧蚀运行即可获得强大的 21B Reka Flash 和 7B 边缘模型,以及我们即将推出的最大核心模型。在运行次数非常有限的情况下找到可靠的方案具有挑战性,并且考虑到搜索空间极其巨大,需要立即更改许多变量。为了做到这一点,人们必须放弃大型科技公司的系统性,而在很大程度上依赖「Yolo」、直觉和本能。
值得庆幸的是,我以及我们团队中的许多人在我们的 ML 职业生涯中已经积累了相当多的「直觉」,可以在很短的尝试时间内得到正确的结果。虽然我们在之前的工作中训练过非常好的模型,但训练基础设施、数据、新想法的融合和其他环境问题的差异仍然可能导致结果的巨大差异。也就是说,强大的先验有助于显著减少搜索空间,这可能就是我们能够通过如此少的试验、资源和实验训练出真正强大的模型的原因之一。