Implementation:Microsoft DeepSpeedExamples NvidiaBert DatasetProvider

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	Data Loading, Distributed Training
Last Updated	2026-02-07 12:00 GMT

Overview

An HDF5-based dataset provider for BERT pretraining that implements distributed file sharding and asynchronous shard prefetching using a process pool executor.

Description

NvidiaBertDatasetProvider implements the BertDatasetProviderInterface to load BERT pretraining data from HDF5 files in the NVIDIA data format. The companion pretraining_dataset class reads HDF5 files containing input IDs, attention masks, segment IDs, masked LM positions, masked LM IDs, and next sentence labels, loading all arrays into memory upon initialization for fast access during training.

The provider manages a collection of HDF5 training files, automatically discovering files containing training in their names from a configured directory. It implements distributed file sharding by computing unique file indices per global rank, ensuring each GPU processes different data partitions. The sharding logic handles cases where the world size exceeds the number of files by introducing remainder-based offsets to distribute files evenly.

Shard prefetching is accomplished through a ProcessPoolExecutor that asynchronously loads the next shard while the current shard is being consumed. The prefetch_shard method submits the dataset creation task to the pool, and get_shard retrieves the result (blocking if necessary). The WorkerInitObj class ensures reproducible random seeding across worker processes.

Usage

Use this provider when running BERT pretraining with NVIDIA-format HDF5 data files. It is an alternative to BingBertDatasetProvider and is typically selected based on the data format available. The provider is optimized for large-scale distributed training with automatic file partitioning across GPUs.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/bing_bert/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, pretraining_dataset

I/O Contract

Inputs

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

Outputs

Name	Type	Description
train_dataloader	DataLoader	PyTorch DataLoader yielding batches of pretraining tensors
sample_count	int	Total number of samples in the loaded shard
batch	list[Tensor]	List containing batch_type, input_ids, input_mask, segment_ids, next_sentence_labels, masked_lm_labels

Usage Examples

from nvidia_bert_dataset_provider import NvidiaBertDatasetProvider

# Initialize the provider
provider = NvidiaBertDatasetProvider(args)

# Prefetch the next shard while processing current
provider.prefetch_shard(index=1)

# Load current shard
dataloader, sample_count = provider.get_shard(index=0)
for batch in dataloader:
    processed_batch = provider.get_batch(batch)
    # ... train on batch ...

# Release shard memory
provider.release_shard(index=0)

# Get the prefetched shard
dataloader, sample_count = provider.get_shard(index=1)