Implementation:Microsoft DeepSpeedExamples Nvidia Bert Dataset Provider

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	Deep Learning, Natural Language Processing, Data Loading
Last Updated	2026-02-07 12:00 GMT

Overview

A BERT dataset provider that loads HDF5 pretraining data with shard-level scheduling, distributed sampling, and asynchronous prefetching for the BingBert GLUE training pipeline.

Description

The NvidiaBertDatasetProvider class implements the BertDatasetProviderInterface to supply BERT pretraining batches from HDF5 data files. It manages a collection of training shards discovered from a configured dataset path, shuffles them randomly, and distributes file access across ranks using a shard-index-to-file mapping that accounts for world size and number of available files. It uses a ProcessPoolExecutor to asynchronously prefetch the next shard while the current one is being consumed.

The underlying pretraining_dataset class (a PyTorch Dataset) reads six HDF5 fields per sample: input_ids, input_mask, segment_ids, masked_lm_positions, masked_lm_ids, and next_sentence_labels. The __getitem__ method reconstructs masked_lm_labels by scattering masked token IDs into a full-length label tensor at the positions indicated by masked_lm_positions, using -1 as the ignore index.

The helper function create_pretraining_dataset wraps dataset creation with a DataLoader that uses either RandomSampler or DistributedSampler, along with a WorkerInitObj to seed each data-loading worker for reproducibility. The provider exposes get_shard, release_shard, and prefetch_shard methods for the training loop to manage data lifecycle.

Usage

Use this provider when training BERT models with the BingBert/GLUE pipeline on NVIDIA-format HDF5 pretraining datasets. It is instantiated with the training arguments and called by the training loop to retrieve DataLoaders for each data shard in sequence, with optional prefetching for pipeline efficiency.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/BingBertGlue/nvidia_bert_dataset_provider.py
• Lines: 1-169

Signature

class WorkerInitObj(object):
    def __init__(self, seed):
    def __call__(self, id):

def create_pretraining_dataset(input_file, max_predictions_per_seq,
                               num_workers, train_batch_size, worker_init,
                               data_sampler):

class pretraining_dataset(Dataset):
    def __init__(self, input_file, max_predictions_per_seq):
    def __len__(self):
    def __getitem__(self, index):

class NvidiaBertDatasetProvider(BertDatasetProviderInterface):
    def __init__(self, args):
    def get_shard(self, index):
    def release_shard(self, index):
    def prefetch_shard(self, index):
    def get_batch(self, batch_iter):
    def prefetch_batch(self):

Import

from nvidia_bert_dataset_provider import NvidiaBertDatasetProvider

I/O Contract

Inputs

Name	Type	Required	Description
args	Namespace	Yes	Training arguments containing config, max_seq_length, max_predictions_per_seq, gradient_accumulation_steps, train_micro_batch_size_per_gpu, local_rank, data_path_prefix, seed, and logger
input_file	str	Yes	Path to an HDF5 file containing pretraining data
max_predictions_per_seq	int	Yes	Maximum number of masked LM predictions per sequence
index	int	Yes	Shard index for get_shard, release_shard, and prefetch_shard operations

Outputs

Name	Type	Description
train_dataloader	DataLoader	PyTorch DataLoader yielding batches of [batch_type, input_ids, input_mask, segment_ids, next_sentence_labels, masked_lm_labels]
sample_count	int	Total number of samples in the loaded shard

Usage Examples

from nvidia_bert_dataset_provider import NvidiaBertDatasetProvider

# Initialize the provider with training args
provider = NvidiaBertDatasetProvider(args)

# Get first shard
dataloader, num_samples = provider.get_shard(0)

# Prefetch next shard asynchronously
provider.prefetch_shard(1)

# Iterate over batches
for batch in dataloader:
    batch_type, input_ids, input_mask, segment_ids, nsp_labels, mlm_labels = batch
    # Training step...

# Release and move to next shard
provider.release_shard(0)
dataloader, num_samples = provider.get_shard(1)