Implementation:Microsoft DeepSpeedExamples PreLN Bert LayerDrop Modeling

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

Overview

Pre-Layer-Normalization BERT model implementation with LayerDrop regularization, DeepSpeed Transformer kernel support, and sparse attention integration for efficient pretraining and fine-tuning.

Description

This module provides a comprehensive Pre-LN BERT implementation with LayerDrop regularization, adapted from the NVIDIA DeepLearningExamples repository. The Pre-LN variant applies layer normalization before the attention and feed-forward sublayers (rather than after), which has been shown to improve training stability for deep transformer models. The LayerDrop mechanism enables stochastic depth by randomly dropping entire transformer layers during training.

The implementation includes all standard BERT components: BertEmbeddings, BertSelfAttention, BertSelfOutput, BertIntermediate, BertOutput, BertLayer, and BertEncoder, along with task-specific heads including BertForPreTrainingPreLN, BertForMaskedLM, BertForSequenceClassification, BertForQuestionAnswering, and others. The BertEncoder supports DeepSpeed Transformer kernels for accelerated training, sparse attention patterns (dense, fixed, bigbird, bslongformer, variable), and gradient checkpointing for memory efficiency.

Key features include JIT-compiled activation functions (fused GELU and tanh with bias), a LinearActivation module that fuses linear projection with activation, and the BertConfig class for model configuration. The progressive LayerDrop is controlled through a theta parameter in the encoder's forward pass, enabling structured regularization during pretraining.

Usage

Use this module when training or fine-tuning Pre-LN BERT models with the BingBertSquad example, particularly when you need LayerDrop regularization for improved generalization, DeepSpeed Transformer kernel acceleration, or sparse attention patterns for long sequences.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/BingBertSquad/turing/modelingpreln_layerdrop.py
• Lines: 1-1652

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):
    def forward(self, batch, **kwargs):

class BertForQuestionAnswering(BertPreTrainedModel):

class BertEncoder(nn.Module):
    def __init__(self, config, args, sparse_attention_config=None):
    def forward(self, hidden_states, attention_mask,
                output_all_encoded_layers=True,
                checkpoint_activations=False,
                progressive_layer_drop=False, theta=0.5):

Import

from turing.modelingpreln_layerdrop import BertForPreTrainingPreLN, BertConfig, BertModel

I/O Contract

Inputs

Name	Type	Required	Description
config	BertConfig	Yes	BERT model configuration with hidden size, layers, heads, etc.
args	Namespace	Yes	Training arguments including DeepSpeed settings and sparse attention config
input_ids	torch.LongTensor	Yes	Token indices of shape [batch_size, sequence_length]
token_type_ids	torch.LongTensor	No	Segment type indices of shape [batch_size, sequence_length]
attention_mask	torch.LongTensor	No	Attention mask of shape [batch_size, sequence_length]
progressive_layer_drop	bool	No	Whether to enable LayerDrop during forward pass
theta	float	No	LayerDrop probability threshold (default: 0.5)

Outputs

Name	Type	Description
total_loss	torch.Tensor	Sum of MLM loss and NSP loss (when labels provided)
prediction_scores	torch.Tensor	MLM logits of shape [batch_size, seq_length, vocab_size]
seq_relationship_score	torch.Tensor	NSP logits of shape [batch_size, 2]
encoded_layers	List[torch.Tensor]	Hidden states from each encoder layer

Usage Examples

from turing.modelingpreln_layerdrop import BertConfig, BertForPreTrainingPreLN

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

# Create model with LayerDrop support
model = BertForPreTrainingPreLN(config, args)

# Forward pass with progressive LayerDrop
output = model(batch, progressive_layer_drop=True, theta=0.5)