重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和新库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

Lightning 与 Megatron Core 之间的桥梁#

在 PyTorch Lightning 中，Strategy 负责管理模型在训练、验证和测试期间的分布式执行。Strategy 通常使用一个可以处理分布式执行的类来包装用户定义的模型。例如，标准的 DDPStrategy（分布式数据并行策略）使用 PyTorch 的 DistributedDataParallel 类来包装模型。此包装器处理跨多个 GPU 或节点的数据分发，在反向传播期间同步梯度，并确保模型参数在所有进程中保持一致。Lightning 中的 Strategy 抽象化了分布式训练的许多复杂性，使用户可以专注于他们的模型架构和训练逻辑，而框架则处理分布式执行的复杂性。

MegatronStrategy#

MegatronStrategy 是一种 PyTorch Lightning 策略，它支持使用 NVIDIA 的 Megatron-LM 框架对大型语言模型进行分布式训练。它旨在通过实现各种形式的模型并行性来处理超出单个 GPU 内存容量的模型。

要使用 MegatronStrategy，您需要使用定义并行设置的参数对其进行初始化

from nemo import lightning as nl

strategy = nl.MegatronStrategy(
    tensor_model_parallel_size=2,
    pipeline_model_parallel_size=2,
    virtual_pipeline_model_parallel_size=None,
    context_parallel_size=1,
    sequence_parallel=False,
    expert_model_parallel_size=1,
)

这些参数决定了模型将在可用 GPU 之间如何拆分。然后，该策略设置必要的分布式环境，并为每种类型的并行性初始化进程组。

该策略还负责配置检查点 IO 接口，该接口处理检查点的保存和加载。有关可以通过 MegatronStrategy 配置的完整选项列表，请参阅文档。

当您创建 PyTorch Lightning Trainer 时，您需要传递此策略

trainer = nl.Trainer(strategy=strategy, devices=8, accelerator="gpu")

MegatronStrategy 利用 Megatron 的分布式检查点系统进行模型 I/O。该系统有效地管理跨多个 GPU 分区的模型的检查点，并在各种并行配置中保持一致性。即使在保存和加载之间 GPU 设置不同，它也能实现正确的模型重建。

MegatronStrategy 包装了用户定义的 training_step、validation_step 和 test_step 方法，以使其与 Megatron 的前向-后向传播实现兼容。此包装过程允许在 Megatron 的分布式执行框架的上下文中执行这些步骤，从而确保在训练循环的每个阶段都正确处理所有形式的并行性。通过这样做，该策略为用户保持了熟悉的 PyTorch Lightning 界面，同时无缝集成了大规模模型训练所需的复杂分布式操作。

MegatronStrategy 采用 MegatronParallel 类来管理用户定义模型的分布式执行。此类将执行过程分解为三个关键步骤

  1. 数据步骤：准备输入数据并将其分发到模型并行组。

  2. 前向步骤：跨分区模型执行前向传播。

  3. 损失减少：计算并减少分布式设置中的损失。

MegatronParallel 利用这些步骤来执行前向-后向传播，这源自用户定义的 training_step、validation_step 和 test_step 方法。它协调数据和梯度通过分区模型的流动，管理 GPU 间通信，并确保正确的梯度同步。这种方法可以在多个 GPU 上实现高效执行，同时保留用户 Lightning 模块的逻辑结构。

MegatronParallel#

MegatronParallel 类是在 Megatron-LM 框架中实现分布式模型并行性的核心组件。它管理模型在多个 GPU 上的执行，并将过程分解为三个关键步骤

  1. 数据步骤：此步骤准备输入数据并将其分发到模型并行组。对于 GPT 模型，它使用 gpt_data_step 函数

    def data_step(self, dataloader_iter):
    return gpt_data_step(dataloader_iter)

    此函数处理

    1. 从数据加载器中获取批次

    2. 将所需的张量移动到 CUDA

    3. 使用 get_batch_on_this_context_parallel_rank 为上下文并行切片批次

    4. 必要时准备打包序列参数

  2. 前向步骤：此步骤跨分区模型执行前向传播。对于 GPT 模型，它使用 gpt_forward_step 函数

    def forward_step(self, model, batch):
    return gpt_forward_step(model, batch)

    此函数

    1. 从批次中准备前向参数

    2. 使用这些参数调用模型的前向方法

    3. 处理标准和打包序列输入

  3. 损失减少：在前向传播之后，此步骤计算并减少分布式设置中的损失。GPT 模型使用 MaskedTokenLossReduction

    def loss_reduction(self, model):
    return model.training_loss_reduction()

    用于验证

    def validation_loss_reduction(self, model):
    return model.validation_loss_reduction()

    这些方法处理

    1. 使用掩码 token 损失计算损失

    2. 减少跨数据并行组的损失

    3. 处理验证的特殊情况（例如，不丢弃最后一个批次）

MegatronParallel 类协调这些步骤以执行完整的前向-后向传播。

通过使用这些特定于模型的功能，MegatronParallel 允许 GPT 模型定义其自己的数据处理、前向传播和损失计算逻辑，同时仍然受益于分布式执行框架。这种方法使研究人员和工程师可以使用熟悉的 PyTorch Lightning 界面来处理大型语言模型，而底层的分布式执行是透明处理的。

MegatronMixedPrecision#

MegatronMixedPrecision 类是 PyTorch Lightning 中 Megatron-LM 模型的专用精度插件。它扩展了标准的 MixedPrecision 插件，以处理使用 Megatron-LM 训练的大型语言模型的特定要求。

from nemo import lightning as nl

precision = nl.MegatronMixedPrecision(precision="bf16-mixed")
trainer = nl.Trainer(strategy=strategy, plugins=precision)
本页内容