三：分布式并行策略

数据并行 (DP) 和分布式数据并行 (DDP) 在实现上有何不同？DDP 为什么更快？¶

数据并行是最基础的分布式训练策略。PyTorch提供了两种实现：DataParallel（DP）和DistributedDataParallel（DDP）。二者虽然目标一致——将数据分片到多张GPU上并行计算——但实现机制和效率有天壤之别。

DP的实现方式：DP采用单进程多线程架构。主GPU（通常是GPU0）负责承载模型副本、处理输入数据的分发、收集各GPU的输出、计算损失、反向传播，最后将梯度散射到各GPU并更新参数。其他GPU仅承担前向传播的计算任务，不参与反向传播的梯度计算和图构建。这种架构带来了几个致命缺陷：

主GPU负载过重：所有数据的聚合、损失的汇总、梯度的分发都经过主GPU，形成通信和计算瓶颈。主GPU显存占用远高于其他卡，因为它需要存储完整的计算图、梯度和优化器状态。
Python GIL（全局解释器锁）限制：DP基于多线程，Python的GIL使得同一时刻只有一个线程能执行Python代码。虽然底层CUDA操作可以并行，但线程间的调度和通信仍受GIL阻碍，导致CPU利用率低下。
通信效率低：每次前向传播结束，各GPU的输出需要拷贝回主GPU；反向传播后，梯度又从主GPU散射回各卡。这些数据搬运涉及多次CPU-GPU通信和显存拷贝，开销巨大。

DDP的实现方式：DDP采用多进程架构，每个GPU上运行一个独立的Python进程。每个进程拥有自己的模型副本、优化器、数据加载器。训练开始时，所有进程通过DistributedSampler获取不同的数据子集，独立执行前向传播、计算损失、反向传播。在反向传播过程中，梯度自动通过高效的集合通信（AllReduce）在GPU间同步，无需经过任何中心节点。更新参数时，所有进程使用相同的全局梯度，因此模型副本始终保持一致。

DDP为什么更快：

消除单点瓶颈：没有主GPU，所有设备地位平等。通信直接在GPU间进行，不经过CPU内存拷贝。NCCL库可以利用NVLink和InfiniBand实现GPU直连通信（RDMA），带宽极高。
多进程避开GIL：每个进程独立运行Python解释器，真正实现并行。CPU计算资源被充分利用。
通信与计算重叠：DDP将梯度同步拆分为细粒度的“桶”（bucket），在反向传播进行的同时，一旦某个桶内的所有参数梯度计算完毕，立即启动该桶的异步AllReduce。这样，梯度通信被隐藏在计算时间之中，几乎不产生额外延迟。
通信次数减半：相比于DP的“收集-散射”两次传输，AllReduce一次完成全局梯度求和并广播到所有设备，通信量更低。

因此，DDP已成为当前分布式训练的唯一标准，DP仅用于极简单或单机单卡的调试场景。

画出 DDP 的梯度同步流程：反向传播 + AllReduce 梯度。¶

DDP的梯度同步是嵌入在反向传播过程中的，而非等整个反向传播结束才启动。这一设计精妙地实现了通信与计算的重叠。

前向传播
  │
  ├─ 各GPU独立前向，计算损失 loss
  │
反向传播开始
  │
  ├─ 计算输出层的梯度
  │
  ├─ 逐层反向传播，计算各层参数的梯度
  │     │
  │     ├─ 当某个“梯度桶”内的所有参数梯度都已计算完成时
  │     │     │
  │     │     └─ 立即启动该桶的 AllReduce（异步）
  │     │           │
  │     │           └─ 通信进行的同时，反向传播继续向更浅层推进
  │     │
  │     └─ 重复此过程，直到所有桶的梯度都完成 AllReduce
  │
反向传播结束
  │
  └─ 所有参数梯度已全局同步，各GPU调用 optimizer.step() 更新参数

关键设计：DDP预先将模型参数分组（默认按层），每个组称为一个“梯度桶”。反向传播从输出层向输入层进行。当某个桶（例如第N层）的所有参数梯度都计算出来后，立即将该桶的梯度加入通信队列，启动AllReduce。此时反向传播继续向前一层推进，AllReduce在后台异步执行。由于深层参数的梯度先被计算，其AllReduce可以与浅层的梯度计算和通信重叠。最终，整个模型的梯度同步延迟几乎完全被隐藏。

DDP 中的 AllReduce 是同步的还是异步的？有没有异步的梯度同步方法？¶

DDP中的AllReduce是异步的。具体来说，每个梯度桶的AllReduce操作被投递到NCCL的通信流（CUDA Stream）中，DDP的主计算流不会等待AllReduce完成，而是继续反向传播计算后续层的梯度。只有当所有梯度桶的AllReduce都完成，DDP才会进入参数更新阶段。这种“投递即忘，最后统一等待”的模式实现了通信与计算的最大重叠。

除了DDP内置的异步AllReduce，还存在更激进的异步梯度同步方法：

异步SGD (Asynchronous SGD)：各工作节点计算完梯度后立即更新自己的本地参数，无需等待全局同步。梯度的接收和发送完全异步，参数服务器在后台更新全局模型。这种方法消除了所有同步等待，但会引入梯度延迟——节点使用的全局参数可能是基于几秒前的旧梯度，导致训练不稳定和精度下降。实践中极少用于大模型训练。
梯度压缩 (Gradient Compression)：通过量化、稀疏化或低秩近似减少通信量。例如1-bit Adam将梯度压缩为1比特符号，极大降低通信延迟，同时保持同步更新的语义。
本地SGD (Local SGD)：每个节点独立更新本地模型多步，然后周期性全局平均。减少了通信频次，但也牺牲了全局同步的精度。

总的来说，DDP的异步AllReduce + 通信计算重叠是当前效率与精度平衡的最佳方案，更激进的异步方法通常只在特定通信瓶颈下考虑。

解释什么是“梯度桶” (Gradient Bucket) 在 DDP 中的作用。¶

梯度桶是DDP实现通信与计算重叠的核心机制。它将模型的所有参数按反向传播的执行顺序分组，每个组就是一个“桶”。默认情况下，桶的划分遵循模型层的顺序（反向传播从后往前），桶的大小由bucket_cap_mb参数控制（默认为25MB）。

桶的作用：

延迟通信以累积更大粒度：单个参数的梯度张量通常很小（如几十KB），若每个梯度都独立发起AllReduce，通信启动开销将极为昂贵。桶将多个小梯度累积成较大的通信块（如25MB），一次AllReduce处理整个桶，大幅降低通信次数和延迟。
实现通信与计算的重叠：当反向传播完成某一个桶内最后一个参数的梯度计算时，DDP立即触发该桶的AllReduce。在AllReduce异步执行的同时，反向传播继续向前推进，计算下一个桶的梯度。这样，前一个桶的通信与后一个桶的计算在时间上重叠，隐藏了通信延迟。
自动管理依赖：DDP内置的Reducer会自动追踪每个桶内的梯度是否都已就绪。一旦就绪，启动AllReduce；所有桶完成后，才允许优化器更新参数。

