Skip to content

llm.trainers.bert.data

NvidiaBertDatasetConfig

Bases: NamedTuple

NVIDIA BERT pretraining dataset configuration.

RobertaDatasetConfig

Bases: NamedTuple

RoBERTa pretraining dataset configuration.

get_dataloader 

get_dataloader(
    dataset: DistributedShardedDataset[Sample],
    sampler: Sampler[int],
    batch_size: int,
) -> DataLoader[Sample]

Create a dataloader from a sharded dataset.

Source code in llm/trainers/bert/data.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def get_dataloader(
    dataset: DistributedShardedDataset[Sample],
    sampler: torch.utils.data.Sampler[int],
    batch_size: int,
) -> torch.utils.data.DataLoader[Sample]:
    """Create a dataloader from a sharded dataset."""
    return torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        sampler=sampler,
        num_workers=4,
        pin_memory=True,
        drop_last=True,
    )

get_dataset 

get_dataset(
    config: NvidiaBertDatasetConfig | RobertaDatasetConfig,
) -> DistributedShardedDataset[Sample]

Load a sharded BERT pretraining dataset.

Source code in llm/trainers/bert/data.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def get_dataset(
    config: NvidiaBertDatasetConfig | RobertaDatasetConfig,
) -> DistributedShardedDataset[Sample]:
    """Load a sharded BERT pretraining dataset."""
    files = get_filepaths(
        config.input_dir,
        extensions=['.h5', '.hdf5'],
        recursive=True,
    )

    dataset = torch.utils.data.Dataset
    params: dict[str, DatasetParams]
    if isinstance(config, NvidiaBertDatasetConfig):
        dataset = NvidiaBertDataset
        params = {file: ((file,), {}) for file in files}
    elif isinstance(config, RobertaDatasetConfig):
        tokenizer = tokenizers.Tokenizer.from_file(config.tokenizer_file)

        dataset = RoBERTaDataset
        kwargs = {
            'mask_token_id': tokenizer.token_to_id('[MASK]'),
            'mask_token_prob': config.mask_token_prob,
            'vocab_size': tokenizer.get_vocab_size(),
        }
        params = {file: ((file,), kwargs) for file in files}
    else:
        raise AssertionError('Unreachable.')

    return DistributedShardedDataset(
        dataset,
        params,
        rank=torch.distributed.get_rank(),
        world_size=torch.distributed.get_world_size(),
    )