Implementation:Microsoft DeepSpeedExamples HelloDeepSpeed Train Bert

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	Natural Language Processing, Model Training
Last Updated	2026-02-07 12:00 GMT

Overview

Standalone BERT training script for the HelloDeepSpeed tutorial that implements a complete MLM training pipeline using PyTorch with RoBERTa architecture on the WikiText dataset.

Description

This module provides a self-contained BERT-style training script as part of the HelloDeepSpeed tutorial series. It implements a complete Masked Language Model (MLM) training pipeline using the RoBERTa architecture from HuggingFace Transformers, training on the WikiText dataset. The script serves as the baseline PyTorch implementation before DeepSpeed integration is added.

The module includes several key components: masking_function() for implementing BERT's MLM preprocessing with configurable mask probability, random replacement, and unmasking probabilities; WikiTextMLMDataset as a map-style PyTorch Dataset class; InfiniteIterator for wrapping data loaders into infinite iterators; collate_function() for batching and padding sequences; and create_data_iterator() for assembling the full data pipeline. The model components include RobertaLMHeadWithMaskedPredict for efficient MLM prediction on only the masked positions, and RobertaMLMModel which combines the RoBERTa encoder with the masked prediction head.

The main train() function orchestrates the full training pipeline including experiment directory creation, model construction, data loading, training loop with gradient accumulation, periodic logging via TensorBoard, and checkpoint saving/loading. The script uses the fire library for CLI argument parsing and supports single-GPU training with configurable hyperparameters.

Usage

Use this script as a learning example for BERT-style MLM training with standard PyTorch. It serves as the baseline before applying DeepSpeed optimizations in train_bert_ds.py. Run it directly with configurable parameters for model architecture, dataset settings, and training hyperparameters.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/HelloDeepSpeed/train_bert.py
• Lines: 1-791

Signature

class WikiTextMLMDataset(Dataset):
    def __init__(self, dataset, masking_function):

class InfiniteIterator(object):
    def __init__(self, iterable):

class RobertaLMHeadWithMaskedPredict(RobertaLMHead):

class RobertaMLMModel(RobertaPreTrainedModel):
    def __init__(self, config, encoder):
    def forward(self, src_tokens, attention_mask, tgt_tokens):

def create_model(num_layers, num_heads, ff_dim, h_dim, dropout):

def train(checkpoint_dir=None, load_checkpoint_dir=None,
          mask_prob=0.15, random_replace_prob=0.1,
          unmask_replace_prob=0.1, max_seq_length=512,
          tokenizer="roberta-base", num_layers=6, num_heads=8,
          ff_dim=512, h_dim=256, dropout=0.1, batch_size=8,
          num_iterations=10000, checkpoint_every=1000,
          log_every=10, local_rank=-1):

def collate_function(batch, pad_token_id):
def masking_function(text, tokenizer, mask_prob, random_replace_prob,
                     unmask_replace_prob, max_length):

Import

from train_bert import train, create_model, RobertaMLMModel, WikiTextMLMDataset

I/O Contract

Inputs

Name	Type	Required	Description
checkpoint_dir	str	Yes (or load_checkpoint_dir)	Directory path to save experiment checkpoints
load_checkpoint_dir	str	No	Directory to resume training from an existing checkpoint
mask_prob	float	No	Fraction of tokens to mask (default: 0.15)
max_seq_length	int	No	Maximum sequence length (default: 512)
num_layers	int	No	Number of transformer layers (default: 6)
num_heads	int	No	Number of attention heads (default: 8)
h_dim	int	No	Hidden dimension size (default: 256)
batch_size	int	No	Training batch size (default: 8)
num_iterations	int	No	Total training iterations (default: 10000)
local_rank	int	No	GPU device index, -1 for CPU (default: -1)

Outputs

Name	Type	Description
experiment_dir	pathlib.Path	Path to the experiment directory with checkpoints and logs
masked_lm_loss	torch.Tensor	MLM cross-entropy loss per training step
model_checkpoint	dict	Saved model state dict, optimizer state, and iteration count

Usage Examples

# Run training from the command line using fire
# python train_bert.py train --checkpoint_dir ./experiments \
#     --num_layers 6 --num_heads 8 --h_dim 256 --batch_size 8

# Programmatic usage
from train_bert import train

experiment_dir = train(
    checkpoint_dir="./experiments",
    num_layers=6,
    num_heads=8,
    ff_dim=512,
    h_dim=256,
    dropout=0.1,
    batch_size=8,
    num_iterations=10000,
    log_every=10,
)