桶的划分策略：通常按模型层从前到后的顺序（即反向传播的逆序）划分。例如，一个Transformer模型可能将最后几个FFN层分为一个桶，中间几个注意力层分为另一个桶。划分的粒度影响重叠效率：桶太大，等待最后一个梯度的时间过长，通信开始晚；桶太小，通信次数多。默认25MB是经过调优的平衡值。

为什么 DDP 中通常对参数分组进行 AllReduce？如何优化通信顺序？¶

分组AllReduce（即梯度桶机制）的目的如上所述：将小梯度累积成大块，以减少通信次数，同时实现通信与计算的重叠。如果不分组，每个梯度都独立AllReduce，会导致：

大量极小的通信操作，网络利用率极低，延迟高昂。
通信与计算无法有效重叠，因为AllReduce的启动和完成散布在整个反向传播过程中，难以管理。

优化通信顺序：DDP默认按照反向传播的自然顺序触发AllReduce，即从深层到浅层。这个顺序本身就有利于重叠——最深的层最先完成梯度计算，其通信可以与浅层计算重叠。在某些情况下，可以通过调整bucket_cap_mb或手动将通信更密集的层（如大矩阵乘法）提前划分到单独的桶中，进一步优化重叠效果。例如，将嵌入层（通常参数量大、位于模型首尾）单独分为一个桶，使它的通信不与中间层混淆。现代框架如DeepSpeed和Megatron-LM实现了更精细的通信调度，但核心思想都是梯度分组与异步通信。

张量并行 (TP) 的核心思想是什么？解决什么问题？¶

张量并行的核心思想是：将单层内的权重矩阵切分到多个GPU上，每张GPU只计算该层输出的一个部分。它解决的是单卡显存无法容纳整个模型层，或者单卡计算单层太慢的问题。

在数据并行和ZeRO中，每张GPU仍然需要处理完整的单层前向和反向计算（即需要看到完整的层权重，即使这些权重是通过AllGather临时收集的）。当模型单层参数量极大（例如GPT-3的某些层超过1B参数）时，单卡的计算和显存都无法承受。张量并行将层内的矩阵乘法切分，使得每张卡只需存储和计算层权重的一部分，从而突破了单卡显存和算力对单层大小的限制。

典型应用场景：当模型隐藏维度极大（如8192甚至更高）时，单层的QKV投影矩阵可能达到数亿参数，单卡显存无法容纳。通过TP将这些矩阵切分到节点内的多张卡上，利用NVLink高带宽进行通信，可以高效训练。

以 Megatron-LM 为例，说明如何对 Transformer 的 MLP 层进行张量并行。（按列切割和按行切割）¶

Megatron-LM是NVIDIA开源的张量并行框架，对Transformer的MLP层和自注意力层提供了标准的切分方案。MLP层由两个全连接层组成：第一层从hidden_dim升维到intermediate_dim（通常是4倍），经GeLU激活；第二层从intermediate_dim降维回hidden_dim。

切分方法：

第一层（升维层）按列切割：将权重矩阵W1 [hidden_dim, intermediate_dim]沿列方向（输出维度）切分为TP份。例如TP=2，则切分为W1_a [hidden_dim, intermediate_dim/2]和W1_b [hidden_dim, intermediate_dim/2]。输入X [batch, seq_len, hidden_dim]被广播到所有TP卡上。每张卡独立计算Y_a = GeLU(X @ W1_a)，Y_b = GeLU(X @ W1_b)。此时中间激活Y的维度从intermediate_dim降为intermediate_dim/TP。
第二层（降维层）按行切割：将权重矩阵W2 [intermediate_dim, hidden_dim]沿行方向（输入维度）切分。对应第一层的列切分，W2_a [intermediate_dim/2, hidden_dim]和W2_b [intermediate_dim/2, hidden_dim]分别放在两张卡上。每张卡计算Z_a = Y_a @ W2_a，Z_b = Y_b @ W2_b。此时每张卡的输出Z形状为[batch, seq_len, hidden_dim]，但只是部分和，需要将两张卡的结果通过AllReduce求和，得到完整的MLP输出。

通信模式：前向传播中，第二层按行切割后需要一次AllReduce来合并部分和。反向传播时，梯度∂L/∂Z本身就是完整的（因为AllReduce得到完整输出），但计算∂L/∂W2时每张卡只需本地梯度，计算∂L/∂W1时则需要将梯度通过AllReduce或ReduceScatter进行聚合。具体通信次数：前向1次AllReduce，反向1次AllReduce（或等效通信）。

张量并行中，MLP 的两次矩阵乘法如何分配？前向和反向的通信模式。¶

延续上例，TP=2的MLP层：

前向传播：

Y_a = X @ W1_a，Y_b = X @ W1_b （各自独立，无通信）
Y_a = GeLU(Y_a)，Y_b = GeLU(Y_b)
Z_a = Y_a @ W2_a，Z_b = Y_b @ W2_b （各自独立）
Z = AllReduce(Z_a + Z_b) （一次AllReduce，求和）

反向传播：

梯度∂L/∂Z已通过AllReduce得到，每张卡持有完整∂L/∂Z。
计算∂L/∂W2_a = Y_a^T @ ∂L/∂Z（本地），∂L/∂W2_b同理。
计算∂L/∂Y_a = ∂L/∂Z @ W2_a^T（本地），∂L/∂Y_b同理。
计算∂L/∂X_a = ∂L/∂Y_a @ W1_a^T（本地，需先经过GeLU反向）。此时每张卡得到的是∂L/∂X的部分贡献（因为Y是部分和，梯度也是部分）。需要将两张卡的∂L/∂X进行AllReduce求和，得到完整∂L/∂X。
同时计算∂L/∂W1_a = X^T @ ∂L/∂Y_a（本地）。

所以，前向和反向各需要一次AllReduce。TP的通信量正比于batch_size × seq_len × hidden_dim，与参数量无关。因此，当hidden_dim很大时，TP通信量很大，但参数量也大，TP节省的显存和计算量也显著。TP最适合节点内高带宽（NVLink）环境。

对自注意力头进行张量并行时，Q、K、V 投影如何切分？输出如何合并？¶

自注意力层的张量并行切分通常按注意力头维度进行，这比切分hidden_dim维度更自然，因为多头注意力本身就是按头独立计算然后拼接的。

切分方法：

