跳到内容

bionemo-example_model

简介

这是一个包含示例模型的极简软件包，该模型利用了 bionemo2 和 nemo 约定。它包含必要的模型、数据加载器、数据集和自定义损失函数。引用的类和函数位于 bionemo.example_model.lightning.lightning_basic 中。

本教程演示了如何创建一个简单的 MNIST 模型。这应该在 BioNeMo 容器中运行。BioNeMo 框架容器可以在 brev.dev 启动器中运行： 点击此处部署。。将此笔记本作为 Launchable 部署大约需要 10 分钟。在撰写本文时，我们正在开发免费层，因此可能需要信用卡。您可以联系您的 NVIDIA 代表以获取信用额度。笔记本和 shell 界面可以通过单击“打开笔记本”启动。（注意：此链接指向 nightly 版本，可能与这些文档不同步。）

在本教程中，我们将重用 BioNeMo example_model 包中的元素。

Megatron/NeMo 模块和数据集是 PyTorch 模块和数据集的特殊衍生产品，它们扩展并加速了 PyTorch 的分布式训练和推理能力。

Megatron/NeMo 的一些区别是

  • torch.nn.Module/LightningModule 更改为 MegatronModule。
  • 损失函数应扩展 MegatronLossReduction 模块，并实现一个 reduce 方法，用于聚合多个微批次的损失。
  • Megatron 配置类（例如 megatron.core.transformer.TransformerConfig）扩展了一个 configure_model 方法，该方法定义了如何初始化和加载模型权重，以符合通过 NeMo2 进行训练的方式。
  • 对常见 PyTorch 类的各种修改和扩展，例如向您的 LightningDataModule 添加 MegatronDataSampler（以及重新采样器，例如 PRNGResampleDataset 或 MultiEpochDatasetResampler）。

损失函数

首先，我们在 bionemo.example_model.lightning.lightning_basic 中定义一个简单的损失函数。这些函数应扩展 MegatronLossReduction 类。前向和后向传递的输出并行发生。应该有一个前向函数来计算定义的损失。需要 reduce 函数。

此处使用的损失函数是 MSELossReduction 和 ClassifierLossReduction。这些函数返回一个 Tensor，其中包含微批次的损失，以及一个包含平均损失的 SameSizeLossDict。这是一个类型化字典，它是为整个批次计算的损失的返回类型，其中所有微批次的大小相同。

数据集和数据模块

用于模型训练的数据集必须与 Megatron 数据集兼容。为了实现这一点，给定索引和 epoch 的输出必须是确定性的。但是，我们可能希望在每个 epoch 中有不同的排序。为了实现这一点，数据集中的项目应该可以通过 epoch 和索引访问。这可以通过使用来自 bionemo.core.data.multi_epoch_dataset 的 EpochIndex 访问数据集的元素来完成。一种简单的方法是使用从 bionemo.core.data.multi_epoch_dataset 导入的 IdentityMultiEpochDatasetWrapper 包装数据集。在本例中，在 bionemo.example_model.lightning.lightning_basic 中，我们使用自定义数据集 MNISTCustomDataset，它包装了 MNIST 数据集的 __getitem__ 方法，使其返回 dict 而不是 Tuple 或 tensor。MNISTCustomDataset 返回 MnistItem 类型的元素，这是一个 TypedDict。

在数据模块/数据加载器类中，必须有一个 data_sampler 属性来打乱数据，并允许采样器与 Megatron 一起使用。这是 nemo2 的特性。nemo.lightning.pytorch.plugins.MegatronDataSampler 是最佳选择。它设置了利用微批处理和梯度累积的能力。它也是构建全局批次大小的地方。

此外，采样器不会打乱您的数据。因此，您需要将数据集包装在数据集打乱器中，该打乱器将顺序 ID 映射到数据集中的随机 ID。这可以使用来自 bionemo.core.data.multi_epoch_dataset 的 MultiEpochDatasetResampler 完成。

这在 MNISTDataModule 中实现。在数据加载器的 setup 方法中，训练集、测试集和验证集 MNISTCustomDataset 被包装在 IdentityMultiEpochDatasetWrapper 中。然后将这些包装在 MultiEpochDatasetResampler 中。有关 MegatronCompatability 以及如何设置更复杂数据集的更多信息，请参阅 docs.user-guide.background.megatron_datasets.md。

我们还定义了 train_dataloader、val_dataloader 和 predict_dataloader 方法，这些方法返回相应的数据加载器。

模型

模型需要是 Megatron 模块。在最基本的层面上，这仅仅意味着

  1. 它们扩展了 megatron.core.transformer.module 中的 MegatronModule。
  2. 它们需要一个 megatron.core.ModelParallelConfig 类型的 config 参数。实现此目的的一种简单方法是从 bionemo.llm.model.config.MegatronBioNeMoTrainableModelConfig 继承。这是一个用于 BioNeMo 的类，支持与 Megatron 模型一起使用，因为 NeMo2 需要这样做。此类也继承了 ModelParallelConfig。
  3. 它们需要定义一个 self.model_type:megatron.core.transformer.enums.ModelType 枚举（ModelType.encoder_or_decoder 是一个不错的选择。）
  4. 需要存在 def set_input_tensor(self, input_tensor)。这在模型并行性中使用。此函数可以是存根/占位符函数。

