Implementation:Microsoft DeepSpeedExamples ConvertTfCheckpointToPytorch

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

Overview

A conversion script that loads a TensorFlow BERT checkpoint and translates its weights into a PyTorch state dictionary, handling variable name mapping and tensor transposition.

Description

convert_tf_checkpoint_to_pytorch.py provides a convert_tf_checkpoint_to_pytorch function that bridges the gap between TensorFlow and PyTorch BERT models. It loads all variables from a TensorFlow checkpoint using tf.train.list_variables and tf.train.load_variable, then maps each variable to the corresponding PyTorch model attribute through a systematic name translation process.

The name mapping logic handles several TensorFlow-to-PyTorch naming conventions: kernel and gamma map to weight, output_bias and beta map to bias, and output_weights maps to weight. Variables with numeric suffixes (e.g., layer_0) are parsed and used as indices into PyTorch module lists. Special cases include transposing kernel tensors (to account for TF vs PyTorch convention) and handling _embeddings suffixed variables. Adam optimizer variables (adam_v, adam_m) and global_step are intentionally skipped.

The script accepts three required command-line arguments: the TensorFlow checkpoint path, a BERT config JSON file, and the output PyTorch model path. It constructs a BertForPreTraining model from the config, assigns the converted weights, and saves the resulting state dictionary.

Usage

Use this script when you have a pretrained BERT model in TensorFlow checkpoint format and need to convert it for use with the PyTorch-based Bing BERT training pipeline. This is a one-time conversion step typically performed before fine-tuning or continued pretraining.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/bing_bert/pytorch_pretrained_bert/convert_tf_checkpoint_to_pytorch.py
• Lines: 1-117

Signature

def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, bert_config_file,
                                     pytorch_dump_path):

Import

from pytorch_pretrained_bert.convert_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch

I/O Contract

Inputs

Name	Type	Required	Description
tf_checkpoint_path	str	Yes	Path to the TensorFlow checkpoint file (without file extension)
bert_config_file	str	Yes	Path to the BERT configuration JSON file specifying model architecture
pytorch_dump_path	str	Yes	Output path for the converted PyTorch model state dictionary

Outputs

Name	Type	Description
pytorch model file	file	Serialized PyTorch state dictionary saved via torch.save at the specified dump path

Usage Examples

# Command-line usage
# python convert_tf_checkpoint_to_pytorch.py \
#     --tf_checkpoint_path /path/to/bert_model.ckpt \
#     --bert_config_file /path/to/bert_config.json \
#     --pytorch_dump_path /path/to/pytorch_model.bin

# Programmatic usage
from pytorch_pretrained_bert.convert_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch

convert_tf_checkpoint_to_pytorch(
    tf_checkpoint_path="/path/to/bert_model.ckpt",
    bert_config_file="/path/to/bert_config.json",
    pytorch_dump_path="/path/to/pytorch_model.bin"
)