AI Infra 综述（二）

AI Infra 综述（二）

5 计算优化 COMPUTATION OPTIMIZATIONS

核心思想： 现代AI芯片（如GPU）的算力非常强大，但要充分利用这些算力，就需要精巧的优化技术。本节主要介绍两大类方法：

  1. 算子优化 (Operator Optimization)
  • 核心算子优化： 针对模型里最耗费计算的核心模块（如注意力机制）进行专门的手动优化，以大幅提升性能。
  • 编译器自动优化： 利用编译器自动分析计算过程，通过利用硬件的大规模并行计算和高效内存访问能力，自动生成更快的代码。
  1. 混合精度训练 (Mixed-Precision Training)
  • 核心原理： 使用较低精度的数据（比如从32位浮点数降到16位）进行计算。这就像用“简笔画”代替“高清照片”，虽然精度略有损失，但计算速度会快很多。
  • 现状与趋势： 目前，16位混合精度训练已成为大模型训练的事实标准。为了追求更高的效率，学术界和工业界还在探索更低的精度，比如8位甚至1位。

5.1 算子优化

算子优化分为手动优化和自动优化两种。手动优化主要集中在最耗费资源的注意力算子上，而自动优化则应用更广泛。

5.1.1 手动优化的注意力算子

问题： 注意力机制（Attention）是Transformer模型的核心，但它的计算量和内存占用会随序列长度的增长呈平方级增加，这严重限制了模型的训练性能和能处理的文本长度。

解决方案： 手动编写高度优化的注意力算子。其中最著名的就是 FlashAttention 系列：

  • 核心思想： 大幅减少对慢速显存（HBM）的读写次数。
  • 实现方式： 它通过一种叫“Tiling”（分块）的技术，将计算拆分成小块，并在GPU高速的片上缓存（SRAM）中一次性完成“矩阵乘法 -> Softmax -> 矩阵乘法”等一系列操作。这避免了中间结果在慢速显存中的反复存取，从而大幅提升速度和内存效率。
  • 其他优化： 还有针对特定硬件或可变长度序列的专门优化，以避免因填充（Padding）而产生的计算浪费。

5.1.2 编译器的自动优化

通过深度学习编译器自动完成优化，主要分为两个层面：

  1. 算子核生成 (Kernel Generation): 自动为单个计算任务（如矩阵乘法）生成高效的底层代码。编译器（如 TVM, Triton）能分析硬件特性，自动地、更好地利用并行计算和数据局部性，生成媲美甚至超越专家手写性能的代码。
  2. 计算图优化 (Graph-level Optimization): 核心技术是算子融合 (Operator Fusion)。由于GPU的计算速度远快于访问显存的速度，频繁读写显存是主要瓶颈。编译器（如 PyTorch 2.0 的 TorchDynamo）会分析整个计算流程，将多个连续的小算子合并成一个大算子。这样，中间结果可以直接留在高速缓存中参与下一步计算，极大减少了对慢速显存的访问，提升了整体性能。

5.2 混合精度训练 (Mixed-precision Training)

核心思想： 用更低精度的数据进行计算，以节省计算、存储和通信成本。

一个简单的比喻：

  • FP32 (32位全精度) 就像一张超高清原图，细节丰富但文件巨大。
  • FP16 (16位半精度) 就像一张高清压缩图，肉眼看起来差不多，但文件小了一半。

在训练模型时，大部分计算并不需要那么高的精度。“混合精度训练”就是，在大部分计算中用“压缩图”（低精度）来加速，但在一些关键步骤（比如最后汇总结果时）换回“原图”（高精度），以保证最终结果的准确性。

为了防止在低精度下，一些非常小的有效数字被直接当成0（信息丢失），通常会使用一种叫 “损失缩放” (Loss Scaling) 的技术。它会先把所有数字都乘以一个很大的数，计算完后再除回去，从而保住这些微小的细节。

5.2.1 16位浮点数 (16-Bit Floating Point)

这是目前大模型训练的主流和事实标准。主要有两种格式：

  • FP16 (半精度浮点数):
  • 优点: 精度相对较高。
  • 缺点: 能表示的数值范围小，在处理非常大或非常小的数时容易“溢出”（overflow/underflow），导致训练不稳定。
  • BF16 (bfloat16):
  • 优点: 能表示的数值范围和FP32一样大，训练过程更稳定，不易溢出。
  • 缺点: 精度比FP16略低，且需要较新的硬件支持（如NVIDIA Ampere架构及之后的GPU，或Google的TPU）。

5.2.2 8位及以下的浮点数 (Sub-8-Bit Floating Point)

这是更前沿的探索，旨在将精度进一步降低到 FP8，以获得极致的性能。

  • 驱动力: 新一代AI芯片（如NVIDIA H100）提供了对FP8的硬件支持。
  • 挑战: 数值范围更小，溢出问题更严重。
  • 解决方案: 设计新的FP8数据格式，并结合自动缩放等技术来动态调整数值范围，以确保训练的稳定性。

5.2.3 低位定点数 (Low-Bit Fixed Point)

这是最极致的方案，不仅降低了位数，还从浮点数变成了定点数（整数），计算和存储效率更高。

  • INT8 / INT4: 使用8位或4位整数进行训练。
  • 1-Bit (例如 BitNet):
  • 做法: 模型权重只用 1-bit 来表示，即只有 +1 和 -1 两种值（或-1, 0, +1）。
  • 如何工作: 这并不意味着整个训练都在1位下进行。实际上，在训练时，系统内部仍然会保留一个高精度的“影子权重”用于梯度更新，保证了训练的稳定性和准确性。这可以看作是混合精度思想的极致应用。

