Implementation:Microsoft DeepSpeedExamples HelloDeepSpeed Train Bert DS

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

Overview

DeepSpeed-integrated BERT training script for the HelloDeepSpeed tutorial that extends the baseline PyTorch training with DeepSpeed engine initialization, distributed training, mixed precision, and ZeRO optimization.

Description

This module is a modified version of train_bert.py that adds DeepSpeed integration for distributed and efficient training. It demonstrates how to convert a standard PyTorch training script to use DeepSpeed's features including engine initialization, distributed data parallel training, mixed-precision (bf16/fp16), and ZeRO optimizer state partitioning. The module serves as the culminating tutorial in the HelloDeepSpeed series.

The module shares the same core components as the baseline script: WikiTextMLMDataset for data loading, masking_function() for MLM preprocessing, RobertaLMHeadWithMaskedPredict for efficient masked prediction, and RobertaMLMModel combining the encoder with the LM head. Key additions include log_dist() for rank-aware distributed logging, is_rank_0() for identifying the primary process, and integration with DeepSpeed's deepspeed.initialize() for engine setup.

The train() function has been enhanced with DeepSpeed-specific features: it uses deepspeed.initialize() to wrap the model and optimizer, supports the dtype parameter for selecting between bf16 and fp16 training, handles distributed checkpoint saving and loading through DeepSpeed's checkpoint API, and supports multi-GPU distributed training with proper rank-based data sharding and logging. The training loop uses DeepSpeed's model_engine.backward() and model_engine.step() instead of manual gradient management.

Usage

Use this script to train a BERT-style MLM model with DeepSpeed optimizations. It demonstrates the minimal changes needed to convert a PyTorch training script to DeepSpeed, making it an ideal tutorial for learning DeepSpeed integration. Launch with deepspeed CLI for multi-GPU training.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/HelloDeepSpeed/train_bert_ds.py
• Lines: 1-863

Signature

def is_rank_0() -> bool:

def log_dist(message: str, ranks: List[int] = [],
             level: int = logging.INFO) -> None:

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

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, dtype="bf16"):

Import

from train_bert_ds import train, create_model, RobertaMLMModel

I/O Contract

Inputs

Name	Type	Required	Description
checkpoint_dir	str	Yes (or load_checkpoint_dir)	Directory to save experiment checkpoints
load_checkpoint_dir	str	No	Directory to resume from an existing DeepSpeed 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	Per-GPU training batch size (default: 8)
num_iterations	int	No	Total training iterations (default: 10000)
local_rank	int	No	Local GPU rank for distributed training (default: -1)
dtype	str	No	Training precision type, "bf16" or "fp16" (default: "bf16")

Outputs

Name	Type	Description
experiment_dir	pathlib.Path	Path to the experiment directory with DeepSpeed checkpoints and logs
masked_lm_loss	torch.Tensor	MLM cross-entropy loss per training step
ds_checkpoint	directory	DeepSpeed checkpoint directory with model and optimizer states

Usage Examples

# Launch with DeepSpeed for multi-GPU training
# deepspeed train_bert_ds.py train --checkpoint_dir ./experiments \
#     --num_layers 6 --num_heads 8 --h_dim 256 --batch_size 8 --dtype bf16

# Single GPU usage
from train_bert_ds 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,
    local_rank=0,
    dtype="bf16",
)