StreamSparkAI

概览/核心摘要 (Executive Summary)

本讲座（Stanford CS336 Language Modeling from Scratch, Spring 2025, 07 Parallelism 1）深入探讨了在多机器环境下训练大规模语言模型所需的并行化策略。核心目标是解决单GPU在计算能力和内存容量上的瓶颈，通过跨机器并行实现对巨大模型的有效训练。讲座首先介绍了并行化的动因（计算需求和内存限制）以及相关的硬件基础，特别是不同层级的网络通信（如GPU间NVLink/NVSwitch与跨机器InfiniBand的速率差异，以及TPU的环形网络拓扑）。接着，重点阐述了三种主要的并行化范式：

1. 数据并行 (Data Parallelism)：核心思想是在多个GPU上复制模型参数，并将数据批次（batch）切分给不同GPU处理。讲座详细介绍了朴素数据并行及其内存开销问题，并重点讲解了Zero Redundancy Optimizer (ZeRO) 的三个阶段（ZeRO-1、ZeRO-2、ZeRO-3/FSDP），它们通过逐步分片优化器状态、梯度和模型参数，显著减少内存占用，同时分析了其通信开销和实现技巧（如ZeRO-3中计算与通信的重叠）。
2. 模型并行 (Model Parallelism)：当模型过大无法在单个GPU内容纳时，需要将模型本身切分。讲座讨论了两种主要方式：
   • 流水线并行 (Pipeline Parallelism)：按层切分模型，不同GPU负责不同层段，通过传递激活值进行协作。挑战在于“流水线气泡”（GPU空闲），可通过微批次（micro-batching）和更复杂的调度（如Zero-Bubble Pipelining）缓解，但实现复杂。
   • 张量并行 (Tensor Parallelism)：在模型内部（如矩阵乘法）进行切分，将单个张量操作分布到多个GPU上。通信开销大，通常用于节点内高速互联的GPU（如8卡NVLink）。
3. 激活并行/序列并行 (Activation/Sequence Parallelism)：针对训练过程中激活值占用大量内存的问题，通过在序列维度上切分操作（如LayerNorm、Dropout）来减少单个GPU的激活内存占用。

讲座强调，实际大规模训练通常结合使用这些策略（所谓的3D或4D并行），并给出了选择和组合策略的经验法则，如优先使用张量并行填满节点内带宽，再用流水线并行或FSDP跨节点扩展以适应模型大小，最后用数据并行进一步扩展计算规模。批处理大小（batch size）被视为一种关键资源，影响各种并行策略的效率。最后，通过Megatron-LM、Llama 3、Gemma 2等实际案例，展示了这些并行策略在业界顶尖模型训练中的应用和效果，并提及了大规模训练中硬件故障等实际挑战。

讲座核心内容总结

一、引言与并行化动机

• 目标：从优化单个GPU吞吐量转向理解训练真正大型模型所需的复杂性和细节，实现跨机器并行。
• 核心挑战：
  • 模型大小 (Memory Concerns)：大型模型参数量巨大（数十亿甚至上万亿），远超单个GPU内存容量。
    • 引用数据：GPU内存虽然在增长，但速度不及模型参数增长。
  • 训练速度 (Compute Concerns)：需要利用多台服务器的计算资源以快速训练模型。
    • 引用数据：单个GPU的浮点运算能力（FLOPS）虽呈超指数增长，但仍不足以快速扩展。训练顶级模型需依赖拥有ExaFLOPS级别算力的超级计算机。
  • 异构通信 (Communication)：不同层级（GPU内部、机器内部、机器之间）的通信速度差异显著，影响并行策略选择。
• 讲座结构：
  1. 网络基础 (Networking Basics)
  2. 并行化策略 (Parallelization Strategies)
  3. 案例研究 (Case Studies)

二、硬件基础与网络通信

• GPU服务器硬件层级：
  • 机器内部 (Intra-node)：通常一台机器包含多个GPU（如8个），通过高速互联技术（如NVIDIA的NVLink和NVSwitch）连接，通信速度极快。
  • 机器之间 (Inter-node)：GPU通过网络交换机（如InfiniBand HDR）与其他机器上的GPU通信，速度相对较慢。
    • 引用数据：讲者提到一个例子中，InfiniBand HDR每通道吞吐量比NVSwitch慢约8倍。
  • 大规模集群：当GPU数量超过一定阈值（如256个GPU），可能需要更复杂的网络拓扑（如叶脊网络 Leaf-Spine Switches），可能引入更多通信瓶颈。