6 内存优化 MEMORY OPTIMIZATIONS

大语言模型（LLM）训练时的内存开销主要源于四个方面：

  1. 模型状态：包括参数、梯度和优化器状态。在主流的混合精度训练中，这部分大约占用 模型参数量的16倍 内存空间。
  2. 激活值：前向计算时产生的中间结果，用于反向传播计算梯度。序列越长，这部分内存越大。
  3. 临时缓冲区：用于存储计算或通信过程中的临时数据。
  4. 内存碎片：频繁的内存申请和释放导致可用内存不连续，无法分配出所需的大块内存。

为解决这些内存瓶颈，主要有四类优化方法：

  1. 激活重计算：以算力换空间。不保存所有激活值，在需要时重新计算它们。
  2. 冗余减少：通过数据分片等技术，减少在多张GPU上重复存储的数据。
  3. 碎片整理：优化内存管理机制，提高内存利用率。
  4. 交换与卸载：将GPU中暂时不用的数据转移到CPU内存或SSD硬盘上，需要时再加载回来。

6.1 激活重计算

在模型训练中，反向传播（backward pass）计算梯度时，需要依赖前向传播（forward pass）产生的激活值（Activations）。对于大语言模型（LLM），存储所有激活值会消耗巨大的GPU显存，成为训练的瓶颈。

激活重计算（也常被称为梯度检查点 Gradient Checkpointing）是一种“以计算换空间”的核心技术。其基本原理是：

  • 前向传播时：不再保存所有的激活值，只选择性地保存其中几个关键节点（称为“检查点”），然后将中间的激活值丢弃以释放显存。
  • 反向传播时：当需要用到某个被丢弃的激活值时，系统会从最近的一个“检查点”开始，重新执行一小段前向计算，以实时生成所需的激活值。

核心权衡：这种方法显著降低了峰值显存占用，但代价是增加了额外的计算量。优化的关键在于找到内存节省和计算开销之间的最佳平衡点。

激活重计算主要分为两大类：

  • 静态策略 (Static Evicting)：预先制定好固定的丢弃和重计算计划。
  • 动态策略 (Dynamic Evicting)：在训练过程中根据实时状态动态决定丢弃哪些激活值。

由于LLM的架构（如Transformer层）高度重复且结构固定，静态策略非常有效，是目前的主流方法。

6.1.1 静态策略 (Static Evicting) - 预定义计划

静态策略根据模型结构，提前设计好检查点的位置。

  • 基本方法 (Selective-checkpointing)：这是一种常见的实践，通常在计算开销和内存占用较大的模块边界设置检查点，如注意力模块（Attention Module）。系统只保存这些模块的输入或输出，模块内部的中间激活值则在反向传播时重算。
  • 内核级优化 (FlashAttention, DistFlashAttn)：这些技术将重计算策略应用得更加精细。例如，它们不仅仅在整个注意力模块外设置检查点，而是优化模块内部的计算流程，只保存最关键的部分（如Attention的最终输出），从而避免了在反向传播时对整个模块进行重计算，极大地降低了长序列训练的计算开销。
  • 白名单机制 (LoongTrain)：通过将注意力模块等关键部分加入“白名单”，在前向传播时强制保存其输出，反向传播时直接复用，避免了对这些计算密集型模块的重计算。
  • 全局最优搜索 (Yuan et al.)：这是一种更系统化的方法，它会分析模型中每个张量的重计算成本和内存成本，通过算法（如寻找帕累托前沿）来自动寻找一个理论上在计算与内存之间达到最佳平衡的全局检查点方案。

6.1.2 动态策略 (Dynamic Evicting) - 实时决策

动态策略在训练时实时监控内存使用情况，并动态决策。

  • 原理：当检测到内存压力较大时，系统会根据某种启发式规则（如张量的访问模式、大小等）选择性地丢弃一些当前未使用的激活值。
  • 优点：灵活性高，不依赖于特定的模型结构。
  • 现状：尽管灵活，但由于实时决策会引入额外的管理开销，并且静态策略对结构规整的LLM效果已经很好，因此动态策略在LLM训练中应用较少。相关研究（如Coop）还会考虑在丢弃时避免产生内存碎片。

6.2 冗余减少

在标准的数据并行 (Data Parallelism) 训练中，为了让每张GPU都能独立计算，系统会在每张GPU上都存储一份完整的模型状态。模型状态主要包括三个部分：

  1. 模型参数 (Parameters)：模型的权重。
  2. 梯度 (Gradients)：用于更新参数的计算结果。
  3. 优化器状态 (Optimizer States)：例如Adam优化器中的动量（momentum）和方差（variance），内存占用通常是模型参数的数倍。

这种完全复制的策略导致了极大的内存冗余，限制了可训练模型的规模。冗余减少 (Redundancy Reduction) 技术的核心目标就是通过分片（Sharding）来消除这些冗余副本。

6.2.1 完全分片 (Fully Sharding)：ZeRO 技术

ZeRO (Zero Redundancy Optimizer) 是一种通过逐步分片模型状态来消除冗余的系统性方法，分为三个递进的阶段：

  • ZeRO-1: 分片优化器状态
  • 将占用内存最大的优化器状态进行分片，均匀地存储在所有GPU上。
  • 每张GPU只持有一部分优化器状态，但仍然保留完整的模型参数和梯度。
  • 效果：显著降低内存占用，是性价比最高的优化阶段。
  • ZeRO-2: 分片优化器状态和梯度
  • 在ZeRO-1的基础上，进一步对梯度进行分片。
  • 每张GPU在反向传播后，只保留与其负责的那部分参数相对应的梯度。
  • 效果：内存占用进一步降低，但通信开销略有增加。
  • ZeRO-3: 分片优化器状态、梯度和模型参数
  • 最彻底的分片阶段，将模型参数本身也进行分片。
  • 每张GPU在任何时候都只持有一小部分模型参数。
  • 工作机制：在前向或反向计算需要特定参数时，通过AllGather通信操作从其他GPU动态获取，计算完成后立即释放这部分内存。
  • 效果：内存效率最高，使训练超大规模模型成为可能。
  • 代价：引入了大量的即时通信，对网络带宽和延迟要求极高。