将Q、K、V的投影权重矩阵沿列（输出维度）切分为TP份，每份包含num_heads/TP个注意力头。
每张卡独立计算自己负责的那些头的Q、K、V。
执行该部分头的自注意力计算（点积、softmax、加权），得到该部分头的输出。
输出投影矩阵W_O沿行切分，每张卡将本地头的输出投影，然后通过AllReduce求和，合并所有头的输出，得到完整的多头注意力输出。

优势：这种切分方式下，每张卡的计算和存储与头数成正比。注意力矩阵也被切分（每张卡只存自己那部分头的注意力矩阵），显存进一步降低。通信模式与MLP类似，前向和反向各需一次AllReduce。

TP 的通信量大小与什么有关？为什么 TP 适合单机多卡高带宽？¶

TP的通信量主要取决于激活值的大小，即 batch_size × seq_len × hidden_dim。每次AllReduce传输的数据量就是这个规模。与数据并行不同，TP的通信量与参数量无关，因为参数是静态切分的，不需要在每次迭代中传输。

为什么适合单机多卡高带宽：

TP的通信频率很高——每个Transformer层的前向和反向都需要至少两次AllReduce。在深层网络中，这意味着几十上百次的AllReduce。如果跨节点进行，即使InfiniBand带宽（如400GB/s）也远低于节点内NVLink（如900GB/s），通信延迟将严重拖慢训练。
TP的通信数据量较大（激活值），需要极大带宽来保证通信不成为瓶颈。单机内8卡A100通过NVSwitch互联，每对卡间带宽可达600GB/s，总带宽足够支撑TP=8的通信需求。
TP切分后的参数和激活显存都降低，使得单卡可以承载更大的单层。因此TP通常作为节点内并行策略，与跨节点的数据并行、流水线并行组合使用，形成3D并行。

解释 TP 中的“f 操作”和“g 操作”，分别对应什么通信？¶

在 Megatron-LM 的张量并行（Tensor Parallelism, TP）实现中，为了统一描述前向和反向传播中的通信模式，定义了f 操作和g 操作。这两个操作是构成张量并行的基本单元。

f 操作（前向传播中的恒等操作）：

前向传播时，f 操作就是不进行任何通信。输入张量直接复制到所有 TP 组的设备上，每张卡独立计算自己负责的部分输出。例如在 MLP 的第一层（升维层）按列切分时，输入 X 广播到所有卡，每张卡计算 Y_i = GeLU(X @ W1_i)，此时不需要通信，就属于 f 操作。
反向传播时，f 操作对应的反向是AllReduce。因为前向时数据被复制，反向时各卡计算的梯度是部分贡献（例如 ∂L/∂X 的部分和），必须通过 AllReduce 求和得到完整的梯度。

g 操作（前向传播中的 AllReduce 操作）：

前向传播时，g 操作是AllReduce。例如在 MLP 的第二层（降维层）按行切分时，每张卡计算 Z_i = Y_i @ W2_i，这些 Z_i 是部分和，需要 AllReduce 求和得到完整的输出 Z。这就是 g 操作。
反向传播时，g 操作对应的反向是恒等操作（不通信）。因为前向已经通过 AllReduce 得到了完整输出，反向时每张卡都持有完整的 ∂L/∂Z，可以直接本地计算梯度。

总结：f 操作和 g 操作是互为“镜像”的。前向的 f 对应反向的 AllReduce；前向的 g（AllReduce）对应反向的恒等。这种设计使得每层 Transformer 的前向只需要一次 AllReduce（在 MLP 的第二层和自注意力的输出投影后），反向也只需要一次 AllReduce（在 MLP 的第一层和自注意力的 QKV 投影后），通信次数最小化。

在自注意力层中：

Q、K、V 投影按列切分（f 操作），前向无通信；输出投影按行切分（g 操作），前向需要 AllReduce。
反向时，输出投影的梯度计算无需通信（因为前向 g 操作的反向是恒等）；Q、K、V 投影的梯度计算需要 AllReduce（因为前向 f 操作的反向是 AllReduce）。

通过这种统一的 f/g 抽象，Megatron 可以灵活组合 TP 的各种切分策略，并自动推导所需的通信。

为什么 TP 的通信成本很高？如果跨节点使用 TP 会有什么问题？¶

TP 的通信成本高，主要体现在两个方面：

通信数据量大：TP 的通信传输的是激活值及其梯度，数据量等于 batch_size × seq_len × hidden_dim（以及相应的 FP16 字节数）。当模型隐藏维度很大（如 8192 或更高）时，每次 AllReduce 传输的数据量可达数百 MB。而且，TP 的通信发生在每个 Transformer 层的内部，层数越多，通信次数越多（每层前向和反向各至少一次 AllReduce）。总通信量与模型层数、隐藏维度、序列长度和批次大小成正比。

通信频率高：在一个深层 Transformer（如 96 层）的训练中，每个迭代步需要进行 96 次前向 AllReduce 和 96 次反向 AllReduce（共计约 200 次）。这种高频通信对带宽和延迟都极其敏感。

跨节点使用 TP 的问题：

带宽瓶颈：节点内 GPU 通过 NVLink 或 NVSwitch 互联，带宽极高（如 A100 可达 600 GB/s 每对 GPU，总带宽达 4.8 TB/s）。而跨节点通信依赖 InfiniBand（IB）或以太网，即使是最新的 IB HDR（200 Gb/s ≈ 25 GB/s），单链路带宽仍远低于 NVLink。TP 的巨量通信会迅速占满跨节点网络，导致通信时间远超计算时间，GPU 大量空闲。
延迟增加：跨节点通信需要经过网卡、交换机等多个环节，延迟远高于节点内的 GPU 直连。高频的 AllReduce 操作会累积显著的延迟开销。
扩展性差：随着 TP 并行度增加，通信量线性增长，跨节点网络的瓶颈效应更加突出。

因此，TP 通常仅用于单机内部（通常 TP 并行度 ≤ 节点内 GPU 数，如 8），利用 NVLink 提供的高带宽支撑其密集通信需求。跨节点时，更倾向于使用数据并行、流水线并行或 ZeRO-3 等通信量较低或频率较低的策略。

流水线并行 (PP) 是如何划分模型的？常见的划分策略（按层均分或根据计算量）。¶

流水线并行（Pipeline Parallelism, PP）将模型的不同层（或层组）分配到不同的 GPU 上。每张 GPU 负责其中一部分层的计算，层与层之间通过点对点通信传递激活值和梯度。

划分策略：

