如何使用 PyTorch 训练具有大文本语料库（200GB）的掩码语言模型？

Question

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

现在我将整个训练数据分成 20 块并创建 20 个 DataLoader，每个数据加载器将相应的数据加载到内存中并使用 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 更适合我的任务，但是文档中有一个注释说采样器和batch_sampler 都与迭代式数据集不兼容，因为此类数据集没有键或索引的概念。

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