6.2.2 部分分片 (Partially Sharding)：通信与内存的权衡

ZeRO-3的全局通信在大规模集群中可能成为性能瓶颈。部分分片是对其的一种优化，旨在平衡内存效率和通信成本。

  • 核心思想：将分片操作限制在一个通信效率更高的范围内（如单台服务器内的GPU），而不是在整个集群的所有GPU上进行全局分片。
  • 工作机制：
  1. 将所有GPU划分为若干个子组 (subgroup)。
  2. 在子组内部，执行类似ZeRO-3的完全分片，所有模型状态（参数、梯度、优化器状态）都被分片存储。组内通信通常通过高速互联（如NVLink）完成，开销很小。
  3. 在子组之间，每个子组可以看作一个独立的单元，共同构成一个数据并行的整体。
  • 权衡：
  • 通信优化：大部分通信被限制在高速的子组内部，跨节点的全局通信频率和数据量显著减少。
  • 内存代价：由于每个子组都需要能够独立工作，因此在子组的粒度上存在状态冗余，内存节省效果不如全局的ZeRO-3。

PyTorch FSDP (Fully Sharded Data Parallel) 就是这种部分分片思想的一个典型实现，它允许用户灵活地定义分片组，以适应不同的硬件拓扑，在通信和内存之间找到最佳平衡。

6.3 碎片整理

在LLM训练过程中，GPU显存会被频繁地申请（分配）和释放。内存碎片（MemoryFragmentation）指的是，在多次分配和释放后，总的可用显存虽然很多，但它们被分割成了许多不连续的小块。

当系统需要申请一块较大的连续内存时（例如，存储一个大的激活值张量），尽管总的空闲内存足够，但由于找不到任何一块足够大的连续空闲空间，导致分配失败，从而引发内存不足 (Out-of-Memory, OOM) 错误。

产生原因：

  • 张量生命周期不同：不同的张量（tensors）在内存中存在的时间长短不一。
  • 频繁的内存操作：像激活重计算、数据卸载等内存优化技术，会引入更频繁和不规则的内存分配与释放请求，加剧了碎片问题。

碎片整理 (Defragmentation) 的目标就是通过优化内存管理来解决这个问题。

6.3.1 基于张量的碎片整理 (Tensor-based Defragmentation)

这类方法从深度学习框架的内存分配器 (caching allocator) 入手，通过优化张量的分配和释放策略来减少碎片。

  • 核心思想：通过分析整个计算图，预测所有张量的生命周期（何时创建、何时销毁）和大小，然后提前规划一个最优的内存布局方案。
  • 具体方法：
  • 优化计算顺序和内存分配 (ROAM)：通过算法调整算子（operators）的执行顺序，并优化张量的分配位置，使得生命周期相似的张量尽可能地分配在一起，从而最大限度地复用内存空间。
  • 2D装箱问题建模 (Imanishi et al.)：将内存分配问题看作一个二维的“装箱”问题（时间轴和内存地址轴），通过启发式算法找到一个能让“箱子”（内存块）排列得最紧凑的方案。
  • 与重计算结合 (MegTaiChi, Coop)：在执行激活重计算时，不仅考虑节省内存，还考虑丢弃哪些张量能更好地维持内存的连续性。

局限性：这些方法通常需要对计算图进行复杂的离线分析，对于LLM这种极其庞大的模型，分析和搜索最优解的成本非常高，扩展性可能受限。

6.3.2 基于虚拟内存管理的碎片整理 (VMM-based Defragmentation)

这类方法不去改变上层的张量分配逻辑，而是利用底层CUDA驱动提供的虚拟内存管理 (Virtual Memory Management, VMM) 功能来解决碎片问题。

  • 核心思想：将物理上不连续的多个小内存块，通过虚拟地址映射，“拼接”成一个逻辑上连续的大内存块。
  • 具体方法：
  • 虚拟内存拼接 (GMLake)：当需要一块大内存时，系统可以找到多个分散的小空闲块，然后在虚拟地址空间中将它们映射成一段连续的地址。上层应用（如PyTorch）看到的是一个完整的大内存块，而无需关心它在物理上是分散的。这种方法几乎没有数据拷贝开销。
  • 可扩展内存段 (PyTorch expandable segments)：允许一个已经分配的内存块在使用过程中“原地”扩大。如果旁边的内存块恰好被释放了，系统就可以直接将当前内存块的边界扩展过去，合并两个空间，而不需要重新分配和拷贝数据。

优势：

  • 透明性：对上层应用完全透明，无需修改模型代码。
  • 通用性：可以和各种内存优化技术（如重计算、卸载）无缝集成。
  • 高效性：利用了硬件和驱动层面的功能，开销小，扩展性好。目前已被主流框架（如PyTorch 2.1）集成。

6.4 卸载

当GPU显存不足以容纳整个模型及其训练状态时，卸载 (Offloading) 技术提供了一种解决方案。其核心思想是：

将GPU中当前不需要的数据（如模型参数、优化器状态、激活值等）临时转移到容量更大但速度更慢的外部存储资源上，如CPU内存或NVMe SSD。当GPU需要这些数据时，再将其加载回来。