按层数均分：最直观的策略，将模型的总层数 L 均匀分配给 P 个设备，每个设备负责 L/P 层。优点是实现简单，层数均分容易。但当不同层的计算量或参数量差异较大时，各设备的计算负载可能不均衡，导致部分设备闲置等待（流水线气泡）。
根据计算量或显存均衡划分：考虑到 Transformer 的首尾层可能包含嵌入层（参数量大但计算简单）或 LM Head（参数量大），以及不同层的激活值大小可能不同，更优的划分是根据每一层的计算时间或显存占用来分配。例如，将嵌入层放在第一个设备上，并将较少的 Transformer 层分配给它，使其与其他设备的计算时间均衡。Megatron-LM 支持自动搜索较优的层分配方案以最小化显存峰值。
手动调优：在实践中，工程师会根据模型特点和硬件拓扑，手动调整每个 PP 阶段的层数和放置。例如，将需要大显存的层（如最后的 LM Head）单独放在一个设备上，并利用其额外的显存来容纳更大的激活值。

什么是“微批次” (Micro-batch)？在 PP 中为什么需要？¶

微批次（micro-batch）是将一个全局批次（global batch）进一步拆分成更小的、流水线可以并行处理的数据块。假设全局批次大小为 64 个样本，可以拆分为 8 个微批次，每个微批次 8 个样本。

PP 中需要微批次的原因：在标准的流水线中，如果不拆分微批次，第一张 GPU 处理完整个 batch 后，将结果传给下一张 GPU，此时第一张 GPU 完全空闲，直到整个 batch 流经所有设备。这造成了巨大的“流水线气泡”（设备空闲）。通过将大 batch 拆分为多个微批次，各设备可以在处理完上一个微批次后立即处理下一个微批次，不同设备同时处理不同微批次的不同阶段，实现时间上的重叠。这显著减少了空闲时间，提高了吞吐量。

微批次的大小受限于单卡的显存（因为需要存储该微批次产生的激活值），因此需要在显存允许的前提下尽量增大微批次大小，或者通过增加微批次数量（即更细的切分）来提高流水线利用率。

画出 GPipe 的调度图，并计算流水线气泡率。¶

GPipe 是最简单的同步流水线并行策略。假设有 4 个设备（P=4），每个设备负责模型的一部分。将每个 batch 切分为 M 个微批次（这里 M=4，微批次 F1~F4 表示前向，B1~B4 表示对应的反向）。

GPipe 调度图：

时间 →
Device 0: F1 F2 F3 F4               B1 B2 B3 B4
Device 1:    F1 F2 F3 F4            B1 B2 B3 B4
Device 2:       F1 F2 F3 F4         B1 B2 B3 B4
Device 3:          F1 F2 F3 F4      B1 B2 B3 B4

GPipe 先完成所有微批次的前向传播，再开始反向传播。反向传播的顺序与前向相反（从最后一个微批次开始）。

气泡时间：气泡是图中设备空闲的时间。对于 GPipe，气泡出现在前向传播和反向传播的“两端”。总的时间步数为 M（前向）+ M（反向）- 1（最后一个前向与第一个反向重叠？）实际上 GPipe 是严格顺序的，每个设备执行 M 个前向，然后 M 个反向。设备 0 空闲了 M-1 个前向时间，设备 P-1 空闲了 M-1 个反向时间。总气泡数为 (P-1) × M。每个微批次前向或反向花费的时间设为 1 单位，总执行步数为 P+M-1（？）更准确的 GPipe 气泡率公式为：

气泡率 = (P-1) / (P-1+M) ——当反向计算时间等于前向时间时。

推导：总时间 = M * P（前向总步数，但流水线执行时间更短）。在 GPipe 调度下，第一个设备完成所有 M 个前向需要 M 时间，然后开始反向，反向也要 M 时间，但最后一个设备在第一个设备开始反向时可能还在执行前向。总体执行时间 = (P-1) + M（前向流水线填充） + (P-1) + M（反向流水线）？更标准的计算是：有效工作时间 = M * P（所有设备的总前向+反向时间），实际耗时 = P + 2(M-1)（？）实际上气泡率直观理解为：P 个设备，M 个微批次，空闲的总时间占比。理想无气泡时总工作量为 P2M（前向+反向），GPipe 的空闲时间为 (P-1)2M？这是错的。

我们重新推导：设每个微批次的前向时间为 1，反向时间也为 1。在 GPipe 中，最后一个设备开始第一个微批次的前向需要等待 (P-1) 时间，所有前向完成需要 (P-1+M) 时间。然后开始反向，反向最后一个微批次在最后一个设备完成需要再 (P-1+M) 时间。因此总时间 T_gpipe = 2*(P-1+M)。理想无气泡时，所有设备连续工作，总时间 = (2M)。气泡时间 = T_gpipe - 2M = 2(P-1)。因此气泡率 = 2(P-1) / 2(P-1+M) = (P-1)/(P-1+M)。例如 P=4,M=4，气泡率 = 3/7 ≈ 43%。

比较 GPipe 和 PipeDream（1F1B 调度）的区别。1F1B 如何减少显存峰值？¶

GPipe：

先完成所有微批次的前向，再开始反向。
需要存储所有微批次的前向激活值，直到反向传播。显存峰值极高（与微批次数量 M 成正比）。
实现简单，但显存不友好，通常需要配合 Gradient Checkpointing。

PipeDream（1F1B，One Forward One Backward）调度：

在预热阶段执行若干前向传播，使得流水线填充。然后进入稳定阶段，每个设备交替执行一个前向和一个反向（1F1B）。
显存峰值大幅降低：因为较早完成的微批次的前向激活值可以尽早被反向消费并释放。显存占用大致与流水线深度（P）相关，而不是微批次数量 M。这允许使用更大的 M 来减少气泡而不导致 OOM。
实现更复杂，但已成为主流 PP 调度算法。

1F1B 减少显存峰值的原理：反向传播可以及时释放激活值。通过限制同时“存活”的微批次数量（大致为流水线深度 P），而非 GPipe 的 M（M>>P），显存峰值控制在一个较低水平。

说明 1F1B 调度中的“预热阶段”和“稳定阶段”。¶

预热阶段：从第一个设备开始，每个设备依次执行其前向传播。这个阶段的目标是让流水线“充满”，使得最后设备也能收到微批次。在这个阶段，设备执行的前向操作多于反向操作（或仅执行前向）。预热阶段执行的微批次数量等于 min(P-1, M)，通常每个设备执行一定数量的前向，直到流水线建立起稳定的批次流。

稳定阶段：预热完成后，流水线进入稳态。每个设备按照 1F1B 模式工作：完成一个前向传播后，紧接着执行一个反向传播（可能是另一个微批次的反向）。理想情况下，每个设备每步交替进行前向和反向，保持流水线持续运转。此时系统中同时存在多个微批次：有些正在前向，有些正在反向。稳定阶段一直持续到所有微批次的前向都已完成。

收尾阶段：当不再有新的前向微批次时，设备逐渐排空流水线中的反向任务，此阶段称为“冷却”。最终完成所有反向。

