NeMo2 并行

NeMo2 代表扩展 pytorch-lightning 功能的工具和实用程序，以支持使用 megatron 模型进行训练和推理。虽然 pytorch-lightning 支持足够的并行抽象以用于适合单个 GPU 的 LLM（分布式数据并行，又名 DDP），甚至可以用于需要在 GPU 小集群之间分片的稍大架构（完全分片数据并行，又名 FSDP），但当您处理非常大的架构并希望获得尽可能高效的预训练和推理时，megatron 支持的并行性是一个很好的选择。

因此，换句话说，NeMo2 除了标准的 DDP 和 FSDP 策略外，还添加了 Megatron 策略。

许多下游约束和约定都受到 megatron 底层约束的驱动。

Megatron 的更深层背景知识

并行化较小模型的其他选项

Megatron 是一个用于支持高级模型并行变体的系统。虽然可以使用分布式数据并行 (DDP) 等系统并行执行普通模型，或者可以使用 Meta 的完全分片数据并行 (FSDP/FSDP2) 训练中等大小的模型，但当您处理非常大的模型并且希望以最大效率训练它们时，您需要 megatron 的某种变体。

DDP 背景知识

当您可以将整个模型都放入集群中的每个 GPU 中时，DDP 是最佳选择。 使用 DDP，您可以通过将 global batch 拆分为更小的 mini-batches（每个 GPU 一个）来跨多个 GPU 并行化您的 global batch。每个 GPU 独立计算其数据子集的前向和后向传递，从而实现最大化利用率。梯度同步在每个批次的后向传递完成后进行，然后进行权重更新，以确保所有 GPU 在下一次迭代中都具有同步参数。这是一个如何在带有小型模型的集群上显示此情况的示例

FSDP 背景知识

FSDP 通过在集群中的 GPU 之间分片（拆分）模型权重来扩展 DDP，以优化内存使用。虽然数据仍然像 DDP 一样在 GPU 之间拆分，但 FSDP 会策略性地同步和广播模型权重的必要分片，以便在正向传递期间及时进行计算。

例如，当计算需要某个层时，拥有该分片的 GPU 会将该权重分片发送到其他 GPU，然后这些 GPU 对该层执行前向计算。计算完成后，FSDP 会释放除拥有该分片的 GPU 之外的所有 GPU 上该层的内存。此过程对每一层迭代进行，直到整个模型都在数据上执行完毕。

请注意，此过程以一种能够执行过大模型的方式并行化存储（假设单层不会太大而无法放入 GPU）。Megatron（下一步）共同定位存储和计算。

以下两张图显示了通过 FSDP 分片的模型的前向传递的两个步骤。

模型并行

模型并行是适用于跨集群并行化模型的各种不同并行策略的总称。下面我们解释 megatron 中实现的几种模型并行变体。如前一节所述，megatron 特定的并行类型（接下来描述）的一个关键优势在于它们共同定位了层的存储和计算。由朴素调度器实现引起的效率低下问题也得到了解决（在关于调度器的部分中讨论）。

流水线并行

流水线并行类似于 FSDP，但分片的模型块也在拥有相关模型权重的节点上并行计算。您可以将其视为碰巧分布在多个子 GPU 上的更大的模拟 GPU。 例如，这包括 parallel_state.is_pipeline_last_stage()，它通常用于判断特定节点是否位于最后一个流水线阶段，您可以在其中计算最终 head 输出、损失等。。 类似地，第一个流水线阶段（例如，您在其中计算嵌入）也有方便的环境查找 parallel_state.is_pipeline_first_stage()。

张量并行

张量并行表示跨 GPU 拆分单层。 这也可以解决某些单个层在理论上可能太大而无法放入单个 GPU 的问题，这将使 FSDP 无法实现。但这仍然有效，因为单个层权重（和计算）是分布式的。 megatron 中的示例包括 RowParallelLinear 和 ColumnParallelLinear 层。

序列并行

在 megatron 中，“序列并行”是指 transformer 的 dropout 和 layernorm 块的并行化。 这个想法大致如下。 首先，请记住，在典型的 transformer 架构中，embedding_dimension 是唯一应用 LayerNorm 的维度。 同样，Dropout（在注意力块之外）是在最后一个嵌入维度上应用的操作。 这两个层在序列维度上是独立的，因此可以在单独的 GPU 上分块处理。 如下图所示，多头 transformer 块中的初始 LayerNorm 是并行执行的。 接下来，收集结果用于自注意力层和线性层（通常设置为张量并行）。 接下来，结果从这些层分散回序列并行节点，这些节点执行 dropout，从先前的序列并行输出执行残差连接，以及 layernorm。 接下来，再次收集这些结果，用于最终的 FFN 和激活层，然后再分散到序列并行 GPU，以用于该 transformer 块的输出。

作为用户，如果您知道您的 transformer 是并行执行的，并且您有自定义损失或下游层，则需要确保为您的损失计算等发生适当的收集操作。

上下文并行

上下文并行 通过并行化注意力机制本身来扩展序列并行，类似于 Ring Attention。 一般来说，如果您使用 transformer，那么对于非常长的输入序列，上下文并行将比序列并行表现更好。 也就是说，由于整个架构中都需要 all-gather 和 reduce scatter 操作，因此如果微批次可以放入单个设备，您应该避免这些类型的并行的一般建议仍然成立。 在全局批次中拆分元素代表集群中 GPU 之间最少的必要通信，因此，如果可以使微批次的训练循环适合一个 GPU，则标准 DDP 应该运行得最快。

混合并行策略

您可以将多种并行类型混合在一起，以实现更高效的结果。 一般来说，应进行实验以确定最佳并行组合。 请观看 Jared Casper 的 YouTube 教程，以获取有关 megatron 并行策略的更多背景知识。

下图演示了混合策略如何产生更大的“虚拟 GPU”，类似地，这意味着您的集群中正在运行的独特微批次更少。 另请注意，虚拟 GPU 的数量是乘法的，因此如果您有 TP=2 和 PP=2，那么您正在使用 2*2=4 个 GPU 创建更大的虚拟 GPU，因此在这种情况下，您的集群大小需要是 4 的倍数。

调度模型并行

您可以通过将微批次拆分为更小的块、在单个 GPU 上执行模型的多个阶段以及在一个微批次进行前向传递时开始计算另一个微批次的后向传递来改进朴素调度。 这些优化可以更好地实现集群 GPU 利用率。 例如，下图显示了 megatron 中更高级的拆分技术（例如，交错调度器）如何在模型并行使用时提供更好的利用率。 同样，当您可以避免使用模型并行（DDP）时，这通常是最佳方法。