• TPU网络设计：
  • Google TPU采用不同的网络方法，如环形网格 (Toroidal Mesh)，每个TPU芯片与其邻居高速连接。
  • 这种设计对于集体通信操作（如All-Reduce）非常高效，即使不采用全连接（all-to-all）拓扑。
• 集体通信操作 (Collective Communication Operations)：
  • 是并行算法的构建模块，理解其性能特征对优化至关重要。
  • All-Reduce：所有机器拥有数据，进行规约操作（如求和），结果分发回所有机器。通信成本约 2 * N (N为数据大小)。
  • Broadcast：一个机器的数据复制到所有其他机器。成本约 1 * N。
  • Reduce：不同机器的数据规约后发送到单个机器。
  • All-Gather：每个机器拥有数据的一部分，将其复制到所有其他机器，最终每个机器拥有完整数据。
  • Reduce-Scatter：每个机器拥有完整数据，对数据进行规约（如按行求和），然后将结果的不同部分分发给不同机器。
  • 重要等价关系：All-Reduce 操作在带宽受限情况下，其最优性能等价于一个 Reduce-Scatter 操作后接一个 All-Gather 操作。
    • 讲者强调：“this could be replaced with two operations, a reduced scatter. And I'll all gather... In the bandwidth limited regime, this is basically the best that you can do, all reduce.”

三、并行化策略的目标

• 将整个数据中心 (Data Center) 视为一个新的计算单元。
• 线性内存扩展 (Linear Memory Scaling)：随着GPU数量增加，可训练的最大模型大小也应线性增长。
• 线性计算扩展 (Linear Compute Scaling)：随着GPU数量增加，用于训练模型的有效计算量也应线性增长。
• 算法实现依赖于调用简单的集体通信原语 (Collective Communications Primitives)。

四、数据并行 (Data Parallelism)

• 核心思想：模型参数在所有GPU上复制，数据批次被切分到不同GPU上。
• 朴素数据并行 (Naive Data Parallelism)：
  • 每个GPU处理一部分数据，计算梯度。
  • 通过 All-Reduce 同步所有GPU的梯度。
  • 进行参数更新。
  • 计算扩展：良好，前提是每个GPU的微批次大小（micro batch size）足以饱和其计算单元。
  • 通信开销：每次批处理需传输 2 * P（P为参数量）的数据（All-Reduce）。若批次大，可掩盖通信。
  • 内存扩展：差。每个GPU需存储完整模型参数和优化器状态。
• 内存瓶颈分析：
  • 实际内存占用远超参数本身，主要来自优化器状态 (Optimizer States)，尤其是Adam等自适应优化器。
    • 引用数据：可能需要存储约16字节/参数，包括参数本身（如BF16占2字节）、梯度（2字节）、Adam主权重（4字节）、Adam一阶矩（2-4字节）、Adam二阶矩（2-4字节）。“you need to store something like five copies of your weights.”
    • 一个7.5B参数模型在64个加速器上，若朴素复制，总内存占用巨大（示例中为120GB）。