推导流水线气泡时间公式：Bubble = (P-1) / (P-1+M) ，并解释。¶

我们针对 1F1B 调度推导气泡率。假设每个微批次的前向计算时间为 1，反向计算时间也为 1（通常反向略长，但为简化设为相等）。总设备数为 P，微批次数量为 M。

理想情况：所有设备 100% 利用率，总有效工作时间为 P * 2M（P 个设备，每个完成 M 个前向和 M 个反向）。完成所有工作的时间应为 2M。

1F1B 调度的实际耗时：在 1F1B 中，流水线的填充和排空需要额外时间。总执行时间 T = (P-1) + M (前向填充) + (P-1) + M (反向排空) ? 更准确地说，1F1B 的调度时间可以计算为：第一个微批次从第一个设备到最后一个设备需要 P 步（前向传递），然后反向回来也需要 P 步。由于微批次流水线执行，最后一个微批次的反向完成时间等于：

前向完成时间：最后一个微批次在第一个设备上开始的时间是 M-1，到达最后一个设备是 (M-1) + (P-1) = M+P-2。
反向完成时间：从最后一个设备开始反向，第一个微批次的反向与最后一个微批次的前向有时间重叠。1F1B 的总时间公式通常为：T = (P-1) + (M) + (P-1) + (M) - 1？实际推导可得 T = M + 2(P-1) 当 M > P 时？经过更精确的 1F1B 调度分析，总时间 = M + 2(P-1) 当 M 足够大时（稳定阶段）。总计算量是 2P*M，气泡时间 = T - 2M = 2(P-1)。气泡率 = 2(P-1) / (M + 2(P-1))。但是通常给出的简化气泡率公式为 (P-1) / (P-1 + M)，这是假设前向和反向混合在一起的时间归一化结果，考虑了流水线深度与微批次比例。

另一种常见推导：在 1F1B 中，设备空闲只发生在流水线的填充和排空阶段。填充阶段需要 (P-1) 步（每个设备等待第一个微批次到达），排空阶段同样需要 (P-1) 步（最后一个微批次反向传递）。总时间 = 2M (理想) + (P-1)（前向填充）+ (P-1)（反向排空）- 一些重叠？实际上 T = 2M + 2(P-1) - 1 ？我们采用广泛引用的结果：对于 1F1B，气泡率 ≈ (P-1)/(P-1+M)。这意味着：当 M 远大于 P 时，气泡率趋近于 0；当 M 较小时，气泡显著。这符合直觉：微批次越多，填充和排空的相对开销越低。

推导：设总有效工作量为 W = P * 2M。实际总耗时为 T = 2M + 2(P-1)。空闲总时间为 2(P-1)。因此空闲时间占比（气泡率）= 2(P-1) / (2M + 2(P-1)) = (P-1) / (P-1+M)。如果考虑反向时间比前向长，公式会有调整，但核心原理不变。

如何降低流水线气泡？增加 M 的代价是什么？¶

降低气泡最直接的方法是增加微批次数量 M。从气泡公式可见，M 越大，分母 (P-1+M) 越大，气泡率越低。但 M 不能无限增加，代价是：

显存占用增加：因为需要维护更多“在途”的微批次及其激活值（即使使用 1F1B，峰值激活值也与流水线深度 P 和每个微批次的大小有关，但通常主要受 micro batch size 和 P 影响，M 增大对峰值影响较小，但需存储更多微批次的输入和输出状态）。实际上，增加 M 通常意味着每个 micro batch 更小，所以总激活值与 batch size 成正比。因此，显存预算限制了 micro batch 的大小下限（不能太小否则 GPU 利用率低），从而限制了 M 的上限。
通信和调度开销：更细粒度的微批次增加了调度频率和点对点通信的启动次数，可能降低效率。
训练稳定性：过小的 micro batch 可能导致梯度噪声过大，影响收敛。通常需要配合梯度累积来维持合理的全局批次大小。

其他降气泡方法：使用交错调度 (Interleaved Pipeline)，让每个设备负责多个不连续的层段，进一步减少气泡。

什么是“交错调度” (Interleaved Pipeline)？如何减少气泡？¶

传统的流水线并行中，每个设备被分配连续的层块（如设备0负责层0-3，设备1负责层4-7等）。交错调度则将模型划分为更多的微阶段（micro-stages），每个设备负责多个不连续的微阶段。例如，一个 16 层的模型，P=4，传统方法每个设备连续 4 层。交错调度可能让设备0负责层0和层8，设备1负责层1和层9，以此类推，使得流水线中不同微批次的处理在设备间更密集地交叉。

减少气泡的原理：交错调度使得每个设备的多个计算块交错排列，填充了传统调度中因依赖关系产生的空闲时间。气泡率公式中的有效 P 变成了 P * v（v 是每个设备的微阶段数），而 M 相应地需要更大。气泡率近似为 (P-1) / (P-1 + v*M)？实际上交错调度可以将气泡率降低到 (P-1) / (P-1 + v*M) 其中 v 是每个设备的虚拟阶段数。这意味着气泡显著减小。

Megatron-LM 支持交错流水线并行，通常设置 num_micro_batches 和 num_model_chunks（每个设备的模型块数）。它要求 M 能被 P 整除等约束，但能极大提升流水线利用率。

序列并行 (SP) 的动机是什么？主要用在什么场景？¶

动机：在训练长序列（如几千甚至几万 tokens）时，激活值显存成为主要瓶颈。自注意力的激活值（尤其是注意力矩阵）随序列长度 L 呈二次增长。即使使用了 Gradient Checkpointing 和 FlashAttention，当 L 极大时，单卡仍然无法容纳单层的激活值。序列并行将序列维度切分到多个 GPU 上，每张卡只存储和处理一部分序列的激活，从而降低单卡显存占用。

应用场景：

训练极长上下文模型（如 8K、16K、32K tokens）。
与张量并行（TP）结合使用，因为 TP 已经切分了隐藏维度，进一步切分序列维度可以叠加节省显存。Megatron-LM 中，TP 和 SP 通常一起开启。

Megatron 中的序列并行是如何实现的？它对 LayerNorm 和 Attention 的处理。¶

Megatron-LM 的序列并行与张量并行结合使用。假设已开启 TP（切分隐藏维度），在此基础上启用 SP，将激活张量的序列维度 seq_len 进一步切分到同一 TP 组内的设备上。

具体实现：

