编写 Megatron-LM 兼容的 Datamodule

Megatron-LM 依赖于训练数据集类中的确定性，以确保跨模型并行 ranks 正确初始化输入张量（参见 NeMo2 并行性）。因此，新的数据集类必须小心地保持所需的确定性。常见操作（如数据增强、掩码等）可能导致 dataset[i] 对于给定索引返回随机结果，从而破坏此 megatron 约定。

多轮次训练

这种限制最明显的训练方案之一是多轮次训练，其中标准训练方案会在每次遇到数据时应用不同的随机掩码或不同的数据增强策略。BioNeMo 提供了许多实用工具，可以更轻松地进行多轮次训练，同时仍然遵守 megatron 的确定性要求。

MultiEpochDatasetResampler 类简化了多轮次训练的过程，其中数据应该在每个轮次重新洗牌，并在每次看到数据时应用不同的随机效应。为了与此重采样器兼容，所提供的数据集类的 __getitem__ 方法应接受一个 EpochIndex 元组，其中包含轮次和索引值。然后可以通过基于轮次值设置 torch 随机种子来执行随机效应

class MyDataset:
    def __getitem__(self, idx: EpochIndex):
        rng = torch.Generator()
        rng.manual_seed(idx.epoch)
        ...

避免 torch.manual_seed

Megatron-LM 在内部处理 torch 播种。在用户提供的数据集中调用 torch.cuda.manual_seed 可能会导致模型并行性问题。有关更多详细信息，请参见 megatron/core/tensor_parallel/random.py#L198-L199。

对于仍然希望通过轮次级别洗牌进行多轮次训练的确定性数据集，IdentityMultiEpochDatasetWrapper 类可以通过包装接受整数索引的数据集并传递来自重采样数据集的 EpochIndex 索引值来简化此过程。

class MyDeterministicDataset:
    def __getitem__(self, index: int):
        ...

dataset = IdentityMultiEpochDatasetWrapper(MyDeterministicDataset())
for sample in MultiEpochDatasetResampler(dataset, num_epochs=3, shuffle=True):
    ...

训练恢复

为了确保在有和没有作业中断的情况下行为相同，BioNeMo 提供了 MegatronDataModule 来保存和加载状态字典以进行训练恢复，并提供了 [WrappedDataLoader][nemo.lightning.data.WrappedDataLoader] 以向 [DataLoader][torch.utils.data.DataLoader] 添加 mode 属性。

class MyDataModule(MegatronDataModule):
    def __init__(self, *args, **kwargs):
        super().__init__()
        ...

    def train_dataloader(self):
        self.update_init_global_step()  # required to set the correct `global_step` for resumption
        return WrappedDataLoader(
            ...,  # any other arguments for DataLoader
            mode="train",
        )

    def val_dataloader(self):
        self.update_init_global_step()  # required to set the correct `global_step` for resumption
        return WrappedDataLoader(
            ...,  # any other arguments for DataLoader
            mode="validation",
        )

    def test_dataloader(self):
        self.update_init_global_step()  # required to set the correct `global_step` for resumption
        return WrappedDataLoader(
            ...,  # any other arguments for DataLoader
            mode="test",
        )

MegatronDataModule

如果用户的数据模块不是继承自 MegatronDataModule，除非用户处理了类似的逻辑，否则他们将看到不重叠的训练曲线。在 MegatronDataModule 中，必须在返回数据加载器之前调用 self.update_init_global_step()，以确保训练使用正确的样本索引恢复，而不是每次都从 0 重新开始。我们建议用户继承自 MegatronDataModule，类似于上面的模式。

WrappedDataLoader

WrappedDataLoader 类是 PyTorch DataLoader 类的包装器，它向数据加载器添加了 mode 属性。仅当 mode 为“train”时，数据加载器才会从最后一个样本索引恢复。val_dataloader 和 test_dataloader 不受影响。

警告：'train' 是 WrappedDataLoader 中 mode 的默认值。如果未设置，用户可能会发现他们的验证/测试数据加载器通过从非零样本索引恢复而改变行为。

测试数据集与 Megatron 的兼容性

BioNeMo 还为测试套件提供了实用函数，以验证数据集是否符合 megatron 数据模型。[assert_dataset_compatible_with_megatron][bionemo.testing.data_utils.assert_dataset_compatible_with_megatron] 函数使用相同的索引调用数据集，并确保输出相同，同时还检查是否使用了 torch.manual_seed。

BioNeMo 中的示例数据集

ESMMaskedResidueDataset 演示了一种利用 EpochIndex 索引在 megatron 数据模型的限制范围内执行轮次级别随机化的方法。