这种技术利用了外部资源巨大的存储容量，使得在有限的GPU上训练超大规模模型成为可能。关键挑战在于隐藏或最小化数据传输带来的延迟，避免GPU因等待数据而空闲。

6.4.1 卸载到 CPU (CPU Offloading)

这是最常见的卸载方式，利用比GPU显存大得多的CPU内存。

静态卸载 (Static Offloading)

在训练开始前就制定好固定的卸载计划，决定哪些数据常驻GPU，哪些常驻CPU。

  • 代表技术：ZeRO-Offload
  • 策略：将最消耗内存的优化器状态和梯度放在CPU内存中，而将计算密集型的模型参数保留在GPU上。
  • 优势：简单有效，能显著扩展可训练模型的规模。
  • 局限：卸载策略固定，可能无法充分利用所有GPU显存，且CPU计算（如参数更新）速度慢，可能成为瓶颈。
  • 优化方向：通过更精细的分析（如Elixir）来决定哪些数据块卸载到CPU，以更高效地利用GPU显存；或者通过流水线（Pipelining）和预取（Prefetching）来重叠计算和数据传输，隐藏延迟。

动态卸载 (Dynamic Offloading)

在训练过程中根据实时内存使用情况，动态地决定何时、卸载哪些数据。

  • 核心思想：更加灵活，按需分配。当GPU显存紧张时，自动将最近最少使用的数据转移到CPU；当需要时，再加载回来。
  • 实现方式：
  • 细粒度管理 (PatrickStar, TSPLIT)：将大张量（Tensor）切分为更小的内存块 (chunks) 或微张量 (micro-tensors) 进行管理。这样可以更灵活、更精确地进行数据移动，只传输必要的部分，从而优化数据传输带宽和延迟。
  • 与流水线结合 (MPipeMoE)：在流水线并行中，动态决定是重新计算一个激活值划算，还是从CPU加载一个卸载的激活值划算。

6.4.2 卸载到 SSD (SSD Offloading)

对于万亿（Trillion）参数级别的巨型模型，即使是CPU内存也可能不够用。这时，就需要将数据进一步卸载到容量更大的NVMe SSD上。

  • 代表技术：ZeRO-Infinity
  • 策略：构建了一个GPU -> CPU -> SSD的三级存储体系。
  • 模型状态（参数、梯度、优化器状态）可以被卸载到CPU内存或NVMe SSD。
  • 激活值因为访问频繁，通常只卸载到CPU内存。
  • 优势：极大地扩展了模型规模的上限，理论上可以训练数十万亿参数的模型。
  • 挑战：SSD的读写速度远慢于内存，数据传输延迟更高，对系统调度和预取技术的要求也更高。
  • 优化方向：
  • 卸载激活值到SSD (Fuyou)：在CPU内存也极其有限的单卡或消费级服务器上，将激活值也卸载到SSD。
  • 近存储计算 (Smart-Infinity)：利用SSD附近的计算单元（如智能网卡DPU）直接在存储端执行一部分计算（如参数更新），减少需要传回GPU的数据量。
  • 混合存储与预取 (MoESys)：针对MoE（混合专家）模型，将不同类型的参数（稀疏/密集）存放在不同层级的存储设备上，并设计高效的预取调度策略来重叠计算和数据加载。

7 通信优化 COMMUNICATION OPTIMIZATIONS

在分布式LLM训练中，不同的并行策略会产生不同模式的网络通信流量：

  • 张量并行 (Tensor Parallelism)：需要在层内计算中进行 AllReduce 通信，对带宽要求极高。
  • 数据并行 (Data Parallelism)：需要在每次迭代结束时同步梯度，同样需要 AllReduce 通信。
  • 流水线并行 (Pipeline Parallelism)：需要在不同阶段（stage）之间传递激活值，通常是点对点（Send/Recv）通信。

通常的部署策略是：

  • 高带宽需求的并行（如张量并行）被限制在高带宽域内，例如同一台服务器内的多张GPU（通过NVLink高速互联）。
  • 较低带宽需求的并行（如数据并行、流水线并行）则可以跨高带宽域进行，例如跨服务器（通过InfiniBand或以太网）。

这种部署策略导致LLM训练的通信流量呈现出清晰的模式和层次性：绝大部分通信发生在小范围的GPU组内，只有一小部分流量需要跨越整个集群。这一洞察启发了新的网络拓扑设计，如轨道优化拓扑 (rail-optimized topology)，通过减少不必要的昂贵核心交换机来降低成本。

本章将介绍用于优化分布式LLM训练中集体通信 (collective communication) 性能的系统和技术。主要分为三个方面：

  1. 集体通信库 (Collective Communication Libraries)：
  • 讨论实现 AllReduce、AllGather 等核心通信操作的底层库。
  • 这些库使用的算法可以是预定义的 (pre-defined)，如经典的 Ring 和 Tree 算法，也可以是根据硬件拓扑自动合成的 (synthesized) 更优算法。
  1. 通信调度 (Communication Scheduling)：
  • 探讨如何重新组织通信操作的执行顺序，使其与计算操作重叠 (overlap)。
  • 目标是让GPU在等待网络通信时也能进行计算，从而隐藏通信延迟，提升训练效率。
  1. 网络内聚合 (In-Network Aggregation, INA)：
  • 介绍一种前沿技术，它利用网络设备（如交换机）自身的计算能力来执行一部分聚合操作（如梯度求和）。
  • 通过在数据传输路径上完成计算，可以显著减少需要传输到GPU的数据量，从而降低网络负载和延迟。