LayerNorm 和 Dropout：这些操作在序列维度上是独立的（对每个 token 独立归一化）。在 SP 模式下，每个设备只拥有序列的一部分 tokens，因此可以独立执行 LayerNorm 和 Dropout，无需通信。这减少了冗余计算。
自注意力：注意力计算需要完整的序列信息（因为每个 token 要与所有 token 交互）。在 SP 中，每个设备只持有部分 token 的 Q，但需要完整的 K 和 V 来计算注意力。因此，在计算 QKV 投影后，需要AllGather 操作收集完整的 K 和 V（在序列维度上拼接），然后计算注意力。注意力输出同样在序列维度上被切分回各设备（通过 ReduceScatter 或类似操作）。更高效的实现（如 Ring Attention）则通过环形通信逐步传输 KV 块，避免一次性 AllGather 的大显存。
MLP 和残差连接：这些也是 token 独立操作，无需通信，各设备独立计算。

通信模式：SP 在自注意力层引入额外的 AllGather 和 ReduceScatter（或等效的 Ring 通信），增加了通信量。但由于 SP 通常与 TP 结合，且长序列时注意力计算成为瓶颈，通信开销往往可以被隐藏。通过 SP，激活显存进一步降低，使得训练超长序列成为可能。

SP 与 TP 组合时，通信是如何优化的？减少冗余通信。¶

序列并行（SP）和张量并行（TP）组合使用时，核心优化目标是最大化激活值显存节省，同时最小化额外通信。两者结合可以叠加降低单卡显存：TP切分隐藏维度，SP切分序列维度，激活值在两个维度上同时被分片。

通信优化手段：

共享通信组：在Megatron-LM中，SP和TP使用相同的进程组。这意味着不需要为SP单独建立通信域，SP的通信操作（如AllGather、ReduceScatter）与TP的通信操作在同一个NCCL组内执行，避免了额外的组初始化和上下文切换开销。
LayerNorm和MLP等token独立操作：这些操作在序列维度上是独立的（每个token单独归一化）。在SP模式下，每个设备只持有一部分序列的token，因此可以本地独立执行LayerNorm、Dropout和MLP，不需要任何跨设备通信。相比没有SP时，TP组内每张卡都需要处理完整序列的这些操作，SP让计算量也得到分摊，且没有通信开销。
自注意力层的通信融合：
在SP中，每个设备只拥有部分token的Q，但需要完整的K和V来计算注意力分数。标准的做法是：在计算QKV投影后，执行一次AllGather将各设备的K和V在序列维度上拼接，得到完整K和V；计算完注意力后，通过ReduceScatter将输出在序列维度上切分回各设备。
由于TP已经对隐藏维度进行了切分，SP的AllGather/ReduceScatter的通信数据量是 batch_size × (seq_len/TP) × hidden_dim，实际上是按TP切分后的维度计算的，因此通信量相比纯SP有所减少。
进一步优化：可以使用Ring Attention代替AllGather，通过环形通信逐步传输KV块，避免一次性收集完整序列带来的显存峰值和通信突发，实现通信与计算的重叠。
消除冗余通信：如果不做优化，SP和TP可能会在不同阶段重复进行类似的AllReduce（如MLP的AllReduce和SP的AllGather可能各自独立）。通过统一通信调度，Megatron将SP的通信嵌入到TP的f/g操作框架中，使得每层总通信次数和通信量经过精心编排，避免重复。

通信模式对比表：

查看内嵌表格

什么是“混合并行”？DP、TP、PP 三者的组合原则。¶

混合并行是将数据并行（DP）、张量并行（TP）和流水线并行（PP）组合在一起，以训练单卡无法容纳的巨型模型。三种并行策略各有侧重，组合的核心原则是扬长避短，逐级扩展。

三者特点对比表：

查看内嵌表格

组合原则：

TP优先节点内：TP的通信发生在每个Transformer层内部，频率极高（每层2次AllReduce），数据量也大。节点内NVLink带宽可达900 GB/s，远高于跨节点IB网络（通常25-50 GB/s每卡）。因此TP几乎总是被限制在单个节点内部，并行度不超过单节点GPU数（如8）。
PP跨节点，但注意气泡：PP的通信是层间的点对点传输，频率低（每个微批次每层边界一次），通信量也较小。PP可以跨节点，但受限于流水线深度和气泡率。通常PP的并行度不宜过大（10~20），以减少气泡开销。
DP作为全局扩展维度：DP的通信是梯度AllReduce，发生在整个训练步结束时，可以利用通信与计算重叠（DDP）。DP可以跨大量节点扩展，且与ZeRO结合可显著降低单卡显存冗余。在3D并行中，DP通常是最大并行维度，负责扩展到数百甚至数千卡。
层级叠加顺序：通常先设定TP（如8卡节点内），再设定PP（如16个流水线阶段），最后将剩余的卡全部用于DP。三个维度的乘积应等于总GPU数量。

组合实例分析：

假设有1024张A100 GPU，单机8卡，共128台节点。

TP=8（占用1个节点，所有节点内部采用相同的TP切分）
PP=16（需要16个TP组串联成流水线，占用16个节点）
DP=1024/(8×16)=8（8个这样的16节点流水线组做数据并行）

这样，DP=8, TP=8, PP=16，实现了3D并行。

对于一个 175B GPT-3 类模型，推荐哪种混合并行方案？为什么？¶

GPT-3 175B的隐藏维度为12288，层数96，头数96。这个规模单卡完全无法容纳，必须使用混合并行。

推荐方案：TP=8, PP=16, DP=64（假设使用8卡节点，总计1024张卡）。或者 TP=8, PP=8, DP=128（总计1024卡）。

原因分析：

TP=8：隐藏维度12288，TP切分后每卡负责的隐藏维度为1536，QKV投影矩阵大小合理。TP限制在单节点内，利用NVLink高带宽。若TP=4也可行（每卡隐藏维度3072），但8能进一步降低单卡显存。通常TP设为节点内卡数（8），最大化利用节点内高速通信。
PP=16：96层模型，若按层均分，每个流水线阶段约6层。PP=16使得每个设备层数较少，减少单卡激活值存储（因为每张卡只需为它负责的层保存激活）。PP并行度不宜过大，否则气泡率增高。16是合理的折中。
DP=64：总卡数 = TP×PP×DP = 8×16×DP。若集群为1024卡，则DP=8。若集群更大，DP可相应增加。数据并行是吞吐扩展的主要手段，DP维度最大。
为什么不是其他配置：
若TP=2, PP=48：PP并行度过高，气泡大，且48个阶段可能导致负载不均。
若只用ZeRO-3而不使用TP/PP：ZeRO-3可以训练175B模型，但ZeRO-3需要频繁AllGather参数，通信压力大，不如3D并行通信效率高，尤其在超大规模集群上，混合并行通常能获得更高的MFU。

实际GPT-3训练（据公开信息）使用了类似Megatron的3D并行，具体配置未完全公开，但上述方案是工业界训练175B量级模型的典型配置。

3D 并行的通信开销如何分解？哪部分最大？¶

3D并行的总通信开销可以分解为DP、TP、PP三部分，它们各自的特点决定了各自的通信瓶颈程度。

通信开销分解表：

