从裸机到700亿参数大模型,这里有份教程,还有现成可用的脚本

选自imbue.com

作者:Imbue 团队

机器之心编译

编辑:panda

我们知道 LLM 是在大规模计算机集群上使用海量数据训练得到的,机器之心曾介绍过不少用于辅助和改进 LLM 训练流程的方法和技术。而今天,我们要分享的是一篇深入技术底层的文章,介绍如何将一堆连操作系统也没有的「裸机」变成用于训练 LLM 的计算机集群。

这篇文章来自于 AI 初创公司 Imbue,该公司致力于通过理解机器的思维方式来实现通用智能。

当然,将一堆连操作系统也没有的「裸机」变成用于训练 LLM 的计算机集群并不是一个轻松的过程,充满了探索和试错,但 Imbue 最终成功训练了一个 700 亿参数的 LLM,并在此过程中积累了许多有用的经验。

本文将深入介绍该团队构建自己的 LLM 训练基础设施的全过程,并会分享他们为方便监控、检查和纠错而编写的诸多工具和脚本。

如果你有心构建自己的 LLM 训练基础设施或好奇 LLM 是如何炼成的,那么这篇文章值得你阅读和收藏。

以下是 Imbue 团队文章原文。

引言

我们这个由研究者和工程师组成的小团队用了几个月时间在自己的基础设施上从头开始训练了一个 700 亿参数量的模型,并且该模型在推理相关的任务上胜过了零样本的 GPT-4o。

今天,我们要分享的是设置所需基础设施的过程:从组合初始集群和安装操作系统到设置在训练期间遭遇错误时自动恢复。我们会详细说明在每一步遭遇到的难题和解决方法。除了这些心得之外,我们还将发布我们一路上开发的许多脚本,以便其他团队能更轻松地为自己的模型训练创建稳定的基础设施。

在整个过程中,我们的工程师团队与 Voltage Park 一起准备好了计算机集群,构建了生产应用的基础。这整个过程包括:

1. 配置各台机器

2. 配置 InfiniBand

3. 确保机器完全健康

4. 诊断常见的训练问题

5. 改进基础设施工具

下面将详细描述每个步骤。

背景:这是如何炼成的

我们执行计算的目标是确保能快速实验大型语言模型。为此,我们需要大量高速 GPU,并且这些 GPU 之间也能高速通信。

本文将重点关注一个集群,其中包含分散在 511 台计算机中的 4088 台 H100 GPU,也就是每台计算机 8 台 GPU。之所以有 511 台带 GPU 的计算机,是因为需要保留一些连接给统一结构管理器(Unified Fabric Manager)节点,其作用是管理 InfiniBand 网络。在 511 台带 GPU 的主机上,每台 GPU 都直接与一块 ConnectX-7 网卡相连,其能以 400 Gbps 的速度与该 InfiniBand 网络中的任意 GPU 传输数据。

我们的 InfiniBand 网络拓扑结构的形式是「fully non-blocking」,即完全无阻塞;理论上讲,这能让 GPU 以最大速度互相通信。为此,我们使用了一种三层式 InfiniBand 网络架构:三层 InfiniBand 交换机。只要连接正确,便能让整个网络都获得这样高水平的吞吐量。下图展示了这个 InfiniBand 网络的概况:

请注意,训练网络时的通信发生在 InfiniBand 上,而不是以太网上。尽管这些机器也连接了以太网,但该网络的作用是传输数据集和检查点等数据。如果使用以太网来发送数据,速度会慢得多,因为数据会先从 GPU 传输到 CPU,然后再通过 100 Gbps 速度的以太网卡发出去。尽管也可以使用名为 RDMA over Converged Ethernet(RoCE)的技术基于以太网进行训练,但那需要在硬件和软件方面都做大量额外工作,并且可靠程度通常也不及 InfiniBand。详情可参阅这篇论文:https://arxiv.org/pdf/2402.15627

另外还有一个仅用于配置和管理的辅助以太网,从而可访问 BIOS(基本输入输出系统)、电源和其他低层机器接口的控制界面。如果没有这个管理网络,我们就必须通过 USB 驱动、键盘和显示器来手动设置每个节点。对于有几百台机器的情况,这并非一种可持续的方法。

要在这个集群上实现高性能训练,就需要每个组件(InfiniBand、以太网、GPU 和节点本身)都近乎完美地工作。如果这 12,000 多个连接中有任何一个有点不稳定,就会拖慢整体的训练运行速度。

本文接下来的内容就是介绍如何让这一切都完美且稳定地运行。

过程:如何将裸机变成完全可运行的集群