注意：本章不讨论通过有损压缩 (lossy compression) 或量化 (quantization) 等方法来减少通信量的技术，因为这些属于数据压缩范畴，而非通信系统优化。

7.1 集体通信

在分布式训练中，多张GPU需要协同工作，例如同步梯度、分发参数等。集体通信 (Collective Communication) 就是指一组GPU共同参与完成的通信操作。

  • 关键操作：包括 AllReduce (所有GPU的数据聚合计算后分发给所有人)、AllGather (所有GPU的数据拼接后分发给所有人)、ReduceScatter (聚合计算后将结果分块发给不同人) 等。
  • 常用库：虽然有通用的MPI库，但在LLM训练中，大家更偏爱针对AI加速器（如NVIDIA GPU）高度优化的专用库，其中最著名的是 NVIDIA的NCCL。

7.1.1 预定义算法 (Pre-Defined Algorithm)

NCCL等库内置了一套经过优化的标准通信算法，并会根据网络拓扑和数据大小等条件自动选择最合适的一种。

  • 环形算法 (Ring Algorithm)
  • 工作方式：将数据切分成小块 (chunks)，然后在GPU之间形成一个环路，像流水线一样依次传递和计算。
  • 优点：可以有效利用网络带宽，避免因等待大数据块而造成的网络空闲。非常适合处理大批量数据。
  • 缺点：随着GPU数量增多，环路变长，通信延迟会增加。
  • 树形算法 (Tree Algorithm)
  • 工作方式：将GPU组织成一个或多个树状结构进行数据聚合和分发。
  • 优点：对于大量GPU的场景，其通信延迟通常低于环形算法。
  • 关注点：更侧重于降低延迟 (latency)。
  • 混合算法 (Hybrid Algorithm)
  • 背景：训练集群的网络通常是异构的，即服务器内部的GPU间通信（如NVLink）速度极快，而服务器之间的通信（如以太网）相对较慢。
  • 工作方式：采用分层策略。例如，先在每台服务器内部用高速连接完成第一轮数据聚合（Reduce），然后再由每台服务器的代表GPU在服务器之间完成聚合，最后再将结果广播回服务器内部的各张GPU。
  • 优点：充分利用硬件拓扑的局部性，最大限度地减少跨节点慢速网络的通信量。

7.1.2 自动合成算法 (Synthesized Algorithm)

这是一种更高级的方法，它不依赖于固定的预定义算法，而是根据当前的硬件拓扑结构，自动生成一个定制化的、最优的通信算法。

  • 核心思想：与其使用“通用”方案，不如为“特定”的GPU连接方式量身打造最高效的通信路径和调度计划。
  • 实现方式：
  • 将通信算法的生成问题，建模为一个数学优化问题（如使用SMT、MILP等求解器）。
  • 或者提供一种领域特定语言 (DSL)，让系统可以编译生成针对特定硬件的优化代码。
  • 或者在运行时探测网络，动态构建通信计划。
  • 优点：在许多情况下，生成的定制化算法性能可以超越通用的预定义算法。

7.2 通信调度

在分布式训练中，GPU 的工作状态通常是 计算 -> 通信 -> 计算 -> 通信 ... 的循环。如果严格按顺序执行，那么在通信时，GPU计算单元就会空闲等待，造成资源浪费。

通信调度的核心目标就是打破这种严格的顺序，通过重新安排任务的执行顺序，使得通信操作可以和计算操作同时进行。理想情况下，当GPU在进行网络通信时，它的计算核心也没闲着，仍在处理其他计算任务。

这种“重叠”技术可以有效隐藏通信带来的延迟，从而提升GPU的利用率和整体训练吞吐量。

7.2.1 先进先出调度 (FIFO-based Scheduling)

这是最基础的调度策略，也被称为“无等待反向传播 (wait-free backpropagation)”。

  • 工作方式：在反向传播过程中，梯度是从模型的最后一层逐层向前计算出来的。一旦某一层的梯度计算完成，系统就立即将其放入一个先进先出 (FIFO) 队列，开始进行网络通信（如AllReduce），而不必等待所有层的梯度都计算完毕。
  • 优化：为了提高网络效率（大块数据传输效率更高），系统通常会等待一小段时间，将几层的梯度打包 (fuse) 成一个更大的数据块再进行通信。PyTorch的DDP就采用了这种策略。
  • 优点：实现简单，能实现基本的计算与通信重叠。

7.2.2 基于优先级的调度 (Priority-based Scheduling)

FIFO策略并不总是最优的。因为反向传播计算梯度的顺序（从后到前）和下一轮前向传播使用参数的顺序（从前到后）是相反的。这可能导致关键的通信任务被延后，阻塞后续计算。

  • 工作方式：不再遵循“先算完先通信”的原则，而是为不同的通信任务设置优先级。
  • 一个简单的优先级策略：越是靠前的层（在下一轮前向传播中越早被使用），其梯度通信的优先级就越高，从而让这部分参数能尽快更新完毕。
  • 更高级的策略 (PACE)：引入抢占 (preemption) 机制。当一个高优先级的（通常是小的）通信任务到来时，可以暂停当前正在进行的大通信任务，让高优先级任务“插队”先走，以避免关键路径被阻塞（即“队头阻塞”，Head-of-Line Blocking）。
  • 特定场景优化 (Lina for MoE)：在混合专家（MoE）模型中，All-to-All通信是关键瓶颈。因此，调度器会优先保证All-to-All的带宽，让AllReduce等其他通信在空闲时“见缝插针”地执行。

7.2.3 基于分解的调度 (Decomposition-based Scheduling)