查看内嵌表格

哪部分最大？

TP通信通常是瓶颈：TP的通信频率极高（96层模型，每层2次，共192次AllReduce），每次传输的数据量也很大（batch × seq_len × hidden_dim）。尽管在节点内利用NVLink，但如此高频的通信仍然可能占用大量带宽，成为训练时间的主要开销。例如，在175B模型训练中，TP的AllReduce时间可能占总训练时间的20%~30%。
DP通信次之：DP的AllReduce发生在整个反向传播期间，但通过“梯度桶”机制和异步通信，可以将其隐藏在计算时间中。在大规模DP下（如跨节点），网络带宽可能成为瓶颈，但可以通过梯度压缩、通信拓扑优化等手段缓解。
PP通信最小：PP只传输层间激活值，通信频率低，且可以通过增加微批次数量来摊销通信开销。PP的通信通常不是瓶颈。

结论：TP通信开销最大，且难以隐藏，因此TP应严格限制在节点内高带宽环境。DP通信次之，但可优化余地大。PP通信开销几乎可忽略。

拓扑感知的并行策略：如何在多机环境中分配 PP 和 TP 以最小化跨机通信？¶

多机环境中，跨节点的通信带宽远低于节点内部。因此，拓扑感知的并行策略旨在将通信密集的并行组放置在节点内部，将通信稀疏的组放在跨节点链路中。

具体原则：

TP 严格节点内：TP组的成员必须在同一物理节点内，利用NVLink高速互联。通常TP并行度等于节点内GPU数（如8），这样TP通信完全在节点内解决，不产生跨节点流量。
PP 尽量相邻节点：流水线并行的通信模式是“相邻阶段”的点对点通信。如果将连续的流水线阶段分配到物理上相邻的节点（如同一交换机下的节点），可以减少跨交换机跳数，降低延迟。例如，PP的阶段0和阶段1放在同一机架的相邻节点，它们之间的通信只经过一个交换机，而不跨越整个网络。
DP 跨节点自由分配：数据并行的通信模式是全局AllReduce，所有参与DP的卡都需要交换梯度。拓扑感知的DP可以利用环形AllReduce或分层AllReduce（节点内Reduce，节点间AllReduce）来优化。但DP的卡可以分布在任意节点，因为其通信模式是全连接的。

最小化跨机通信的策略：

为每个流水线阶段分配一个完整的节点，节点内部的TP组处理该阶段的层。这样，节点内部是TP通信，节点之间是PP通信。
如果PP阶段数超过节点数，可以将多个连续阶段放在同一节点内（共享节点内的GPU），它们之间的PP通信变为节点内传输。
DP的AllReduce通过NCCL的跨节点通信进行，已经针对拓扑进行了优化（如使用RDMA和GPUDirect）。更激进地，可以配置NCCL使用多环拓扑来增加跨节点带宽利用率。

解释 PyTorch 中的 DistributedDataParallel 如何与模型并行（如 Megatron）结合。¶

在PyTorch中，DistributedDataParallel (DDP) 负责数据并行，而模型并行（TP/PP）由上层框架（如Megatron-LM）实现。两者结合的关键在于层次化通信组管理。

实现方式：

通信组划分：Megatron-LM首先根据并行配置创建多个进程组：
TP组：参与张量并行的设备列表。
PP组：参与流水线并行的设备列表（通常与TP组交叉，形成网格）。
DP组：剩余设备用于数据并行。一个DP组内包含多个模型副本，每个副本是一个完整的TP+PP组合。
模型构建与包裹：
在每个设备上，根据其所在的TP和PP组，构建它负责的那部分模型层。
将构建好的模型实例用DDP包裹。DDP会自动识别DP组内的所有设备，并在反向传播时进行梯度AllReduce。
梯度同步：
在每个训练步中，TP和PP内部的通信由Megatron-LM控制（f/g操作和层间Send/Recv）。
当反向传播完成，DDP层拦截梯度，并启动DP组内的梯度AllReduce。
对于使用ZeRO的情况，DeepSpeed会接管优化器状态分片等操作，但DDP的基本梯度同步逻辑类似。

关键技术点：

DDP的Reducer会识别出模型中的哪些参数需要全局梯度同步（即那些在DP组内所有副本中都存在的参数）。在模型并行中，部分参数可能只存在于TP组内（如切分后的权重），DDP只对完整的模型参数的梯度进行AllReduce（这些参数在每个DP副本中都是完整的，尽管在TP组内是切分的，但DDP层面看到的是完整的逻辑参数）。
Megatron-LM和DeepSpeed都实现了与DDP的无缝集成，用户在调用deepspeed.initialize()或Megatron的并行初始化时，内部会自动处理通信组的创建和DDP的包裹。

什么是“弹性训练” (Elastic Training)？动态调整并行 worker 数量。¶

弹性训练是指在训练过程中，可以动态地增加或减少参与训练的GPU数量，而无需从头开始重启训练。这对于大规模集群中的容错（节点故障）和资源动态分配（抢占式实例、集群共享）至关重要。

核心机制：

动态进程组管理：当有节点加入或离开时，训练框架自动检测变化，并重新初始化分布式通信组（如NCCL组）。所有剩余的worker协商新的全局rank和world size。
状态重新分布：模型参数、优化器状态和数据加载器的状态需要根据新的并行度重新分布。例如，数据并行的批次大小需要根据新的worker数量调整；如果使用ZeRO，参数分片需要重新平衡。
无缝恢复：当节点故障时，弹性训练可以从最新的检查点恢复，并在剩余的节点上继续训练。当新节点加入时，可以从检查点重新分配状态，实现动态扩展。

实现框架：

PyTorch Elastic (torchrun)：提供torchrun命令行工具，支持动态worker管理。通过etcd或c10d存储协调信息，在worker数量变化时自动重启进程。
DeepSpeed Elasticity：在DeepSpeed中，通过deepspeed.elasticity模块实现，支持ZeRO优化的弹性调整。
Kubernetes + Volcano：在云原生环境下，利用K8s的Pod自动伸缩和调度器实现弹性训练。

应用场景：使用云厂商的竞价实例（Spot Instance）降低成本时，实例可能随时被回收，弹性训练可以自动迁移到新实例继续训练；在多租户集群中，动态分配空闲GPU资源。

为什么数据并行在大规模训练中是基础？可否只用 PP 和 TP 不用 DP？¶

数据并行是基础的原因：

吞吐量扩展：数据并行通过增加GPU数量，线性地提高每个训练步处理的样本数（全局批次大小），从而提升整体训练吞吐量。这是最直接、最高效的扩展方式。
通信效率：DP的梯度通信可以通过“梯度桶”和通信计算重叠隐藏延迟，且通信量与参数量成正比，不随序列长度或隐藏维度爆炸。相比TP的高频激活通信，DP的通信压力更可控。
成熟优化：DP有大量的优化技术（如ZeRO、梯度压缩、异步通信），能够有效降低显存和通信开销。