配置各台机器

在通过管理网络建立了与集群的初始以太网连接之后,就获得了基板管理控制器(BMC/baseboard management controller)的访问凭证。BMC 是一种远程监控主机系统的专用服务处理器,并且通常连接到一个分立的网络。它能让我们就像是亲身在现场一样操作机器,并还额外提供了硬件健康状况、BIOS 设置和电源管理的 API。

配备好这些组件后,我们就可以撸起袖子,开始设置集群了。

步骤 0:先配置好一台机器

我们首先使用 iDRAC(戴尔的基板管理控制器)在一台服务器上安装 Ubuntu 22.04,然后再基于这个操作系统设置其他一切。iDRAC 的一项能力是允许从本地计算机安装和启动 ISO 镜像,并通过浏览器提供一个虚拟控制台。理想情况下,这是该过程中唯一的手动安装步骤。

步骤 1:在每台机器上安装操作系统

在配置完第一台机器之后,继续安装 Ubuntu 的 Metal-as-a-Service (MAAS) 软件以帮助配置剩余的服务器。使用预启动执行环境协议(PXE)启动和自动化 iDRAC 工具,可指示每台机器从网络启动并配置 MAAS 以响应 PXE 启动请求。在执行初始的网络启动时,服务器会通过动态 IP 分配协议 DHCP 从 MAAS 获得一个 IP 和一个初始内核,而无需在本地驱动器上安装任何东西。这是用于自动重复执行操作系统安装的基本环境。从理论上讲,我们只需等待第一次启动,然后一切都会被处理好。但实际上,MAAS 与 BMC 的整合并不可靠,因此我们使用 iDRAC API 来事先收集每台机器的 MAC 地址(一种唯一的物理硬件标识符)。

在这整个训练过程中,MAAS 通常是椎栈中比较可靠的组件。但是,我们在开始时遇到了一些我们的设置特有的问题。举个例子,在配置前几台机器时,由于时钟相差太大,HTTPS 证书验证问题导致无法通过 apt 安装任何东西。与此相关的是,由于 MAAS 服务器必须负责很多事情(DHCP 服务器、用于将主机名解析成 IP 的 DNS 服务器、主机和官方 Ubuntu 软件包服务器之间的 HTTP 代理、NTP 服务器、cloud-init 配置管理、用于将 MAC 地址连接到 IP 到主机名再到自定义元数据的 ground truth 数据库),因此我们很难从根源上解决这些问题。此外,还有 MAAS 配置生命周期的学习曲线问题,因为是设计目标是处理管理绿地部署(greenfield deployment)的复杂性以及节点的逐步迁移和各种调试 / 不健康的中间状态。

步骤 2:诊断损坏的机器

我们发现大约 10% 的机器都无法启动,主要原因是服务器的物理问题。这是设置大型 GPU 集群的常见情况。我们遇到的情况包括:没接网线或接错了网线、iDRAC 中的硬件问题、电源单元损坏、NVME(快速非易失性内存)驱动损坏、内部线路缺失、网卡或 GPU 无法显示。我们自动检查了这些问题,将一些机器退回给了戴尔以重新测试,并为数据中心工作人员提交相应的工单。我们自己上手配置集群的一个优势是:在等待维护某些机器的同时就能立即使用健康的机器。

步骤 3:最小可行可观察机器

我们继续在每台服务器上进行如下设置:

1.Docker(以便更轻松地运行服务和训练作业)

2. 数据中心 GPU 驱动

3.Prometheus 节点导出工具(用于导出稳定的硬件 / 操作系统指标数据流)

4.DCGM 导出工具(用于从英伟达 GPU 导出额外的指标数据,如 GPU 状态、时钟、利用率)

5. 所有非操作系统驱动的 RAIDZ ZFS 池,这让机器在某个驱动失效时也能继续工作,同时还能免费提供透明的压缩(这对纯文本数据集和重复性日志尤其有用 —— 相比于不使用该工具,使用该工具通常能将可使用的空间增大 10 倍)

然后我们运行基本的 GPU 诊断以确定 GPU 是否大体功能正常 —— 不正常的通常会在几个小时内出现硬件问题。

在此期间,当我们试图同时在全部 400 个节点上安装软件包时,我们遭遇了带宽瓶颈。这是我们第一次在数据中心部署的多个组件上收到高温过热警报。这首批发热问题基本上都通过固件更新得到了解决。

步骤 4:单节点的 GPU 训练

下一步是确保每台机器都能够单独处理真实的 GPU 工作负载。很多机器都无法做到这一点,问题包括:

