Implementation:Microsoft DeepSpeedExamples BingBert DeepSpeedTrain

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

Overview

Main training script for Bing BERT pretraining with DeepSpeed distributed training integration.

Description

This module implements the end-to-end BERT pretraining pipeline using DeepSpeed as the distributed training backend. It orchestrates model initialization, data loading, checkpointing, validation, and the main training loop for large-scale BERT pretraining as used in Bing search.

The script provides checkpoint management through functions for saving and restoring training state, including epoch, global step, and data sample counts. It supports both the Bing BERT dataset provider and the NVIDIA BERT dataset provider for flexible data pipeline configurations.

The training loop handles gradient accumulation, learning rate scheduling (including warmup strategies), distributed data parallelism, and periodic validation. It integrates with DeepSpeed for mixed-precision training, ZeRO optimization, and efficient gradient communication across multiple GPUs.

Usage

Use this script as the main entry point for launching BERT pretraining with DeepSpeed. It is typically invoked via DeepSpeed launcher with a configuration JSON file specifying optimization parameters, batch sizes, and distributed training settings.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/bing_bert/deepspeed_train.py
• Lines: 1-600

Signature

def checkpoint_model(PATH, ckpt_id, model, epoch, last_global_step, last_global_data_samples, **kwargs)
def load_training_checkpoint(args, model, PATH, ckpt_id)
def get_dataloader(args, dataset: Dataset, eval_set=False)
def pretrain_validation(args, index, model)
def master_process(args)
def train(args, index, model, optimizer, pretrain_dataset_provider, finetune=False)
def update_learning_rate(args, config, current_global_step, optimizer)
def report_step_metrics(args, lr, loss, step, data_sample_count)
def get_arguments()
def construct_arguments()
def prepare_optimizer_parameters(args, model)
def prepare_model_optimizer(args)
def load_checkpoint(args, model)
def run(args, model, optimizer, start_epoch)
def main()

Import

from deepspeed_train import (
    checkpoint_model, load_training_checkpoint, train,
    prepare_model_optimizer, construct_arguments, main
)

I/O Contract

Inputs

Name	Type	Required	Description
args	Namespace	Yes	Parsed command-line arguments including model config, data paths, and DeepSpeed config
index	int	Yes	Current epoch index for training loop iteration
model	BertMultiTask	Yes	The BERT multi-task model wrapped by DeepSpeed engine
optimizer	Optimizer	Yes	DeepSpeed-managed optimizer
pretrain_dataset_provider	DatasetProvider	Yes	Provider that yields training data shards per epoch

Outputs

Name	Type	Description
model	DeepSpeedEngine	Trained model wrapped in DeepSpeed engine after training completes
checkpoint	dict	Saved checkpoint containing epoch, global step, and data sample counts
metrics	float	Training loss and validation loss reported via logger and TensorBoard

Usage Examples

# Launch BERT pretraining with DeepSpeed
# Command line:
# deepspeed deepspeed_train.py --deepspeed --deepspeed_config ds_config.json

# Programmatic usage of checkpoint functions
checkpoint_model(
    PATH="/output/checkpoints",
    ckpt_id="epoch_1",
    model=model,
    epoch=1,
    last_global_step=1000,
    last_global_data_samples=64000
)

epoch, step, samples = load_training_checkpoint(args, model, "/output/checkpoints", "epoch_1")