在1万块H100 GPU上训练模型涉及到复杂的硬件架构和并行化策略

  每周跟踪AI热点新闻动向和震撼发展 想要探索生成式人工智能的前沿进展吗？订阅我们的简报，深入解析最新的技术突破、实际应用案例和未来的趋势。与全球数同行一同，从行业内部的深度分析和实用指南中受益。不要错过这个机会，成为AI领域的领跑者。点击订阅，与未来同行！ 订阅：https://rengongzhineng.io/

在使用1万块H100显卡时，目标是尽可能将网络规模和批处理大小最大化。为此，涉及多个并行化策略：

1. 批处理并行化（数据并行）：将不同批次的数据分配到多个GPU。
2. 层内并行化：将一个神经网络层拆分到多个GPU上。
3. 层间并行化：将不同的层分配到不同的GPU上，比如第1到N层在GPU1上，第N+1到N+10层在GPU2上。

通过不断优化并行化，确保所有GPU的高效利用，最大化计算资源。

Checkpointing与计算/记忆权衡

为了执行反向传播，需要在前向计算时保存一些中间变量（如save_for_backward）。但当网络足够大时，可以选择释放这些变量，腾出内存以容纳更大的批处理量，并在需要时重新计算。类似于FSDP的技巧，可以将权重分片保存在不同GPU上，节省内存，等需要计算时再从其他GPU获取这些权重。

GPU间的快速通信

在多个GPU之间传输数据时，尽量让通信与计算并行进行。例如，当第N层完成反向传播时，第N-1层的反向传播仍在计算中，这时可以让第N层的GPU们进行梯度的all-reduce操作，从而减少等待时间。

网络拓扑结构的发现与利用

在大规模同步SGD训练中，必须快速传输大量的梯度和优化器状态。这通常涉及多个网络节点，并需要考虑复杂的网络拓扑结构。通过树状的all-reduce算法，可以将通信复杂度降低至log(n)，同时尽量减少光纤连接节点间的通信延迟。

NCCL库可以智能发现网络拓扑并优化通信操作，比如all-reduce或scatter/gather。同时，网络设备（如交换机和NIC）的路由算法也需要优化，以实现负载均衡。此外，交换机的HBM内存在处理大规模通信时也非常关键，因为数据包排队时需要足够的缓存空间。

故障恢复

在拥有1万块GPU的环境中，硬件和软件故障几乎是常态。GPU、NIC、网络线缆等可能随时出问题。有些故障可以快速检测到，而有些只能通过延迟判断，比如NCCL的all-reduce操作卡住了。

另外，大规模训练时偶尔会出现内存位翻转导致的数据损坏，进而引发损失爆炸。这种现象在小规模训练中也有，但非常罕见，难以预先通过软件检测。虽然部分硬件具备内置校验功能，但像H100这样的GPU还没有这一特性。

为应对这些故障，建议频繁保存模型状态，并快速从故障中恢复。通常，模型状态会首先保存到CPU内存，然后再异步保存到磁盘或远程存储中。模型状态保存时可以分片处理，不需要每个GPU保存所有权重，而是保存一部分，其他部分可以从其他GPU的检查点中恢复。这样即使某个节点发生故障，恢复时间也会大大缩短。