GPU 相关的错误,这类问题基本都可通过将 GPU 卡重新插入卡槽来解决:将 200 磅重的服务器从机架上滑出来,移除机盖和 GPU 之间的所有线缆,然后取出 GPU,再重新装上 GPU,之后再重新接上线缆并把服务器推回机架。

根据 Ubuntu 服务器日志,GPU 和 PCIe 总线或网卡之间的许多线缆都发出了这样的报错:「limited width: x4 < x16」。在更新 PCIe 交换机总线固件后,我们发现大约四分之一的主机需要重新安装内部 PCIe 线缆 —— 大概是因为外壳和 GPU 之间的线缆相当脆弱,这意味着每当对 GPU 进行维护时,这些线缆都会被推挤或拔掉。

还有一些杂项故障也影响了几台主机。戴尔通过固件升级帮助我们解决了一些问题:

NVMe 驱动器没有显示故障,但触摸时会锁定整台机器。

硬盘驱动器在 Linux 下以随机顺序显示,这给 MAAS 造成了混乱,并会导致操作系统被安装在错误的驱动器上。

温度读数错误,这会导致风扇一直全速运转。其原因可能是英伟达驱动有问题,这通过降级驱动版本而得到了解决。

CPU 的动态调频失控,将工作内核限制为 2 GHz。

直接的 GPU-GPU 通信(GDR 或 GPUDirect RDMA Peer Memory Client)无法成功应用。

配置 InfiniBand

步骤 0:安装 UFM

InfiniBand 的一个优势是其中心化的设计,这样一来整个网络就有了一个大脑。因此,对于整个网络结构中的 320 个网络交换机,我们只需处理其中一个实例。我们的首个任务是搞清楚哪台交换机连接了哪些机器,然后将其与接线图关联起来,并根据交换机的物理位置重新命名它们。

步骤 1:重新布线

一开始,UFM 无法检测到那 320 台交换机,更别说本应出现在网络结构中的主机了。在与我们的数据中心合作伙伴商讨之后,我们确认这些交换机已通电并接好了线,但依然无法检测到。在检查网络布线列表后,我们注意到该网络结构的顶层设计有误:这个结构不是统一的,而是分成了八个没有公共路由路径的互相脱离的网络。在重新接线之后,我们添加了检查步骤,以验证所有物理连接是否与新设计一致。

步骤 2:一万次温度告警(alert)

在解决了物理接线问题之后,InfiniBand 成功建立了与网络结构中所有 InfiniBand 交换机的联系。但是,几乎每个交换机端口都开始报告温度过高,有时候超过 70 ℃,即便它们还没有传输数据。我们发现这个问题源自同一机架中交换机之间的开放空间,这导致热空气回流到了前面。我们的数据中心合作伙伴帮助我们快速诊断出了该问题并制定了合适的解决方法。

步骤 3:1800 次告警

许多端口还有很高的错误率,或在正常和损坏状态之间来回变动,这被称为「抖动(flapping)」。这些问题只有在实际使用这些端口时才会出现,所以很难预先检测,因为我们的整个结构由 10,000 条高度冗余的链路组成。我们的数据中心合作伙伴帮助清洁和重新安装告警的端口,我们在等待替换时禁用了剩余的警报收发器。

尽管 InfiniBand 能弹性地应对硬件故障,但一旦大约 10% 的结构开始出现问题,自适应路由等功能就无法可靠地运行,无法解决偶尔丢失链路的问题。

在此期间,我们成功使用 100 到 200 台机器运行了多节点训练。我们的流程比较即兴:我们有时会随机启动一组节点,观察它们的性能,然后尽力让其中尽可能多的节点保持运行。该方法可让我们找到该 InfiniBand 网络结构中一组可靠的子集,但难度却很大,因为每次都需要改变用于训练的节点集合,由此改变默认的 InfiniBand 链路。

步骤 4:InfiniBand 疯狂燃烧

为了更高效地诊断 InfiniBand 问题,我们专门为整个集群设计了一个工作负载,其作用是同时将尽可能多的数据推送经过整个结构中的每个端口。这不同于在整个集群上运行一个大型的 all-reduce 工作负载,这需要使用 NCCL 来优化各个节点之中的通信,方法是使用 NVLink 经由 Server PCIe Module (SXM) 插槽来实现 GPU 通信。

相反,我们选择了一种蛮力方法,并轻松取得了成功。UFM 会在大多数端口的数据传输量超过理论容量的 97% 时开始发出警报,同时一些交换机会暂时宕机。我们认为能坚持到当天结束时的每个端口都是足够稳健的,其余的都被禁用或移除以待维修。

