怎样在 10k 个 H100 GPU 上训练模型？

编者按： 怎样在 10,000 个 H100 GPU 上训练大模型？如何充分利用每一块 GPU 的算力？如何在这个复杂的 GPU 网络中高效传递数据？当不可避免的硬件故障发生时，又该如何快速恢复训练进度？我们今天为大家带来的文章中，作者为我们揭示了应对这些挑战的关键策略。

作者 | Soumith Chintala

编译 | 岳扬

我的好友 Francois Fleuret 提出了上述问题。我迅速总结了一些在大规模训练领域中相当普遍的知识，内容分为三部分。

1. 首先，是如何将尽可能大的神经网络和 batch-size 适配到那 10000 张 H100s 上，这个步骤涉及到并行处理和使用节省内存的各种技巧。
2. 其次，是如何在这些 GPU 之间尽可能高效地传递模型状态信息（state）。
3. 最后，是如何在遇到硬件或软件故障时，尽可能迅速地恢复系统。

01 如何将尽可能大的神经网络和 batch-size 适配到那 10000 张 H100s 上

1.1 并行策略

1. 在数据批次（batches）上进行并行处理（数据并行（data parallel））
2. 在神经网络层上进行并行处理（比如，将一层神经网络层分布到多个 GPU 上进行计算）
3. 对神经网络的不同模型层进行分割，以便它们能够在不同的 GPU 上运行（比如，前 N 层运行在 GPU1 上，第 N+1 层到第 N+10 层运行在 GPU2 上）

持续优化并行策略，直到所有 GPU 都能被高效利用，达到最高利用率。

1.2 Checkpointing / Compute vs memorize

• 在执行前向传播时，需要保存一些中间结果以便后续计算反向传播（save\_for\_backward）。然而，当神经网络规模非常大时，为了处理更大的数据批次，更有效的方法是释放这些中间结果，待到需要计算反向传播时再重新计算。
• 类似 FSDP 这样的技术，通过在单个 GPU 上只保留模型的分片来节省内存。当需要其他权重时，会从其他 GPU 聚合模型的完整权重。

02 尽可能高效地在 GPU 集群间传递模型状态信息

2.1 Communication overlap 策略：

在需要 GPU 间通信时，应尽可能早地启动通信过程：

• 例如，当第 N 层完成反向传播后，在第 N-1 层还在进行反向传播计算时，负责第 N 层的所有 GPU 可以同时开始执行梯度全归约操作。

2.2 探索并利用网络底层拓扑结构：

在多个计算节点间传递大量模型状态信息（如梯度、优化器状态信息）是一项复杂的任务。在使用 Sync SGD 时，需要尽可能快地集中传输这些状态信息。

网络中可能包含多层交换机，并具备 RDMA 能力（可以直接将 GPU 内存中的数据复制到网卡，完全绕过 CPU 内存），同时拥有前端和后端网卡（前端网卡连接到如 NFS 之类的存储系统，后端网卡则将 GPU 连接到集群中的其他 GPU）。

因此，在执行 all-reduce 或 scatter/gather 等通信操作时，充分利用这些网络信息至关重要。例如，通过树形归约算法（tree-reduce），all-reduce 操作的时间复杂度可以降低到O(log(n))；同时，网络光纤连接节点间的不同类型光纤对常数因子的影响，对于减少整体延迟时间也是非常重要的。

像 NCCL 这样的库能够智能地识别底层网络拓扑，并在执行 all-reduce 和其他通信操作时加以利用。

在这样的大规模计算中，我们还必须调整交换机和网卡中的数据包路由算法，以实现有效的负载均衡。交换机也需要大量的 HBM 内存（不仅仅是 GPU 需要），因为当数据包排队等待时，需要在某个地方排队而不会被丢弃——这就是交换机级别的 HBM 内存。

03 如何在遇到硬件或软件故障时，尽可能迅速地恢复系统？

故障是不可避免的，涉及GPU、网卡、电缆等多种硬件。有些故障能够迅速被发现，而有些则可能因为某个节点没有按时响应（比如 NCCL 的 all-reduce 操作卡住了）才被察觉。我们开发了多种工具来监控机群的健康状况，并尽可能快地将故障节点从机群中移除。这可不是一件容易的事。

在这种规模下，内存位随机翻转导致的隐性数据损坏概率增加，可能导致训练 loss 值异常升高。虽然这种问题在小规模系统中很少见，但在大规模系统中则可能频繁发生。在软件层面提前检测这种问题非常困难。一些硬件设备配备了内置校验和的电路，可以在计算后进行校验 —— 这样，一旦发生位翻转，硬件就能触发中断。但 H100 和之前的 NVIDIA GPU 都不具备这一功能。

为了应对这些故障，我们需要尽可能频繁且迅速地保存模型状态信息；一旦发生故障，我们也要能够迅速恢复并继续训练。通常，我们会迅速将模型状态信息另存到 CPU 内存的一个独立线程中，并在后台将数据从 CPU 内存写入到磁盘或远程存储系统。我们还以分片的形式保存模型状态信息（利用了 torch.distributed 的 checkpointing 功能），也就是说，不是每个 GPU 都需要保存完整的模型权重；每个 GPU 只需保存一部分权重 —— 其余部分可以通过其他 GPU 的分片 checkpoints 来恢复。

Thanks for reading!

Hope you have enjoyed and learned new things from this blog!

About the authors

Soumith Chintala

Cofounded and lead @PyTorch at Meta. Also dabble in robotics at NYU. AI is delicious when it is accessible and open-source.

END

本期互动内容 🍻

❓还记得你第一次配置分布式训练环境时的经历吗？有什么想对新手说的建议？

原文链接：

https://soumith.ch/blog/2024-10-02-training-10k-scale.md.html