• Zero Redundancy Optimizer (ZeRO)：旨在减少数据并行中的冗余内存。
  • ZeRO Stage 1: 优化器状态分片 (Optimizer State Sharding)
    • 将优化器状态（如Adam的一阶和二阶矩）分片到不同GPU。每个GPU仍持有完整的模型参数和梯度。
    • 流程：
      1. 各GPU计算其数据子集的完整梯度。
      2. 对梯度进行 Reduce-Scatter：GPU i 获得其负责的那部分参数（对应其拥有的优化器状态分片）的全局梯度总和。
      3. GPU i 使用其本地的优化器状态分片和收到的梯度总和，更新其负责的那部分参数。
      4. 对更新后的参数分片进行 All-Gather，使所有GPU恢复完整的最新参数。
    • 通信成本：Reduce-Scatter + All-Gather，与 All-Reduce 相同。在带宽受限情况下，内存节省“几乎免费”。
  • ZeRO Stage 2: 梯度分片 (Gradient Sharding)
    • 在Stage 1基础上，进一步将梯度也进行分片。
    • 在反向传播过程中，当计算完一层梯度后，立即通过 Reduce 操作将其发送给负责该参数分片的GPU，并释放本地梯度内存。
    • 通信成本：总量仍为 2 * P，但由于逐层同步，可能引入额外开销。
  • ZeRO Stage 3: 参数分片 (Parameter Sharding) - 即 FSDP (Fully Sharded Data Parallel)
    • 将模型参数、梯度、优化器状态全部进行分片。
    • 在计算过程中按需请求参数。
    • 前向传播：All-Gather 当前层权重 -> 执行前向计算 -> 释放该层权重。
    • 反向传播：All-Gather 当前层权重 -> 执行反向计算 -> Reduce-Scatter 梯度 -> 释放权重。
    • 通信成本：增加到约 3 * P。
    • 效率关键：通过重叠计算与通信 (Overlapping Communication and Computation) 来隐藏延迟，即在当前层计算时预取下一层参数。
    • 引用数据：一个例子中，使用FSDP可以将8xA100 80GB节点上能容纳的模型从约6B参数提升到50B参数。
• 数据并行的关键资源：批处理大小 (Batch Size)。
  • 并行度不能超过批处理大小。
  • 批处理大小存在收益递减效应（OpenAI的“临界批处理大小”研究）。

五、模型并行 (Model Parallelism)

• 核心思想：当模型过大无法放入单个GPU时，将模型本身切分到不同GPU上。主要传递激活值而非参数。
• 流水线并行 (Pipeline Parallelism)
  • 概念：沿模型深度（层）切分，不同GPU负责一部分连续的层。
  • 挑战：流水线气泡 (Pipeline Bubble)，即GPU空闲等待。
    • 朴素实现中，GPU利用率仅为 1/N (N为流水线阶段数)。
  • 优化：
    • 微批次 (Micro-batching)：将大批次切分为小批次，使流水线各阶段能重叠工作，减小气泡。气泡大小与 (阶段数-1)/微批次数 成正比。
    • Zero-Bubble Pipelining (如DeepSpeed DualPipe)：更复杂的调度，将权重梯度计算等可延迟任务安排在气泡时间内执行。实现极为复杂。
      • 讲者趣闻：“actually there's two people in the group that understand how the pipeline parallel in our infra works. One person left. And so there's a single load bearing person in our training infra.”
  • 优点：节省参数和激活内存，点对点通信对慢速网络链路友好。
  • 缺点：消耗批处理大小以减小气泡，实现复杂。
• 张量并行 (Tensor Parallelism)
  • 概念：沿模型宽度（隐藏层维度）切分，将单个大运算（如矩阵乘法）分解为多个小运算，在不同GPU上并行执行。
  • MLP示例：Y = GeLU(XA)，Z = Dropout(YB)。
    • 将权重矩阵A切分为A1, A2，B切分为B1, B2。
    • 前向：输入X复制给两个GPU -> XA1, XA2 -> (GeLU) -> Y1, Y2 -> Y1B1, Y2B2 -> All-Reduce 合并 Z1, Z2 得到Z。
    • 反向：类似地，梯度在某个点需要 All-Reduce。
  • 通信：每层有同步点（如 All-Reduce），通信量大，需要高速互联。
  • 经验法则：通常用于单节点内部 (within a single node)，如8个通过NVLink连接的GPU。扩展到节点外性能会急剧下降。
    • 引用数据：HuggingFace教程显示，TP扩展到16个设备时吞吐量下降42%，32个设备下降65%。
  • 优点：不产生流水线气泡，不消耗批处理大小，实现相对简单。
  • 缺点：通信开销大。

六、激活内存与序列并行 (Activation Memory & Sequence Parallelism)

• 激活内存问题：即使使用模型并行，激活值仍可能占用大量内存，尤其对于大模型和长序列。
  • 某些激活（如LayerNorm、Dropout、Attention和MLP的输入）在标准张量并行下不会随并行度T线性缩减。