步骤 5:GPUDirect RDMA

要让 GPU 通信时不产生 CPU 计算开销,我们启用了一个名为 GPUDirect RDMA 的功能,其允许 InfiniBand 网卡之间直接通信。这涉及两个关键步骤:

1. 启动一个额外的核模块

2. 确保 PCIe Access Control Service (ACS) 被禁用,以防止 immediate hangs(立即挂起)

步骤 6:扩增「黄金」服务器

要使用最新硬件构建 GPU 集群,一个经验法则是:每周都有大约 3% 的机器出问题,要做好准备。

但是,需要说明一点:并不是每台机器都统一有 3% 的几率发生故障,而是少量不对付的机器反复出现各种不同问题,直到将它们妥善修复。这就凸显了在同一网络结构中配备大量机器的优势。因此,我们的做法不是随便找些机器来运行大规模训练,就像打地鼠一样看哪些出问题,而是专注于扩增已知可靠的服务器,也就是「黄金」服务器。

步骤 7:维护

InfiniBand 的维护主要涉及到响应 UFM 警报、更换故障线缆和收发器,以及偶尔诊断更困难的错误(比如交换机故障)。导致大规模维护的因素通常有两个:

1. 固件更新,尤其是集群中仅有一半完成更新时,这可能导致 UFM 状态损坏并必需重启所有 InfiniBand 交换机上的 UFM。

2.GPU 盒同时大规模重启,这可能会向 UFM 状态灌入大量更新,并同样需要重启 UFM 服务。

确保机器完全健康

在此过程中,我们发现了单台机器的多种故障或减慢训练速度的方式。其中许多故障模式并不会立即显现,因此我们编写了许多健康检查脚本,以检查主机是否足够健康。我们在这里发布了这些代码:https://github.com/imbue-ai/cluster-health

请注意,这些健康检查中的很多都特定于我们的运行时环境,并不一定与基础硬件相关,也不一定容易修复或自动化。这是设计决定的:为了实现让我们的机器准备好训练的总体目标,我们想要一个可以直接了当地回答「是」或「否」的单一入口点,并且可以概括总结任意数量的细微细节。

GPU 健康检查

我们检查了 GPU 数量是否正确、ECC(错误更正代码)检查是否已启用以及是否存在 ECC 错误。我们还检查了 NVLink 拓扑(将 GPU 彼此连接起来)是否已启动且无错误。

磁盘空间健康检查

我们检查了主机的磁盘空间利用率是否超过 95%。

Docker 健康检查

我们检查了 Docker 能否在连接了 GPU 的情况下运行容器(即 NVIDIA Container Runtime 是否正常工作),还检查了与监控 / 分析相关的 Docker 容器是否已激活并获得了正确的主机权限。

Dmesg 健康检查

我们检查了 dmesg 中是否存在硬件 Xids 或 SXid 错误(由 NVIDIA GPU 或 GPU 间 NVIDIA 交换机引发的故障)。我们还读取了所有 dmesg 日志行,以验证它们是否都可以归类到「常见 / 预期日志行」列表中。

iDRAC 健康检查

我们检查了机器上的 iDRAC 错误,其中忽略了非致命错误消息。这是戴尔计算机特有的检查,所以没有被包括在我们开源的代码中。

磁盘健康检查

我们检查了 zpool 是否已安装,Docker 是否已正确连接到它,以及它是否真的可以在不锁定 CPU 的情况下触及它。

InfiniBand 健康检查

我们检查了 InfiniBand 的错误率是否会增加和 / 或驱动固件是否过时。

Nvlink 健康检查

我们检查了机器上的 NVLink 错误。实践中看,这似乎不会导致训练失败,但可能会降低训练速度。

GDR 健康检查

我们检查了机器上的 GDR 是否已启用。

VBIOS 健康检查

我们检查了 GPU 的 VBIOS 版本以及 H100 基板固件是否是最新的。

Flint 健康检查

我们使用 flint 和 hca_self_test 检查了 Mellanox OFED 驱动、网卡固件和收发器固件的版本是否正确,以及它们是否针对英伟达驱动进行了正确编译。

PSB 健康检查

我们查询了 PCIe 设备,以检查 GPU、PSB(PCIe 交换机总线)和网卡之间的连接速度和宽度是否符合我们的预期。我们还检查了交换机固件是否为最新版本。该脚本由戴尔而非 Imbue 开发,所以我们目前无法共享它。

除了这些快速健康检查,我们还进行了一些更复杂的健康检查,包括:

