编者按： 怎样在 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

询问AI