这是最灵活、最前沿的调度思想。它将原本粗粒度的计算和通信任务，分解成许多更小的微任务 (fine-grained tasks)，从而创造出更多的重叠机会。

  • 核心思想：把“大石头”敲成“小石子”，可以更灵活地填充计算和通信之间的“缝隙”。
  • 分解的维度：
  1. 通信分解 (Communication Decomposition)：将一个大的AllReduce或AllGather操作，分解成一系列底层的点对点Send/Recv操作。这样，系统就可以将这些小的Send/Recv操作穿插在计算的间隙中。
  2. 计算分解 (Computation Decomposition)：
  • 与通信重叠：将一个大的矩阵乘法计算也切分成小块。每算完一小块结果，就立刻开始对这一小块进行通信，而不是等整个大矩阵算完。
  • 梯度计算解耦 (Out-of-order backprop)：反向传播会产生两种梯度：用于更新本层权重的权重梯度 (weight gradient) 和用于计算上一层梯度的输出梯度 (output gradient)。传统方法是同时计算两者。分解调度则将它们解耦，可以灵活地先计算其中一个，从而更自由地安排权重梯度的通信时机，以实现更好的重叠效果。

通过这种细粒度的分解和重排，系统可以实现接近完美的计算与通信重叠，最大限度地提升硬件利用率。

7.3 网络聚合

核心思想：让网络设备帮忙计算

在标准的AllReduce中，所有GPU将各自的梯度数据发送到一个（或多个）GPU上，由GPU完成求和计算，然后再将结果分发回去。这个过程会产生大量的数据流量，并占用GPU的计算资源。

网络内聚合 (INA) 的核心思想是：不再让GPU承担全部的聚合计算任务，而是利用网络设备（如交换机）自身的计算能力，在数据传输的“途中”就完成一部分或全部的聚合操作（如求和）。

效果：

  • 减少网络流量：交换机向上级网络或目标GPU发送的是聚合后的结果，而不是原始的多份数据，从而显著减少了网络主干链路的负载。
  • 降低GPU负担：GPU从繁重的聚合计算中解放出来，可以更专注于模型本身的计算。
  • 降低延迟：数据在网络路径上就被处理，减少了来回传输的延迟。

7.3.1 基于以太网的聚合 (Ethernet-based Aggregation)

这类方法主要依赖于可编程交换机 (Programmable Switches)，通过在交换机上运行特定的程序来实现聚合功能。

  • 代表技术：SwitchML
  • 工作方式：在交换机上实现AllReduce逻辑。由于交换机的内存很小，无法存储整个梯度，它采用流式聚合 (streaming aggregation) 的方式，一次只处理一小部分数据，处理完就向上发送，从而在有限的硬件上完成对大张量的聚合。
  • 局限性：早期的实现（如SwitchML）通常不支持直接进行浮点数运算，需要先将浮点数转换成特殊的整数格式，处理起来比较麻烦且可能损失精度。同时，与主流训练框架的集成也比较困难。
  • 后续发展：
  • 支持浮点数 (FPISA)：通过在可编程交换机上直接实现浮点数计算逻辑，解决了数据类型转换的问题。
  • FPGA辅助实现 (NetReduce)：使用FPGA（现场可编程门阵列）作为交换机的“外挂”，在FPGA上实现聚合逻辑，灵活性更高。
  • 面向多租户环境 (PANAMA, ATP)：研究如何在共享的集群环境中，为多个同时运行的训练任务公平且高效地分配网络内聚合资源。

7.3.2 基于 InfiniBand 的聚合 (Infiniband-based Aggregation)

这是目前在大型LLM训练集群中最成熟、最常用的方案。

  • 代表技术：NVIDIA SHARP (Scalable Hierarchical Aggregation Protocol)
  • 性质：这是NVIDIA提供的一项专有技术，内置于其Mellanox InfiniBand交换机和部分NVIDIA GPU/NVSwitch中。
  • 工作方式：
  1. 交换机内部集成了专用的计算单元（如FPU，浮点运算单元）。
  2. 当GPU们发起一个AllReduce操作时，它们的数据包会被交换机识别。
  3. 交换机使用其内置的计算单元，直接对流经它的数据包进行聚合计算（如求和）。
  4. 然后，交换机只将聚合后的结果发送出去。
  • 优势：
  • 硬件原生支持：直接在硬件层面实现，性能极高，可以达到线速（line rate）处理。
  • 功能完善：支持多种集体通信操作（AllReduce, Reduce等），支持多种数据类型（16/32/64位整数和浮点数）。
  • 易于使用：与NVIDIA的通信库NCCL无缝集成。开发者在使用NCCL时，只需开启相应选项，系统就会自动利用SHARP功能，无需修改代码。
  • 生产就绪：已经在最新的InfiniBand（如NDR）和NVSwitch中广泛部署，是许多大型LLM训练集群的标准配置。

8 容错能力 FAULT TOLERANCE

大语言模型（LLM）的训练具有两大特点：

  • 时间长：通常需要持续数周甚至数月。
  • 规模大：使用数千甚至上万张GPU组成的庞大集群。

训练系统包含从底层硬件到上层软件的众多复杂组件。

LLM训练普遍采用同步训练 (synchronous training) 模式，这意味着所有GPU必须步调一致。这种模式的致命弱点在于：

任何一个组件（a single point of failure）发生故障，都会导致整个训练任务暂停。

因此，为如此漫长且庞大的训练过程提供强大的容错机制 (fault tolerance mechanisms)，确保其可靠性，是至关重要的。

8.1 大语言模型故障分析