通过 PyTorch 初始化矩阵计算,以及测量 NVLink 带宽和 GPU 计算速度和内存。我们设置了适当的 GDR 标志来测试 InfiniBand 和 NVLink。

使用 ib_write_bw 和 –use_cuda 通过 IB 卡发送数据并测量 PCIe 和 InfiniBand 卡带宽。这个过程持续了较长时间(约 15 分钟),以确保能找出抖动的 InfiniBand 链路。

运行多节点诊断运行以检查 NCCL 初始化能力以及它是否会随机停顿。如有停顿,则我们的分叉版 NCCL 代码会添加额外的日志记录。这需要 12 到 24 小时才能检测到问题,因此我们通常只对新节点或我们怀疑存在问题时运行此操作。

检查 DCGM 导出是否存在任何 GPU 时钟节流事件(不包括预期的 gpu_idle 和 power_cap)。为了检查这些电源事件,最好的方法是运行多节点训练,以同时检查所有 GPU、InfiniBand 卡以及 CPU 和磁盘。

诊断常见的训练问题

一旦硬件能正常工作,就可以开始训练了。

这一节将基于我们在我们的集群上运行大型语言模型训练的经验,分享一些具体的调试步骤和见解。

启动时崩溃

从某种程度上讲,这是所能遇到的最好的错误,因为其(理论上)很容易重现和迭代。

我们首先检查了我们的代码是否在正确的版本、配置和环境变量上运行。虽然很基础,但我们发现这一点至关重要:确保启动训练过程是可复现且容易检查的。一大原因是 Docker 镜像缓存或不透明秘密配置等中间抽象可能会造成混淆。

我们执行的另一个基础检查是确保所有机器都在线,并且可以轻松地聚合和检查所发出的栈跟踪记录或日志。我们使用了 Loki、Prometheus 和 Grafana 软件栈,但任何合适的日志聚合或跟踪 SaaS 的工具都可以。由于这些训练运行过程本质上是同步的和分布式的,因此第一个错误往往就会导致一连串的不相关错误。在这里,健康检查还可以帮助立马检测出硬盘驱动器损坏或 GPU 缺失或无效等错误。

我们构建了一个在发生故障时自动重启的系统,这使得日志和错误聚合变得更加重要,以避免混淆来自不同重启的错误。我们遇到的一些常见错误包括:

1.「Forward order differs across ranks: rank 0 is all-gathering 43 parameters while rank 1228 is all-gathering 1 parameters」这样的错误。我们发现这是 PyTorch 完全分片数据并行(FSDP)实现的一个奇怪特性,可通过重启解决。