以下模型在 bionemo.example_model.lightning.lightning_basic 中实现。

ExampleModelTrunk 是一个基础模型。这返回一个张量。ExampleModel 是一个模型，它通过添加一些线性层来扩展基础模型，它用于预训练。这返回基础模型和完整模型的输出。

ExampleFineTuneModel 通过添加分类层来扩展 ExampleModelTrunk。这返回一个关于 10 个潜在数字的 logits 张量。

模型配置

模型配置类用于实例化模型。这些配置必须具有：1. 一个 configure_model 方法，该方法允许 Megatron 策略在并行计算环境设置后延迟初始化模型。这些配置还处理加载微调案例的起始权重。此外，这些配置告诉训练器您想要使用哪个损失函数以及匹配的模型。2. 一个 get_loss_reduction_class 方法，用于定义损失函数。

以下配置在 bionemo.example_model.lightning.lightning_basic 中实现。

在这里，定义了一个基础通用配置 ExampleGenericConfig。PretrainConfig 扩展了此类。这在其中定义了模型类和损失类

class PretrainConfig(ExampleGenericConfig["PretrainModel", "MSELossReduction"], iom.IOMixinWithGettersSetters):

    model_cls: Type[PretrainModel] = PretrainModel
    loss_cls: Type[MSELossReduction] = MSELossReduction

类似地，ExampleFineTuneConfig 扩展了用于微调的 ExampleGenericConfig。

训练模块

拥有一个从 lightning.pytorch.LightningModule 继承的训练模块会很有帮助，它可以组织模型架构、训练、验证和测试逻辑，同时抽象出样板代码，从而实现更轻松、更可扩展的训练。此包装器可用于配置中指定的所有模型和损失组合。在 bionemo.example_model.lightning.lightning_basic 中，我们定义了 BionemoLightningModule。

在本例中，training_step、validation_step 和 predict_step 定义的训练、验证和预测循环独立于前向方法。在 nemo 中

  1. NeMo 的策略覆盖了 train_step、validation_step 和 prediction_step 方法。
  2. 策略的训练步骤将调用模型的前向方法。
  3. 该前向方法然后调用 MegatronParallel 的包装前向步骤，后者包装了模型的前向方法。
  4. 该包装的前向步骤然后在 MegatronCore 调度器内部执行，该调度器从 MegatronParallel 类调用 _forward_step 方法。
  5. 然后，此处调用 training_step、validation_step 和 prediction_step 函数。

此外，在这些步骤中，我们记录验证、测试和训练损失。这与 https://lightning.ai/docs/torchmetrics/stable/pages/lightning.html 类似地完成。然后可以将这些日志导出到 wandb 或其他指标查看器。对于更复杂的跟踪，可能需要使用 pytorch 回调：https://lightning.ai/docs/pytorch/stable/extensions/callbacks.html。

此外，loss_reduction_class()、training_loss_reduction()、validation_loss_reduction(), 和 test_loss_reduction() 是根据配置中定义的内容定义的。此外，configure_model() 是根据配置定义的。

训练模型

在 bionemo.example_model.lightning.lightning_basic 中，定义了一个 checkpoint_callback 变量。这启用了类似 .nemo 文件的检查点。

其余函数在训练脚本中定义：pretrain_mnist.py、finetune_mnist.py 和 predict_mnist.py。

我们指定 nemo.lightning.MegatronStrategy 类型的训练策略。此策略使用 NVIDIA 的 Megatron-LM 框架实现模型并行性。它支持各种形式的并行性，包括张量模型并行性、流水线模型并行性、序列并行性和专家并行性，用于高效训练大型语言模型。

我们指定 nemo.lightning.Trainer 类型的训练器，它是 pytorch lightning 训练器的扩展。这是指定设备、验证间隔、最大步骤、最大 epoch 数以及日志记录频率的位置。

我们指定一个 nemo-logger。我们可以设置 TensorBoard 和 WandB 日志记录，以及额外的记录器。在这里，我们指定来自 lightning.pytorch.loggers 的 CSVLogger。

我们现在可以继续训练。第一个预训练脚本是 bionemo/example_model/training_scripts/pretrain_mnist.py

然后，我们使用 BionemoLightningModule、MNISTDataModule、训练器和 nemo_logger 训练模型。

此脚本将打印最终模型的位置

然后我们可以运行微调脚本

python src/bionemo/example_model/training_scripts/finetune_mnist.py ---pretrain_ckpt_dirpath <pretrain_directory>

这里的一个细微之处是，在配置文件中，我们指定了初始检查点路径，以及要跳过的键。在之前的模型检查点中，我们没有标记为“digit_classifier”的头部，因此我们将其指定为要跳过的头部。此脚本将打印微调目录的位置.

最后，我们可以使用以下命令运行分类任务

python src/bionemo/example_model/training_scripts/predict_mnist.py  --finetune_dir <finetune_dir>.

如果配置了 TensorBoardLogger，则可以在 TensorBoardLogger 中查看结果，也可以在 CSVLogger 创建的 CSV 文件中查看结果。