• 序列并行 (Sequence Parallelism)：
  • 针对上述未被张量并行有效切分的逐点操作（point-wise operations）。
  • 沿序列长度 (Sequence Dimension) 进行切分。例如，LayerNorm或Dropout操作在序列的不同位置上是独立的，可以将长序列切分给不同GPU处理。
  • 需要额外的 All-Gather 和 Reduce-Scatter 操作来同步跨序列分片的结果。
• 最终激活内存：结合张量并行(TP, 并行度T)、序列并行和激活重计算（如FlashAttention），每层最小激活内存可达 SBH * 34 / T (S:序列长度, B:批次大小, H:隐藏层大小)。

七、其他并行策略

• 上下文并行 (Context Parallelism) / 环形注意力 (Ring Attention)：
  • 用于处理超长上下文的注意力计算。
  • 将Queries分配给不同机器，Keys和Values在机器间以环形方式传递计算。利用了FlashAttention的分块计算思想。
• 专家并行 (Expert Parallelism - MoE)：
  • 将MLP层替换为多个小型的“专家”MLP，并进行稀疏激活。
  • 概念上类似张量并行（将一个大MLP切分为多个小MLP并分散），但增加了路由（routing）的复杂性，因为专家激活是稀疏且动态的。

八、组合并行策略 (3D/4D Parallelism)

• 资源权衡：内存、带宽、计算、批处理大小都是有限资源，需平衡。
• Google TPU书籍的图示：最优并行策略取决于 批处理大小 / GPU数量 的比率。
  • 比率过小：通信受限。
  • 比率适中：混合FSDP和模型并行（主要是张量并行）可达计算受限。
  • 比率够大：纯FSDP即可计算受限。
• 通用经验法则 (Rule of Thumb)：
  1. 适配模型和激活内存 (Fit Model & Activations)：
     • 首先使用张量并行 (Tensor Parallelism)，直到用尽单机内GPU的高速互联（如8卡）。
     • 若仍无法容纳，则跨机器使用 ZeRO Stage 3 (FSDP) 或 流水线并行 (Pipeline Parallelism)。
  2. 扩展计算能力 (Scale Compute)：
     • 使用数据并行 (Data Parallelism) (如ZeRO Stage 1或FSDP的DP部分)扩展到所有可用GPU。
  3. 小批处理大小优化：
     • 若全局批处理大小受限，可使用梯度累积 (Gradient Accumulation) 来增大有效批处理大小，减少同步频率，提高通信效率。
• 讲者提到可能存在“五维并行”，但尚不清楚第五维具体指什么。

九、实际案例研究

• Megatron-LM (2021)：
  • 展示了从1.7B到1T参数模型的训练配置，实现了40-52%的理论峰值FLOPS。
  • 策略：张量并行最高到8，然后根据模型大小增加流水线并行度，数据并行度相应调整。
  • 激活重计算有助于使用更大批次，从而掩盖流水线并行开销。
• 近期语言模型：
  • Olmo, DeepSeek (v1), E AI：通常采用FSDP (或ZeRO Stage 1) + 张量并行 + 流水线并行。
  • DeepSeek (v3)：16路流水线并行，64路专家并行，ZeRO Stage 1数据并行。
  • E AI Lightning (MoE)：用专家并行替代张量并行。
  • Llama 3 (Meta)：
    • 采用张量并行(8)，上下文并行(CP，用于长序列)，流水线并行，数据并行。
    • 并行优先级（带宽需求从高到低）：TP -> CP -> PP -> DP。
    • 提及大规模训练的挑战：训练过程中遇到148次GPU故障中断，占总中断30%。更可怕的是静默数据损坏。
  • Gemma 2 (Google TPU)：
    • 采用ZeRO Stage 3 (FSDP的TPU等效实现) + 模型并行 + 数据并行。TPU的互联特性使其能更好地支持模型并行。

十、结论

• 扩展到一定规模以上必须依赖多GPU、多节点并行。
• 没有单一的万能解决方案，需要结合数据并行、模型并行（流水线、张量）、激活/序列并行等多种方法。
• 存在简单且可解释的经验法则来指导实践中的并行策略选择与组合。