How to train a Masked Language Model with a big text corpus(200GB) using PyTorch?

Viewed 263

Recently I am training a masked language model with a big text corpus(200GB) using transformers. The training data is too big to fit into computer equiped with 512GB memory and V100(32GB)*8. Is it possible to find a elegant way to train model with big data?

Now I split entire training data into 20 pieces and create 20 DataLoaders, each dataloader loads corresponding data into memory and uses DistributedDataParallel to train the model. The code is below.

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

By reading PyTorch Doc I find a dataloader called IterableDataset, maybe IterableDataset is more suitable for my task, but there is a NOTE in the Doc said neither sampler nor batch_sampler is compatible with iterable-style datasets, since such datasets have no notion of a key or an index.

So I wonder if it is possible to use IterableDataset, DistributedSampler and DistributedDataParallel to train model, or other elegant method to train model instead of splitting data. Thanks a lot.

0 Answers
Related