Implementation:Microsoft DeepSpeedExamples BERT Checkpoint Conversion

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	Deep Learning, Model Conversion
Last Updated	2026-02-07 12:00 GMT

Overview

Utility for converting TensorFlow and HuggingFace BERT checkpoints to DeepSpeed format, including Q/K/V weight merging for transformer kernel acceleration.

Description

convert_bert_ckpt_to_deepspeed.py provides conversion functions to load pretrained BERT checkpoints from TensorFlow or HuggingFace format and map them into DeepSpeed's internal model structure. This is essential for using pretrained BERT weights with DeepSpeed's optimized transformer kernels, which use a different weight layout than standard BERT implementations.

The module provides three main conversion pathways. load_tf_weights_in_bert_kernel() loads TensorFlow checkpoints and maps them to DeepSpeed's kernel-optimized format, merging separate query, key, and value weight matrices into fused attn_qkvw and attn_qkvb tensors. load_hf_weights_in_bert_kernel() performs the same Q/K/V fusion for HuggingFace PyTorch checkpoints. load_hf_weights_in_bert_torch() loads HuggingFace checkpoints into DeepSpeed's non-kernel (standard PyTorch) model format, which uses dense_act layers instead of standard dense layers for certain components.

All conversion functions handle vocabulary size alignment: DeepSpeed BERT models require vocabulary sizes aligned to 8, so the conversion pads the word embedding matrix with zeros when the source vocabulary is smaller. The top-level convert_ckpt_to_deepspeed() function dispatches to the appropriate loader based on checkpoint type and kernel configuration.

Usage

Use this module when loading pretrained BERT weights (from TensorFlow or HuggingFace) into a DeepSpeed-optimized BERT model, particularly when using the DeepSpeed transformer kernel for accelerated training.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/BingBertSquad/convert_bert_ckpt_to_deepspeed.py
• Lines: 1-340

Signature

def convert_ckpt_to_deepspeed(model, ckpt_type, ckpt_path, vocab_diff, kernel_enabled):
    ...

def load_tf_weights_in_bert_kernel(model, ckpt_path, voc_size_diff):
    ...

def load_hf_weights_in_bert_kernel(model, ckpt_path, voc_size_diff):
    ...

def load_hf_weights_in_bert_torch(model, ckpt_path, voc_size_diff):
    ...

def set_data(param, array):
    ...

Import

from convert_bert_ckpt_to_deepspeed import convert_ckpt_to_deepspeed

I/O Contract

Inputs

Name	Type	Required	Description
model	nn.Module	Yes	DeepSpeed BERT model instance to populate with converted weights
ckpt_type	str	Yes	Checkpoint format: "HF" for HuggingFace or "TF" for TensorFlow
ckpt_path	str	Yes	Path to the source checkpoint file or directory
vocab_diff	int	Yes	Difference between DeepSpeed vocabulary size and source vocabulary size (for padding alignment)
kernel_enabled	bool	Yes	Whether DeepSpeed transformer kernel is enabled (determines weight layout)

Outputs

Name	Type	Description
model	nn.Module	The same model instance with weights loaded from the source checkpoint

Usage Examples

Convert HuggingFace Checkpoint to DeepSpeed

from convert_bert_ckpt_to_deepspeed import convert_ckpt_to_deepspeed

# With transformer kernel enabled (fused Q/K/V weights)
model = create_deepspeed_bert_model(config)
convert_ckpt_to_deepspeed(
    model=model,
    ckpt_type="HF",
    ckpt_path="/path/to/pytorch_model.bin",
    vocab_diff=2,  # e.g., 30528 - 30522 for 8-aligned vocab
    kernel_enabled=True
)

# Without transformer kernel (standard PyTorch layout)
convert_ckpt_to_deepspeed(
    model=model,
    ckpt_type="HF",
    ckpt_path="/path/to/pytorch_model.bin",
    vocab_diff=2,
    kernel_enabled=False
)