大规模训练 LLM 会带来独特的基础架构挑战，尤其是在作业跨越数千个 GPU 并长时间运行的情况下。这些作业运行的时间越长，遇到计划外中断或资源波动的可能性就越大。即使设备不经常可用，也会对紧密互联的集群产生巨大影响，导致给定训练运行速度减慢。
对于大规模训练作业而言，根据可用 GPU 的数量弹性调整作业是提高产量的有效方法。在 AI 训练背景下，Goodput 是衡量已完成的推动收的有用工作的关键指标，而不仅仅是原始硬件吞吐量。
如今，有效的弹性扩展方法包括丢弃数据副本、利用快速检查点重启或更换热备份。这些方法使 LLM 作业能够适应 GPU 可用性变化，同时保持整个系统的平衡。但是，在训练作业在可用性降低的情况下运行期间，它们也会造成一定程度的吞吐量损失和更高的成本。
最近一篇关于 Nonuniform Tensor Parallelism (NTP) 的论文介绍了一种前瞻性的实验性框架，该框架基于这些方法构建，可最大限度地减少吞吐量开销。结合潜在的动态功率提升来抵消任何性能损失，吞吐量保持稳定，将中断转化为可管理和可恢复的事件。
NTP 的核心贡献是能够通过防止瞬变设备问题阻碍大型互联训练任务，维持高 Goodput。通过动态调整张量并行度并智能地重叠必要的数据重新分片，NTP 可更大限度地减少损失的时间和计算工作量。
这可确保集群将更多的运行时间用于执行有用的工作，即使硬件条件发生变化，也能保持训练运行的整体完整性和效率。
大规模训练面临的挑战
AI 模型训练是一项并行工作，涉及数千个 GPU。并行化这些工作负载的一种常见技术是张量并行(TP) ，其中神经网络的层被分割到一组紧密合的 GPU 上。该组中的 GPU 数量与纵向扩展域不谋而合，纵向扩展域通过 NVIDIA NVLink 等高速互联技术互联。在 NVIDIA Blackwell 和 NVIDIA Blackwell Ultra 系统上，NVLink 以 1800 GB/s 的速度连接多达 72 个 GPU，支持单跳内的 all-to-all 通信。
图 1. 与 NVIDIA GB300 NVL72 相比，NVIDIA Blackwell Ultra HGX B300 NVL8 可纵向扩展域大小在典型场景中，前沿 LLM 在一组机架上进行训练，每个机架都容纳多个服务器。单个服务器机架构成了纵向扩展域，并构成了单个 TP 组。然后，数据并行 (DP) 会在多个此类扩展域中复制模型，每个域处理一批不同的数据。纵向扩展领域中 GPU 状态的变化可能会影响 TP 组的效率。由于同一 TP 组中的 GPU 共享紧密合的计算，因此一台设备出现问题可能会降低训练性能或需要暂时重新平衡以保持进度。
随着数据中心架构不断发展以支持更大的纵向扩展领域 (从 8 个 GPU 扩展到 72 个 GPU 及以上) ，最大限度地提高每个正常运行设备的正常运行时间成为实现高产出的关键。
该论文指出，系统的设计可保持绝大多数处于活动状态的 GPU 持续处理，而不是允许局部的瞬时中断来决定整体训练吞吐量。通过适应这些硬件波动，集群可以在大规模模型训练期间保持优化 Goodput 所需的高效资源使用。
训练通常涉及管道，其中每个阶段都依赖于先前步骤的及时完成。如果一个 GPU 在 TP 组内遇到延迟，它可能会减慢该组内的同步或处理速度，导致系统恢复或重新分配工作负载之前暂时停止或降低吞吐量。
NTP 如何维护训练
NTP 的基本理念是确保 DP 副本即使在瞬态硬件中断期间也保持高效。从最新的检查点恢复后，模型会自动根据可用硬件调整其配置，同时保持部分功能以保持工作流的移动，而不会完全丢失副本的贡献。
下文将深入探讨 NTP 如何实现这种弹性。
动态 TP 度适应性
当纵向扩展域中的 GPU 发生中断时，系统会识别受影响的组，并自动重新配置其张量并行性，使其仅利用剩余的功能 GPU。例如，如果一个包含 8 个 GPU 的 TP 组出现 1 次丢弃，它可以动态切换到 7 度的 TP。该模型分片会继续计算，从而防止其贡献完全丢失。该组中的其余 GPU 会增加其各自的工作负载，从而使训练工作能够在资源子集遇到问题时保持较高的良品率和可用性。
提升功率以实现性能补偿
降低 TP 度不足以维持全球吞吐量。GPU 较少的 DP 副本本身运行速度较慢，导致整个 DP 系统停止运行，等待最慢的副本。为了解决这一问题，该研究提出了一种包含改进的电气和散热能力的机架设计。这种设计能够在可用性降低的情况下提升纵向扩展领域的功率。通过动态增加受影响域中提供给活动 GPU 的功率，可以暂时提高时钟频率和计算吞吐量，提升功率。
TP 度降低的 DP 副本可以有效地跟上其他功能齐全的副本的步伐，防止全局同步瓶颈，并确保集群的 Goodput 在局部硬件波动的情况下保持高度优化。
高效重新分片
动态调整 TP 度需要一种有效的机制，以便在剩余 GPU 之间重新分配模型的张量分片。NTP 采用了一种巧妙的重新分片技术，与其他计算阶段重叠。
通过在向后计算和参数同步阶段执行这种重新分片，可将正常副本的用度降至最低，通常降至 1% 以下。这种精心的调度可更大限度地提高计算效率，无缝地维持最佳的 Goodput，而不会使适应机制本身成为性能瓶颈。
图 2. NTP 支持梯度重新分片，以便与反向计算重叠并同步，从而经济高效地隐藏NTP 为扩展 AI 训练构建了更灵活的路径
这项工作强调了共同设计硬件和软件对于解决大规模 AI 训练中固有挑战的重要性。NTP 与先进的机架设计集成，可提供动态功率提升所需的电气和散热空间，充分展示了经过深思熟虑的硬件创新如何与复杂的软件解决方案相辅相成。
硬件和软件之间的这种共生关系对于在新一代 AI 系统中实现更高水平的性能、效率、稳定的 Goodput 和弹性至关重要。作为一项前瞻性的实验性功能，NTP 展示了将弹性直接烘焙到并行策略中的可能性。
在此基础上，我们已经在进行研究，将这些概念扩展到非统一专家并行 (NEP) ，优化标准张量并行不太理想的多专家模型 (MoE) 的恢复能力。
查看 NVIDIA Resiliency Extension (NVRx) 或 Nonuniform Tensor Parallelism 自述文件 中提供的生产就绪型容错和弹性功能，详细了解其最近加入 NVIDIA Megatron Core 开发者分支的情况。
