Implementation:Microsoft DeepSpeedExamples Nvidia PreLN Bert Modeling

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

Overview

NVIDIA-optimized Pre-Layer-Normalization BERT model implementation with distributed training support, DeepSpeed integration, and task-specific heads for pretraining and fine-tuning.

Description

This module implements the Pre-LN (Pre-Layer-Normalization) variant of BERT based on the NVIDIA DeepLearningExamples codebase, adapted for use with DeepSpeed. In the Pre-LN architecture, layer normalization is applied before the self-attention and feed-forward sublayers rather than after, which improves training stability and enables training deeper models without learning rate warmup issues. The encoder includes a FinalLayerNorm applied to the output of the last transformer layer.

The module provides the full BERT model stack including BertEmbeddings with word, position, and token-type embeddings; BertSelfAttention for multi-head self-attention; BertEncoder with support for gradient checkpointing and DeepSpeed Transformer kernels configured with pre_layer_norm=True; and all standard task-specific heads. The key model class is BertForPreTrainingPreLN which combines the BERT encoder with masked language modeling and next sentence prediction heads.

Performance features include JIT-compiled fused activation functions (bias_gelu, bias_tanh) that combine bias addition with activation in a single kernel, the LinearActivation module for fused linear+activation operations, and support for distributed training via torch.distributed. The module supports loading pretrained weights from TensorFlow checkpoints and HuggingFace-format model archives.

Usage

Use this module for Pre-LN BERT pretraining and fine-tuning in the BingBertSquad pipeline when training stability is important, particularly for deeper models. The Pre-LN variant is preferred over the standard post-LN BERT for large-scale pretraining as it provides more stable gradients and better convergence properties.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/BingBertSquad/turing/nvidia_modelingpreln.py
• Lines: 1-1539

Signature

class BertConfig(object):
    def __init__(self, vocab_size_or_config_json_file, hidden_size=768,
                 num_hidden_layers=12, num_attention_heads=12,
                 intermediate_size=3072, hidden_act="gelu",
                 hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1,
                 max_position_embeddings=512, type_vocab_size=2,
                 initializer_range=0.02):

class BertModel(BertPreTrainedModel):

class BertForPreTrainingPreLN(BertPreTrainedModel):
    def __init__(self, config, args):

class BertEncoder(nn.Module):
    def __init__(self, config, args):

class BertForQuestionAnswering(BertPreTrainedModel):

class LinearActivation(Module):
    def __init__(self, in_features, out_features, act='gelu', bias=True):

Import

from turing.nvidia_modelingpreln import BertForPreTrainingPreLN, BertConfig, BertModel

I/O Contract

Inputs

Name	Type	Required	Description
config	BertConfig	Yes	BERT configuration with model hyperparameters
args	Namespace	Yes	Training arguments with DeepSpeed and kernel configuration
input_ids	torch.LongTensor	Yes	Token indices of shape [batch_size, sequence_length]
token_type_ids	torch.LongTensor	No	Segment type IDs (0 for sentence A, 1 for sentence B)
attention_mask	torch.LongTensor	No	Padding mask of shape [batch_size, sequence_length]
checkpoint_activations	bool	No	Whether to use gradient checkpointing for memory savings

Outputs

Name	Type	Description
total_loss	torch.Tensor	Sum of MLM and NSP losses when labels are provided
prediction_scores	torch.Tensor	Masked language model logits [batch_size, seq_len, vocab_size]
seq_relationship_score	torch.Tensor	Next sentence prediction logits [batch_size, 2]
encoded_layers	List[torch.Tensor]	List of hidden state tensors from each encoder layer

Usage Examples

from turing.nvidia_modelingpreln import BertConfig, BertForPreTrainingPreLN

# Create Pre-LN BERT configuration
config = BertConfig(
    vocab_size_or_config_json_file=30522,
    hidden_size=768,
    num_hidden_layers=12,
    num_attention_heads=12,
    intermediate_size=3072
)

# Initialize Pre-LN model
model = BertForPreTrainingPreLN(config, args)

# Forward pass with gradient checkpointing
input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
attention_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])

outputs = model(input_ids, token_type_ids, attention_mask,
                checkpoint_activations=True)