2.GPU 内存不足(OOM)错误,看起来像这样:「CUDA out of memory. Tried to allocate …」通过多次检查我们的配置和代码并撤销近期的代码修改(由于启动期间 PyTorch 设备规格不正确而导致过多使用 GPU#0),我们解决了这些问题。

3.CPU/RAM 内存不足(OOM)错误,这些错误在错误日志中不太容易发现,并且通常能通过 Docker 容器外的主机的 dmesg 日志检测出来。当 OOM Killer 调用停止一个分叉进程或同级网络(network peer)时,我们可以看到它们主要表现为 CalledProcessError 或 ConnectionError。当从 dmesg 检测到了 OOM Killer 调用时,我们更倾向于直接放弃健康检查,并重启该机箱。我们还检查了我们的代码路径是否有足够的手动垃圾收集(下面有一部分介绍了如何禁用它),并还检查了是否有意外尝试进行计算或将张量移动到 CPU 上。

在训练过程中崩溃

首要任务是让系统能自动运行,让其能自动重新运行所有健康检查,然后在没发现不健康主机时重启运行。我们遇到了一些随机的硬件错误,包括 Xid 和 SXid 错误;这些错误可能会导致运行崩溃,却不会发出有意义的 Python 栈跟踪记录。行重映射等一些问题可通过重启恢复。不可纠正的 ECC 错误等另一些问题则往往需要硬件维护或更换零件。

此外,我们还观察到格式错误的训练数据也会导致崩溃。举个例子,如果语料库中有一个非常大的单个文档,就可能导致 GPU 或 CPU 出现内存不足错误。为了防止出现这个问题,我们采用了完全确定式的数据加载器 —— 通过与 epoch 或步数相关联,让每一次崩溃都可轻松复现。我们发现禁用数据加载或替换假数据(例如全零数据)有助于确认错误的根本原因是否是数据。

最后,通过指标聚合方法记录网络和一般节点的健康统计数据也很有帮助。以太网短暂断开或磁盘空间不足等问题可能不会显示为有用的错误消息,但却可以很轻松地与已收集的数据相关联。

没有栈跟踪信息的挂起(之后可能会有超时问题)

由于这些问题缺乏有帮助的信息,加上很难可靠地复现,因此这类错误的调试工作着实让人沮丧。

其中最令人难忘的错误类型伴随着这样的报错信息:

Watchdog caught collective operation timeout:WorkNCCL (SeqNum=408951, OpType=_ALLGATHER_BASE, … , Timeout (ms)=600000) ran for 600351 milliseconds before timing out

并且训练运行中的所有 GPU 工作器都发出了这样的报错信息。

这意味着一台或多台主机未能完成 NCCL 操作或者 NCCL 和 InfiniBand 连接崩溃,导致其他所有主机同时卡在了某个张量运算上,直到达到 NCCL_TIMEOUT 超时时间。不幸的是,受 NCCL 软件库本质所限,我们很难找到究竟是哪台主机出了问题。

我们对 NCCL 软件库的日志记录做了一些修改,参见我们的分叉版:https://github.com/boweiliu/nccl 。从而在崩溃发生时可能更好地揭示正在执行的消息或操作,从而确定阻止运行的可能是哪台主机或 GPU。

请注意,为了识别行为不当的主机,我们通常需要搞清楚哪些主机没有生成某些日志消息。缺少此类消息表明该主机上的工作器已落后或已崩溃。

其他没有可用错误消息的无响应情况通常与硬件相关问题有关,比如之前提到的 Xid/SXid/ECC 错误会导致英伟达驱动或英伟达 Docker 通信驱动锁定。为了区分 NCCL 挂起与驱动挂起以及 Python 代码中的竞争条件或死锁,我们使用 Py-Spy 和 GNU Project Debugger(GDB)等工具来实时调试遇到的停滞进程。使用此方法可发现一个特定问题:由于 Python 线程设置中的配置错误,我们无法在某些主机上正确启动八个多线程 NCCL GPU 进程,这些进程在 PyTorch 之前的初始化代码阶段遇到了竞争条件。

训练减速(由 MFU 测量)

缺乏工具让这类问题比前一类更让人沮丧。除了使用 Py-Spy、栈跟踪检查和 GDB 之外,我们还采用了 NVIDIA Nsight 和 profiling 工具,其中一些工具在高度分布式的设置中很难使用。

遗憾的是,导致一般减速或让速度低于之前演示的模型浮点数利用率(MFU)的原因有很多。

首先,事实证明多次检查配置、代码和环境变量是有用的。我们经历过的错误包括:运行了错误的模型、批量大小错误、UFM 或 NCCL 设置出错、CUDA_DEVICE_MAX_CONNECTIONS 出错。这都会导致性能无法达到最优。

我们还发现测量瞬时(即每批次)MFU(而不是平滑或窗口平均值)很有用,因为未平滑处理的 MFU 曲线通常有助于诊断问题类别。导致训练速度下降的问题包括:

从非常低的 MFU(低于预期的十分之一)立即开始训练并保持稳定

这多半是 InfiniBand 网络连接的硬件问题,例如 T2 或 T3 层的交换机死机。GPU 和 NIC 之间的硬件问题也可能导致该情况,对此 dmesg 会这样报错:PCIe x16 lanes limited by …

从预期 MFU 的 30% 立即开始训练并保持稳定

其原因可能是一台主机的 GDR 设置不正确(NVIDIA 对等内存)或 GDR 环境变量不正确。

从预期 MFU 的约 60-80% 立即开始训练并保持稳定

最常见的原因是 InfiniBand 链路质量不行或故障,尤其是单台 GPU 出现了与 InfiniBand NIC 相关的故障,导致 NCCL 尝试经由本地 NVLink 路由流量并在同一主机上的另一台 GPU 上使用 NIC。CPU 节流也可能导致该问题,这需要调整某些主机的 BIOS 设置。

在处理某些数据批次时突然大幅减速(下降 10 倍),并且经常发生这种情况

这基本上都与检查点或评估有关 —— 可通过检查 epoch 数或步数来验证。恼火的是,如果设置了在 MFU 异常时自动告警,就会出现许多误报。

在处理某些数据批次时突然大幅减速(下降 10 倍)

这种情况是随机发生的并且相当罕见(大概每 15 分钟一次),并且之后马上就会完全恢复到良好的 MFU。

最常见的原因似乎是其他需要大量 CPU 计算的工作负载被调度到了一台运行中的主机上。我们发现,与其构建分析工具来识别特定的主机,通过 PID 来粗略地监控 CPU 会更容易。其原因可能是偶尔出现的网络连接问题,比如数据加载器遭遇瓶颈。我们监控了数据加载、检查点和任何非 NCCL 代码的指标数据并添加了 Python 代码计时日志,事实证明这非常可靠。

MFU 在运行过程中逐渐减慢,但每次重启后又会回到 100%

理论上讲,其原因可能是交换机上的热量积累,但我们从未见过这种情况。不过,我们使用 Python 和 NVIDIA 分析器确定:性能下降的原因似乎是自动垃圾收集。

在调试解决这些减速问题时,我们发现吞吐量几乎必然会周期性地下降。随着训练地推进,这种下降会对分布式运算带来越来越多的影响。这让我们猜想下降的原因可能与自动垃圾收集有关 —— 我们通过分析和测试验证了这个猜想。当我们禁用了自动垃圾收集,并在所有主机上设定只在特定的间隔内收集垃圾,这种吞吐量「下降」就消失了。

我们使用了一种基于 ZeRO-3 的同步分布式训练算法 FSDP。在阻塞操作期间,运行垃圾收集的单个工作器进程可能会减慢其他所有工作器的速度。如果有数百个工作器进程,就可能导致速度大幅下降。

一开始性能良好,然后突然下降(达到预期的 70%),并且以高频持续(每 15 秒)

我们观察到这与英伟达 GPU 的「时钟节流原因」相关,这可通过对英伟达 DCGM 进行适当的设置来解决。发热问题(GPU 高温或主机冷却风扇故障 / 效果下降)或电源故障会导致该问题。另外,当我们同时最大化所有 8 台 GPU 利用率和 8x NIC InfiniBand 利用率以及 CPU/RAM/ 磁盘时,我们的一些具有特定电源硬件的主机会出现电压问题,但只有全部使用它们(通常只在实际训练运行期间)时才会出现这种情况。

性能优良但噪声比平常情况多(高频白噪声方差在预期 MFU 的 90% 和 100% 之间)

这也与 InfiniBand 硬件有关,但通常是由于网络中较高层的链路出现一定程度的性能下降或抖动,而不是冗余度较低的主机到 T2 层。

不幸的是,很多这类问题都难以定位到某台具体主机,而与 InfiniBand 相关的问题尤其难以定位,这是由于 InfiniBand 交换机技术的拓扑感知特性。InfiniBand 似乎更偏好 InfiniBand fat-tree 设计中的相邻主机,而 UFM 能以不对称的链路速度路由数据包。

以下是用于调试吞吐量问题的简单摘要 / 流程图 / 完备性检查表:

这套系统之前能正常工作吗?

你最近做了什么修改(比如合并代码、更新驱动)?

你运行的主机是否健康?你的所有依赖服务是否都运行正常,包括第三方的 SaaS,比如 Docker Hub、GitHub 等等?

你能确定现在运行的代码、环境、配置、版本、主机列表、排名顺序、随机种子与上次完全相同吗?(如果能实现这样的检查的话。)

问题可复现吗?

与其他事物有何关联?其他进程?每日 crontab 定时任务?主机或 DCGM 或 UFM 指标?

你的工具是否能正确度量这些指标?

在运行约简版的代码(使用更小的模型、假数据、不保存或加载检查点)时,问题是否依然存在?

改进基础设施工具

完成了以上步骤之后,就能在训练模型时实现优良性能了…… 至少在某个地方出故障之前是这样。

本节将介绍一些用于确保训练持续稳定的工具和系统,同时最好尽可能地少地需要人类干预。由于我们这个团队很小,因此我们没有足够的人手来进行人工维修,所以我们也希望能尽可能地自动化这个过程。

我们在训练过程中遇到的所有问题几乎都可归因于机器或网络组件故障。这类故障在大型集群中很常见,因此我们的做法是自动禁用出故障的机器和网络组件并发送维修请求。

机器故障

我们开发了一个系统,可在运行崩溃时自动从最近的检查点重启。在这个重启过程中,首先是对每台可用机器进行健康检查,然后基于其传递的健康检查结果对每台机器进行分类;然后尝试在最健康的机器上重启训练。

网络组件故障

我们观察到的所有网络故障都可通过 UFM 检测到,并会被记录到 UFM 事件日志中,因此响应方式也很简单:解析 UFM 日志并采取相应措施。

UFM 事件系统非常复杂,包含数十种事件类型。但在实践中,我们发现只有少数事件有问题,主要与链路故障或符号错误技术较高有关。在识别出这些事件后,我们可以编写脚本来解析 UFM 事件日志、禁用与最近事件相关的链路和端口、为这些网络组件申请维护工单、维护完成后重新启用这些组件。

本地镜像文件系统

对于这些大型分布式训练,人们很早就发现集群与以太网的数据交换速度是一大瓶颈。一条共享以太网连接的带宽大约为 10Gbit/s;如果有数百个工作器同时下载数据集和模型检查点,那么这点带宽会很快饱和。

为此,我们决定在我们集群内部构建一个本地文件系统,以作为云存储的镜像,其本质上就是一个缓存空间,可以减少从 S3 读取的文件量。为了解决集群流失问题(即因为维修原因而禁用或更换机器的情况),我们为每份文件都准备了三个副本,并使用了一致性哈希以均匀分配负载,从而在集群流失期间最大限度地减少文件移动。由于集群的磁盘空间有限,所以我们必须开发多种工具来跟踪文件的生命周期和清除不再有用的文件。

本地分布式 Docker 注册表

我们使用了 Kraken,这是一个可点对点传输 Docker 镜像的出色开源软件。这个软件几乎没出现过任何问题,我们还是挺惊讶的,毕竟我们的任务和实现都很复杂。工具地址:https://github.com/uber/kraken

各种性能监控工具

我们设置了默认的 Torch 分析器以及英伟达的 Nsight Systems。后者可帮助我们了解前向 / 反向通过以及 NCCL 通信所需的确切时间,并进一步帮助我们确定给定的模型大小和工作器数量是否会成为瓶颈。然而,Nsight Systems 有点难用,因为其需要在特权模式下运行 Docker,这需要禁用与性能监控事件相关的安全检查,并且保存其配置时往往需要停止整个训练进程。

此外,我们也编写了工具来检测训练速度慢的数据批次并理解其可能原因。我们发现这很有用。其中最有用的工具的作用是监控每一批次所需的时间并在某批次过于慢时丢弃该工作器的栈跟踪 —— 这让我们可以更轻松地找到硬件或软件有些小问题的主机。

将机器分成不同的组别以定位故障主机

在使用该集群的前几个月(那时候健康检查还不如现在这般透彻),我们经常遇到这种情况:在一组机器上训练时出现故障,但并不清楚究竟是哪台机器有问题。为了找到故障主机,我们开发了一些工具,可轻松地把一组机器分成不同的小组,然后在每个机器小组上运行更小的训练。

举个例子,如果一个在 48 台机器上运行的训练失败了,那么就在 6 组各 8 台机器上进行更小规模的训练,然后在 8 组各 6 台机器上运行更小规模的训练。通常情况下,这两个阶段中只有一次运行会失败,这让我们有信心得出结论:在两个阶段中都出问题的机器是有问题的。

反思和学习到的经验教训

在设置和维护基础设施的过程中,我们获得了一些有用的经验教训:

一种有用的做法是交换机器的位置。在运行时,使用多于所需机器数量 10-20% 的机器会很有帮助,这样就能在机器故障时轻松重启训练了。设置集群网络连接时让每台机器都与其他每台机器紧密相连,这样一来我们就能使用这些机器中任意可工作的子集。

对遇到的每一个硬件或软件故障,编写测试和自动化解决方案是值得的,因为训练中遇到的每一个问题都会再次出现。类似地,对于每一个含混不清的报错消息,都值得编写更能解释该错误的工具。

可重复性是优秀科研的关键。我们立马就采用的一大原则是:「一次只修改一个地方」,即便最简单的地方也是如此。

信任,但也要验证。每当我们引入外部工具或加入新人员(无论是来自公司内还是公司外)时,我们都会仔细检查他们声称的东西,尤其是当后续步骤依赖于这些声称的东西时。

总结

训练大型语言模型一开始就需要复杂的基础设施。我们之所以选择深入参与基础设施的设置细节,是因为我们相信完全理解我们操作的系统是非常重要的,也因为我们认为这样做的效率更高。

现在,经历过整个流程之后,我们很高兴我们采用了这样的方法 —— 事实证明,能完全控制我们的基础设施以及有能力轻松地在每个抽象层级上进行调试具有至关重要的价值。虽然这个过程需要大量的监督和迭代,但它让我们可以深入地理解其底层工作流程、构建一系列用于确保主机健康的工具、学习如何让系统自动运行以确保训练持续平滑,最终构建起一套让我们可以快速迭代训练前沿语言模型的基础设施。

这个基础设施构建流程体现了我们研究和构建 AI 智能体的基础方法论:探究细枝末节,不断改进现有流程,并构建有用的工具和系统使我们这个积极奋进的团队能够应对更大的挑战。