python - 如何使用 PyTorch 训练具有大文本语料库（200GB）的 Masked Language Model？

Question

最近我正在使用transformers训练一个带有大文本语料库（200GB）的掩码语言模型。训练数据太大，无法装入配备 512GB 内存和 V100(32GB)*8 的计算机。是否有可能找到一种优雅的方式来用大数据训练模型？

现在我将整个训练数据分成 20 份并创建 20 DataLoaders，每个数据加载器将相应的数据加载到内存中并用于DistributedDataParallel训练模型。代码如下。

def get_dataloader(rank, world_size, tokenizer, part_of_data, args):
    train_data_df = get_DataFrame(args, part_of_data)
    train_dataset = CustomDataset(train_data_df, tokenizer, args.max_len)

    data_collator = transformers.DataCollatorForLanguageModeling(
        tokenizer=tokenizer, mlm=True, mlm_probability=0.15
    )

    train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset,
        num_replicas=world_size,
        rank=rank,
    )

    train_loader = torch.utils.data.DataLoader(
        dataset=train_dataset,
        batch_size=args.batch_size,
        shuffle=False,            
        num_workers=2,
        pin_memory=True,
        collate_fn=data_collator,
        sampler=train_sampler)

    return train_loader

通过阅读 PyTorch Doc，我发现了一个名为IterableDataset的数据加载器，也许 IterableDataset 更适合我的任务，但是文档中有一条注释说sampler 和 batch_sampler 都不兼容 iterable-style 数据集，因为这样的数据集没有键的概念或索引。

所以我想知道是否可以使用IterableDataset, DistributedSamplerandDistributedDataParallel来训练模型，或者其他优雅的方法来训练模型而不是拆分数据。非常感谢。