来自工业界的真实数据显示，LLM训练中的故障极其频繁：

  • Meta OPT (175B)：在992张A100上训练，两周内中断超过40次。
  • ByteDance (MegaScale)：在12,288张GPU上训练，几周内发生上百次故障。
  • Meta LLaMA3：在16,384张H100上训练，54天内中断466次。

这个现象的核心原因是规模效应：即使单个节点的故障率很低（如1.5%），当集群规模扩大到上千个节点时，整个集群每天发生故障的概率会激增到80%以上。

故障来源多样，但主要可以归为几类：

  1. 硬件故障 (主要元凶，占比约80%)：
  • GPU 问题：最常见，如CUDA错误、ECC内存错误。新一代GPU（A100/H100）由于技术新、功耗高，初期故障率更高。
  • 网络/互联问题：如NVLink错误、NCCL超时错误。用户最常看到的“NCCL error”通常根源于硬件问题。
  • 环境问题：如机房过热导致GPU故障。
  1. 软件问题：
  • 训练框架的Bug、数据处理流程出错、库依赖冲突等。
  1. 模型自身的不稳定性：
  • 训练过程中出现损失尖峰 (Loss Spikes)、梯度爆炸、数值溢出等问题。
  1. 外部因素：
  • 数据中心断电、冷却系统故障等。

故障会导致大量的GPU资源浪费，主要体现在两个方面：

  1. 故障恢复时间 (Failure Recovery)：
  • 传统恢复方式是回滚到上一个检查点 (Checkpoint) 并重启。
  • 但在重启前，运维人员需要花费大量时间去定位、诊断和修复故障硬件，这个过程可能长达数小时甚至数天，导致训练长时间停滞。
  1. 性能下降 (Performance Degradation)：
  • 掉队者 (Stragglers)：集群中某些节点因为硬件或网络异常，计算速度变慢。由于同步训练需要等待最慢的那个节点，这会严重拖慢整个集群的训练速度。

一个惊人的案例：Meta训练175B的OPT模型，理论上只需要25天，但实际花费了约57天。这意味着超过一半 (56%) 的时间都被用于处理各种故障，凸显了容错问题对资源利用率和训练效率的巨大影响。

8.2 异常检测

快速检测和诊断LLM训练中的故障是维持训练稳定和高效的关键。这个过程被称为异常检测。

8.2.1 统计监控 (Statistical Monitoring)

这是一种被动式、持续性的检测方法，通过在训练过程中实时收集和分析各种指标来发现异常。

  • 工作模式：
  1. 数据收集：在每台GPU节点上部署一个监控进程，负责收集硬件状态、运行时统计数据等信息。
  2. 心跳机制：监控进程定期将收集到的数据作为心跳消息发送到一个中央监控节点。
  3. 分析与告警：中央节点分析收到的数据。如果某个节点长时间没发心跳，就认为它已宕机；如果某些指标出现异常（如GPU利用率突然下降），则触发告警。
  • 监控哪些关键指标？
  • GPU 相关指标 (通过NVIDIA DCGM工具收集)：
  • 计算利用率：SM（流多处理器）的利用率、占用率等，反映GPU是否在全力工作。
  • 内存健康：如 row-remapping 计数，可以指示GPU显存是否存在物理缺陷。
  • 互联带宽：PCIe和NVLink的流量速率，反映节点内数据传输是否正常。
  • 网络相关指标：
  • RDMA流量：监控InfiniBand网络的流量指标。
  • 连接状态：如RDMA连接的IP、QP号，TCP超时等。
  • 集体通信活动：PyTorch内置的NCCL flight recorder和NCCLX可以记录AllReduce等通信操作的内部状态和耗时，帮助诊断NCCL超时等网络问题。
  • 网络带宽健康检查：定期在节点对之间测试网络带宽，以发现性能下降的网卡或链路。
  • 训练日志 (Training Logs)：分析训练过程中打印的错误日志。
  • 机器学习辅助：利用机器学习算法分析历史监控数据，自动检测异常模式。

Google的TPU超级计算机也采用类似的机制，通过一个名为healthd的守护进程实时监控芯片间互联、PCIe等链路的健康状况。

8.2.2 主动验证 (Proactive Validation)

这是一种主动式、预防性的检测方法，旨在训练任务开始前或训练间隙对系统进行“体检”，以提前发现潜在问题。

  • 核心思想：与其等故障发生后再去诊断，不如提前运行一系列测试来验证系统各组件是否健康。
  • 工作模式与挑战：
  • 权衡：需要在验证的全面性和占用的时间之间做出权衡。过于全面的测试会占用宝贵的训练时间。
  • 轻量级测试 (MegaScale)：开发一套轻量级的测试工具，如节点内网络和NCCL测试，快速诊断常见问题。
  • 分层策略 (Vela)：采用两级策略。在所有节点上定期运行轻量级测试；只有当节点空闲时，才运行更深入、更耗时的侵入式测试。
  • 飞前检查 (Google TPU)：在用户任务开始前，强制执行一次“飞前检查”，包括端到端的功能测试和硬件健康扫描。
  • 基准测试套件 (SuperBench)：提供一个全面的基准测试工具集，可以对单个硬件组件进行评估，并允许用户根据需求选择测试项，以平衡验证时间和潜在风险。

8.3 基于检查点的恢复 (Checkpoint-Based Recovery)

检查点 (Checkpointing) 是LLM训练中最基础的容错方法：定期将模型状态（参数、优化器状态等）保存下来。一旦发生故障，可以从最近的一个检查点恢复训练，而不是从头开始。

但这带来一个两难的困境：

  • 频繁保存：会产生巨大的I/O开销，拖慢训练速度。
  • 不频繁保存：一旦故障，会丢失大量已完成的训练进度，浪费计算资源。