能否只用PP和TP不用DP？

理论上可以，但极端低效。只用PP和TP意味着只有一个模型副本（DP=1）。此时虽然模型可以被切分到多卡，但每张卡每次只能处理一个微批次的数据。为了增加吞吐，必须增大微批次大小，但这受到单卡显存限制。而且，PP和TP都无法像DP那样通过增加副本来线性提升吞吐。

实际后果：

训练速度极慢：即使你用1000张卡做TP=8, PP=125，但只有一个数据流经过，设备利用率极低，大部分时间在等待流水线填充。
无法有效利用大规模集群：为了充分利用千卡集群，必须创建多个数据并行副本。DP提供了水平扩展（scale out）的能力，而TP/PP提供垂直扩展（scale up）以容纳单模型大小。

因此，DP是吞吐扩展的基础，TP/PP是容纳大模型的必要手段。三者结合才能同时实现“大模型”和“高吞吐”。

张量并行的切分粒度可以到行或列，还可以更细吗？¶

张量并行的切分粒度可以更细，除了按行和按列切分矩阵外，常见的还有：

按注意力头切分：将不同的注意力头完全分配到不同的设备上。这是Megatron-LM中对自注意力层的默认切分方式。每个设备负责num_heads / TP个完整头的计算，包括Q、K、V投影和注意力计算。这种方式实现简单，通信量也较小（仅需在输出投影时AllReduce）。
按块切分 (Block-wise)：将矩阵划分为多个二维块，分布到不同的设备上。这种更细粒度的切分可以改善负载均衡，但通信模式会变复杂（可能需要多次AllGather和ReduceScatter）。通常在实际中不常用，因为按列/按行切分已经足够高效。
混合切分：同时结合列切分和行切分，例如对权重矩阵先分块，再在每个块内进行列/行切分。

更细粒度的代价：更细的切分会导致更多的通信启动次数和更小的通信块，降低带宽利用率。因此，当前主流的TP实现（Megatron）选择了按列切分+按行切分（对MLP）和按头切分（对注意力）的组合，这已经被证明在通信和计算效率之间取得了良好平衡。

限制因素：切分粒度也受限于矩阵维度的整除性。例如，隐藏维度必须能被TP整除；注意力头数也必须能被TP整除。如果TP度与模型维度不匹配，则需要填充或更复杂的切分策略。

为什么 TP 通常要求卡数为 2 的幂或特定数？(如 8 卡)¶

原因一：模型维度的整除性。张量并行需要将隐藏维度、注意力头数等切分到各设备。模型设计时，这些维度通常取2的幂（如512, 1024, 2048）或与2的幂密切相关。为了不引入额外的填充（padding），TP并行度也取2的幂，能自然整除这些维度。例如，隐藏维度4096，TP=8整除；若TP=6则不整除，需要特殊处理。

原因二：NCCL通信效率。NCCL的集合通信算法（如Ring AllReduce、Tree AllReduce）在2的幂规模下经过高度优化，能够达到最优的带宽利用率和最低的延迟。在非2的幂规模下，通信拓扑可能非对称，效率下降。

原因三：硬件拓扑对齐。单节点内GPU通常为8卡（A100/H100服务器），且内部通过NVSwitch全互联。TP=8刚好占满一个节点，TP通信完全在节点内解决，避免了跨节点TP。如果TP=4，则一个节点内可以放2个TP组。

但并非绝对：TP=2, 4也是常见配置。对于某些模型，如果注意力头数不是2的幂（如头数=48），TP=6也可行，只要整除即可。但出于效率考虑，主流大模型和框架默认推荐2的幂。

在 TP 中，如果注意力头数为 32，TP=8，每张卡几个头？¶

32头 ÷ 8 = 4个头。每张卡负责4个注意力头的完整计算。

具体地，Q、K、V的投影权重按列切分，每张卡获得一部分列，对应4个头的输出维度。每张卡独立计算这4个头的自注意力，然后输出投影权重按行切分，每张卡计算部分输出后通过AllReduce求和合并。

流水线并行中，如何确定每个阶段 (Stage) 的层数？负载均衡如何做？¶

目标：使各流水线阶段的总计算时间尽可能相等，以减少气泡。

方法：

Profiling：在单卡上运行完整的模型前向和反向，精确测量每一层的计算耗时（FLOPs或实际时间）。同时考虑不同层之间的激活值大小差异可能带来的通信时间。
按计算量均分：将模型所有层的计算时间累加，除以流水线阶段数P，得到每个阶段的目标计算时间。然后通过动态规划或贪心算法寻找分割点，使得每个阶段的总计算时间尽可能接近该目标值。动态规划可以精确求解最小化最大阶段时间问题（类似于“划分数组为P份使最大和最小”）。
考虑首尾特殊层：嵌入层和输出层通常参数量大但计算模式不同。可以将它们单独分配给第一个和最后一个阶段，并相应调整其分配的Transformer层数量。
显存约束：某些层可能需要更多显存（如注意力矩阵大的层），可能导致OOM。因此，划分时还需考虑每阶段的峰值显存，确保不超过设备显存。
交错调度下的负载均衡：如果使用交错流水线，每个设备负责多个不连续的模型块，负载均衡可以在更细粒度进行。

Megatron-LM的实现：提供了--num_layers_per_virtual_pipeline_stage等参数，允许用户手动指定层分配，或者使用自动搜索脚本找到最优划分。

如果各层计算量不同，如何通过“动态规划”进行非均匀切分以减少气泡？¶

问题建模：假设模型共有L层，第i层的计算时间为c[i]。需要将模型划分到P个流水线阶段，每个阶段负责连续若干层，使得所有阶段中最大计算时间最小化。这等价于经典的“分割数组为P段，最小化最大段和”问题。

动态规划解法：

定义dp[i][k]为前i层划分为k段的最小化最大段和。
状态转移：dp[i][k] = min_{j < i} { max(dp[j][k-1], sum(c[j+1..i])) }
时间复杂度O(L²P)，对于L=96, P=16完全可行。

求解后：得到每段的层范围，按照该划分进行PP分配。由于每段计算量均衡，流水线气泡主要由P本身引起，而非额外的不均衡气泡。

实践中的近似：由于精确测量每层计算时间有误差，且通信时间也会影响，通常采用贪心：从输入层开始，累加计算量，当累加值达到目标平均值附近时，切分一个阶段。这种方法简单，也能获得很好的负载均衡效果。

通过动态规划或贪心进行非均匀切分，是减少流水线气泡、提高训练效率的关键步骤，尤其在模型层数不是PP并行度整数倍时尤为重要。