Implementation:Microsoft DeepSpeedExamples SQuAD Baseline Training

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	NLP, Question Answering
Last Updated	2026-02-07 12:00 GMT

Overview

BERT SQuAD question answering training script using NVIDIA Apex for mixed-precision FP16 training as a baseline without DeepSpeed optimizations.

Description

nvidia_run_squad_baseline.py implements fine-tuning of BERT for the SQuAD (Stanford Question Answering Dataset) extractive question answering task using NVIDIA Apex AMP for mixed-precision training. This serves as the baseline implementation against which the DeepSpeed-optimized version is compared.

The script implements the full SQuAD training pipeline. The SquadExample class represents a single QA example with question text, document tokens, and answer span positions. The read_squad_examples() function parses SQuAD JSON files into examples, handling whitespace tokenization and answer span validation. The convert_examples_to_features() function performs WordPiece tokenization with sliding window document chunking (controlled by doc_stride), creating InputFeatures with input IDs, attention masks, segment IDs, and span positions.

Training uses BertForQuestionAnswering from the Turing NVIDIA modeling module, with BertAdam optimizer and Apex AMP for FP16 training. The evaluation pipeline computes start and end logits, applies n-best prediction extraction with _get_best_indexes() and _compute_softmax(), and writes predictions in SQuAD JSON format. The script supports distributed training via PyTorch's DistributedDataParallel and integrates with TensorBoard summary writing via utility functions.

Usage

Use this script as the baseline for BERT SQuAD training with standard Apex mixed-precision. Compare its performance against the DeepSpeed version (nvidia_run_squad_deepspeed.py) to measure DeepSpeed's training acceleration.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/BingBertSquad/nvidia_run_squad_baseline.py
• Lines: 1-1158

Signature

class SquadExample(object):
    def __init__(self, qas_id, question_text, doc_tokens,
                 orig_answer_text=None, start_position=None, end_position=None):
        ...

class InputFeatures(object):
    def __init__(self, unique_id, example_index, doc_span_index, tokens,
                 token_to_orig_map, token_is_max_context, input_ids,
                 input_mask, segment_ids, start_position=None, end_position=None):
        ...

def read_squad_examples(input_file, is_training):
    ...

def convert_examples_to_features(examples, tokenizer, max_seq_length,
                                 doc_stride, max_query_length, is_training):
    ...

def main():
    ...

Import

# This is a standalone training script, run directly:
# python -m torch.distributed.launch nvidia_run_squad_baseline.py ...

I/O Contract

Inputs

Name	Type	Required	Description
--bert_model	str	Yes	Path to pretrained BERT model or HuggingFace model name
--output_dir	str	Yes	Directory for predictions and checkpoints
--train_file	str	No	Path to SQuAD training JSON file (required if --do_train)
--predict_file	str	No	Path to SQuAD evaluation JSON file (required if --do_predict)
--do_train	flag	No	Run training phase
--do_predict	flag	No	Run prediction phase
--max_seq_length	int	No	Maximum sequence length after tokenization (default: 384)
--doc_stride	int	No	Stride for sliding window over long documents (default: 128)
--max_query_length	int	No	Maximum query token length (default: 64)
--train_batch_size	int	No	Training batch size (default: 32)
--learning_rate	float	No	Learning rate for BertAdam (default: 5e-5)
--num_train_epochs	float	No	Number of training epochs (default: 3.0)
--local_rank	int	No	Local rank for distributed training

Outputs

Name	Type	Description
predictions.json	file	SQuAD-format predictions mapping question IDs to predicted answer text
nbest_predictions.json	file	Top-N predictions with probabilities for each question
model checkpoint	directory	Saved model weights in output_dir

Usage Examples

Train BERT on SQuAD v1.1 with Apex FP16

# Distributed training with Apex mixed-precision
python -m torch.distributed.launch --nproc_per_node=4 \
    nvidia_run_squad_baseline.py \
    --bert_model bert-large-uncased \
    --do_train \
    --do_predict \
    --train_file /data/squad/train-v1.1.json \
    --predict_file /data/squad/dev-v1.1.json \
    --output_dir /output/squad_baseline \
    --train_batch_size 24 \
    --learning_rate 3e-5 \
    --num_train_epochs 2.0 \
    --max_seq_length 384 \
    --doc_stride 128 \
    --fp16