本节探讨的各种方法，都是为了解决这个困境，目标是实现既快又频繁的检查点保存。

8.3.1 持久化检查点 (Persistent Checkpointing)

这类方法将检查点保存在非易失性存储上，如SSD或云存储，确保数据在系统彻底崩溃后依然存在。保存过程通常分两步：

  1. 快照 (Snapshot)：将模型状态从GPU显存拷贝到CPU内存。
  2. 持久化 (Persist)：将CPU内存中的快照写入磁盘或云存储。

a) 同步检查点 (Synchronous Checkpointing)

  • 工作方式：暂停整个训练过程，专心执行快照和持久化。
  • 缺点：在保存期间，所有GPU都处于空闲状态，造成严重的资源浪费。这是最简单但最低效的方式。

b) 快照时暂停的检查点 (Snapshot-Stall Checkpointing)

这是对同步方式的重大改进，也是目前主流框架（如DeepSpeed, Megatron-LM）采用的方案。

  • 工作方式：
  1. 短暂暂停训练，只为了完成快照步骤（将数据从GPU拷贝到CPU）。这个过程通常很快，只需几秒钟。
  2. 立即恢复训练。
  3. 在后台异步地、并行地执行耗时最长的持久化步骤（将CPU内存的数据写入磁盘），这个过程与后续的训练同时进行，互不干扰。
  • 优化：
  • 分块与多线程 (TorchSnapshot)：将快照切成小块，用多个CPU线程并行写入磁盘，进一步加快持久化速度。
  • 优化恢复过程 (MegaScale)：恢复时，只让一个GPU从磁盘读取完整的检查点，然后通过高速的内部网络广播给其他GPU，避免所有GPU都去抢占磁盘I/O。

c) 异步检查点 (Asynchronous Checkpointing)

这是最理想但实现也最复杂的方式，目标是几乎不暂停训练。

  • 工作方式：通过精密的流水线 (Pipelining) 设计，将快照和持久化两个步骤都与训练过程重叠起来。
  • 挑战：需要非常小心地处理数据依赖和一致性问题，确保在保存参数时，它不被下一轮的梯度更新所“污染”。
  • 实现：通常采用分层流水线 (layer-wise pipelining) 的方式，在模型计算到某一层时，异步地保存另一层的状态。

8.3.2 内存中检查点 (In-Memory Checkpointing)

这类方法认识到，将检查点写入远程磁盘是速度的主要瓶颈。

  • 核心思想：不把检查点写入慢速的磁盘，而是将其保存在其他计算节点的CPU内存中，或者专门的内存存储系统 (如Redis) 中。
  • 优点：
  • 速度极快：内存到内存的传输速度远高于内存到磁盘，I/O开销大大降低。
  • 频率极高：可以实现非常高频率的检查点保存（例如每分钟一次），即使发生故障，也几乎不损失训练进度。
  • 容错增强 (REFT)：借鉴RAID磁盘阵列的思想，采用纠删码 (erasure coding) 技术。即使保存检查点的某个内存节点也宕机了，仍然可以从其他节点的数据中恢复出完整的检查点。
  • 局限性：无法应对整个集群断电等灾难性故障。因此，最佳实践是采用混合策略：
  • 高频率地保存到内存，用于应对常规的单点故障。
  • 低频率地保存到持久化存储，用于应对灾难性故障。

8.4 无需检查点的恢复 (Checkpoint-Free Recovery)

传统的容错方法是“停止-回滚-重启”，即训练暂停，从上一个检查点恢复，这个过程非常耗时。

无需检查点的恢复 (Checkpoint-Free Recovery) 是一种更高级的容错策略，其核心思想是：当故障发生时，不暂停和回滚训练，而是动态地修复问题，让训练以极小的中断甚至无中断地继续下去。

这依赖于强大的自动故障检测机制来及时发现问题。

8.4.1 实时迁移 (Live Migration)

  • 利用的资源：数据并行训练中天然存在的模型冗余。在数据并行下，我们本来就有多份完全相同的模型副本在不同的GPU组上运行。
  • 工作流程：
  1. 当某个节点（GPU）发生故障时，系统自动检测到故障。
  2. 系统动态地重构集群，例如将故障节点踢出，并从资源池中拉取一个健康的备用节点。
  3. 从一个健康的、正在运行的模型副本那里，将最新的、实时的模型状态拷贝到这个新加入的备用节点上。
  4. 训练在新配置的集群上继续进行。
  • 关键点：这种方法避免了从慢速的磁盘加载陈旧的检查点，而是直接从内存中的“活”数据进行恢复，因此速度快得多，中断时间极短。

8.4.2 模块冗余 (Module Redundancy)

  • 核心思想：这是一种“热备份”策略。它不是在故障后才去恢复状态，而是在正常训练时就准备好了备胎。
  • 工作流程：
  1. 在正常训练时，系统会利用一些空闲的计算资源（例如流水线并行中的“气泡”时间），去冗余地执行一部分计算。这个“冗余模块”就像一个影子，实时跟随着“主模块”的状态。
  2. 当主计算模块发生故障时，系统检测到问题。
  3. 系统立即将原本要发往故障模块的计算任务，重定向（route）到这个已经准备就绪的冗模块。
  • 关键点：恢复动作不是迁移状态 (migrating state)，而是重定向计算 (rerouting computation)。因为冗余模块一直在同步运行，所以切换几乎是瞬时的，可以实现真正的“无中断”恢复。这是比实时迁移更极致的容错方案。
posted @ 2025-12-02 00:20  Orzjh  阅读(0)  评论(0)    收藏  举报