你敢相信吗?一位前谷歌大佬,离职成立公司,不到一年,从头训练出了“GPT3.5”/“Gemini Pro”,注意,后者是多模态大模型!
本文主人公Yi Tay,是一位市面上非常抢手的高性能大模型的大拿。他曾在谷歌Google Brain担任高级研究科学家,专注于大型语言模型和人工智能的研究。在Google任职期间,曾经为业内许多知名的大型语言模型做出了贡献,例如PaLM、UL2、Flan-{PaLM/UL2/T5}、LaMDA/Bard、MUM等。
另外,Yi还参与了大型多模态模型如ViT-22B和PaLI-X的研究,负责了新模型PaLM-2和PaLM API的建模工作。
去年3月,Yi离开了谷歌,创办了一家大模型公司Reka,一直追求打造出令人惊叹的前沿生成模型。
不到一年的时间,从一张卡都没有,到推出了可以匹敌GPT3.5/Gemini Pro的大模型Reka。大模型训练、多模态大模型何其艰难?这期间,究竟发生了怎样的离奇的事情呢?
Yi Tay 就此分享了几点挑战和教训,比如GPU问题百出,不如TPU、野鸡代码的折磨,“多一些YOLO,少一些原则”等等,非常有意思,值得诸君深思。
买算力如同买彩票!
训练模型的首要条件是获取计算能力。这看起来再简单不过了,但事实上这就好比买彩票一样。
你的算力供应商是不固定的,集群和加速器及其连接的质量也随着他们各自不同的厂商而带来巨大差异。
你可能会问,那我们就选择同一型号的GPU、TPU加速器等,集群也配置对等,不就完事了吗?
事实总是啪啪打脸,当我们对不同的服务提供商进行采样时,我们被结果震惊到了:
即使对于相同的硬件,即 GPU(我们用的H100),质量的差异也很大。请注意,这里的硬件指的是整体集群质量,而不一定是芯片或加速器本身。
这就跟买彩票一样。基本上:
并非所有硬件都是一样的。不同硬件提供商的集群质量差异如此之大,以至于这实际上是一场彩票,与训练好模型需要经历多少痛苦有关。简而言之,算力就是大模型时代的硬件彩票。
事情具体是这样的,我们从多家计算提供商那里租用了一些集群,每个集群都有数百到数千个芯片。
集群问题百出,有的还说得过去,只是一些烦人的问题,只需花费一些 SWE 时间就可以解决,而有的集群则完全不可用,每隔几个小时就会失败,而且让人抓马的是,原因也各种不一样。
比如,一些集群的节点每 N 小时就会出现故障,其问题就包括到底是布线问题(其中 N 过小)还是GPU 硬件出错误等。
最为让人百思不得其解地是,同一厂商的每个集群在稳健性方面也可能存在巨大差异。
同时,即使其他一些集群可能拥有更稳定的节点,它们也可能会受到 I/O 和文件系统较差的影响,甚至保存检查点也可能导致超时或大量时间减少集群利用率。
除了这些,有的不同供应商来源的算力甚至需要完全不同的软件层才能运行,并且对于自带代码库的团队非常不友好 ,因为这就意味着需要额外的迁移成本来运行实验或大型作业。
最令人沮丧的部分?几乎不可能真正提前知道,尤其是在一切逼疯你的问题全都来临时,究竟需要获得什么样的硬件?体验的鲁棒性/容错性又该如何预估和保证?
总之一句话,情况没有最差的,只有更差的!
就比如,供应商也有延迟和放鸽子的情况,你无法判断供应商是否能按时交货,而且仅仅只是发货延迟了几个月,供应商自己也尴尬,他们也无法从其他来源采购,这种延迟情况从数周到数月不等。此外,某些提供商还会意外地删除你的检查点文件。
这还没完,对于不同的集群,您还会获得不同的模型失败率利用率 (MFU)!如果不幸找到一个节点布线不良或存在其他问题的提供商,这就会导致一场价值不菲的算力浪费。
再有,当团队成员开始跨集群传输大量数据时,具有非常次优文件系统的系统的训练运行 MFU 将会下降。
此外,所有的服务供应商的服务态度也都分三六九等,都有着不同级别的支持,从礼貌到漠不关心,从“chatgpt 式”的预设回复到将每一件出错的事情归咎于用户。
一整套搞下来,最大的感触就是,我们尝试过的每个集群,都感觉它们有自己的氛围、挣扎和失败模式。几乎每个集群都需要针对自己的一系列问题进行热修复——有些问题比其他问题更容易忍受。
总之,就是为故障做保险非常重要,但关键之处在于如何为所有集群找到快速修复的方法。
在过去的几个月里,我们构建了很多东西只是为了确保东西可用,例如,围绕监控的工具、高效的检查点和各种其他优化,甚至安装我们的自定义文件系统以实现可扩展的数据存储,并且这只是实际需求的冰山一角。
这些工具组合带来了 MFU 的显着改进,同时还最大限度地减少了糟糕硬件带来的停机时间。
跟TPU相比,GPU简直菜鸡
在 Reka,我们的模型大部分时间都在 GPU 上进行训练。就我个人而言,在 Reka 之前的 Google 生活中,我一直在使用 TPU 进行大型语言模型训练。CUDA 和nccl对我来说是最陌生的东西。(我从一位曾经在 Nvidia 工作的同事那里得知后者的发音为“Nickel”)
与我在谷歌使用 TPU 的经历相比,GPU 的故障率让我完全大吃一惊。事实上,我印象中TPU 即便在大规模运行时也没有失败过,可能的确是谷歌基础设施非常出色,拥有着绝对稳健性,而且谷歌拥有着一支专门的硬件团队。
事实上,UL2 20B模型(在 Google)是通过让任务无特别维护地情况下运行一个月来进行训练的,它从未失败过。如果这要是换成 GPU ,那么它肯定会在最初几天内就会宕掉。
也就是说,我认为这可能更多地取决于管理加速器的硬件团队的能力,而不是底层芯片。拥有良好的硬件支持(来自计算提供商)非常重要。这在很大程度上取决于他们的实际能力,这强化了“硬件彩票”的概念。
GPU 领域感觉很奇怪。感觉多节点训练更像是事后的想法,而不是作为 TPU pod 上的一等公民的分布式训练。在 GPU 领域,感觉好像不同的提供商以不同的方式连接它们以实现多节点训练,这导致不同地方的工作方式存在很大差异。
虽然我不是硬件专家,但这就是我的印象。
多集群设置的痛苦
我职业生涯的大部分时间都花在了 Google 基础设施上,它主要运行在Borg、Xmanager和Colossus上。基本上任何集群都是这样的配置。然而,创业后,才意识到不同集群必须要实际设置新环境,这种概念对我来说是陌生的。
在当今世界,拥有多个加速器池集群似乎是不可避免的,除非专门在一个位置为大量集群构建。更具体地说,GPU 供应(或缺乏)也自然导致了这种集群采购模式,其中事物本质上是碎片化的。训练大型模型还需要大量 TB 的数据,即使只是移动它们也会带来很多不便。同时,在超大规模的情况下复制数据通常也不让人望而却步的。
显然,这里的理想情况是某种专门用于将任务发送到不同服务器的编排层。我相信许多AI优先的大型公司通常都配有某种基础设施来改善人工智能研究人员的工作质量。然而,对于一个精益的新初创公司来说,一开始就不可能构建这种复杂而精美的训练基础设施。
目前,我们最终开发了许多内部工作流程来缓解其中的许多问题,并继续朝着世界一流实验基础设施的黄金标准迈进。
当然,有人告诉我,这种杂乱的设置或多或少是非顶级/大公司的常态。
野鸡代码的痛苦
众所周知,我一直以来最喜欢的代码库是T5X和Mesh Tensorflow(名为tensor ftw),但这些选项很快就变得不可行,理由有三点:1)它们在 Google 之外没有得到那么多的支持,2)它们有点被弃用了, 3)他们对我们团队中非 xoogler 的人不友好。
我们最终选择了一些普通的、看似稳定且更流行的东西(即 pytorch),对团队中的大多数人(除了我)来说更容易使用。在最初的几个月里,我对 pip、git、docker 和所有这些野生的东西感到困惑。
话又说回来,我并不能 100% 确定在外部使用 Google 代码库有多稳定或多用户友好。
坦率地说,我不得不说外部代码库的质量明显落后于我在 Google 习惯的代码库。主要是因为 Google 内部的代码库往往是由 “ML 摇滚明星”自己编写(例如,Noam Shazeer、Barret Zoph、Adam Roberts、Hyung Won Chung 等人),并且与我在外部尝试过的代码库相比,感觉更好(例如,优越的氛围) 。特别是,当我涉足其他公司构建的东西时,我发现自己对代码质量非常恼火。
另外,我从来不知道更改模型并行性的能力不是自动(免费)的,直到某些代码库要求我编写一个转换器来更改模型的并行性。对我来说这绝对是一个WTF时刻。
另一个引人注目的事情是,这些代码库对大规模编码器-解码器训练甚至 prefixLM 训练的支持非常少。为此,即使 Flash Attention 也一直拒绝为 prefixLM 训练(即自定义掩码)提供支持,尽管出于某种原因对其 github 问题有合理的需求。
我知道我应该使用 Jax。一位朋友刚刚因为我使用 pytorch 而羞辱我,但这是一家初创公司,我们需要行动快速。我并不为此感到自豪。
少一点原则,多一点Yolo
系统地扩展模型通常需要有原则地从小到大,即分多个阶段(1B->8B->64B->300B等)进行实验,然后选出获胜者并不断扩展它们。在初创公司中,我们执行这些大规模扫描来检查 hparam 的计算量要少得多。最后,我们不得不进行多次Yolo run(幸运的是结果很好)。
YOLO,you only live once。即开始就全套运行,而不是一小步一小步地训练。
最终,我们只需要极少数的较小规模和较短的烧蚀运行,即可获得强大的 21B Reka Flash 和 7B 边缘模型(以及我们即将推出的最大核心模型)。在运行次数非常有限的情况下找到可靠的配方具有挑战性,并且考虑到搜索空间极其巨大,需要立即更改许多变量。
为了做到这一点,人们必须放弃大型科技公司的系统性,而在很大程度上依赖“Yolo”、直觉和本能。
我(以及我们团队中的许多人)在我们的 ML 职业生涯中已经建立了相当多的这种直觉,可以在很短的尝试时间内得到正确的结果。虽然我们在之前的工作中训练过非常好的模型,但训练基础设施、数据、新想法的融合和其他环境问题的差异仍然可能导致结果的巨大差异。
也就是说,强大的先验有助于显着减少搜索空间,这可能是解释为什么我们能够通过如此少的试验、资源和实验来训练真正强大的模型的最简单的解释之一。
写在最后:荒野中起舞,痛并快乐
在荒野中弄清楚事情是一次有趣的经历。不幸的是,这并不是无痛的。计算稀缺和不可靠的计算提供商使事情比预期困难得多,但我们很高兴我们凭借强大的技术实力渡过了难关。
总而言之,这只是我们如何在不到一年的时间里创办一家公司、筹集一些资金、购买一些芯片并让模型的性能匹配到 Gemini pro/GPT 3.5,并超越了许多其他公司的故事的一小部分,而一切都从头开始。
参考链接:
https://reka.ai/reka-flash-an-efficient-and-capable-multimodal-language-model/