Implementation:Microsoft DeepSpeedExamples BingBert NvidiaPreLN LayerDrop

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	BERT Modeling, Regularization
Last Updated	2026-02-07 12:00 GMT

Overview

Pre-LayerNorm BERT model with LayerDrop regularization, extending the NVIDIA modeling variant with stochastic depth for improved generalization.

Description

This module extends the Pre-LayerNorm BERT architecture with LayerDrop, a structured dropout technique that randomly drops entire transformer layers during training. This acts as a form of stochastic depth regularization, encouraging the model to be robust to the removal of individual layers and improving generalization.

The implementation mirrors the standard Pre-LN BERT architecture with the same comprehensive set of model classes (BertModel, BertForPreTrainingPreLN, BertForMaskedLM, BertForSequenceClassification, etc.) but modifies the BertEncoder to probabilistically skip layers during the forward pass. During inference, all layers are used.

Like the base Pre-LN variant, this module includes DeepSpeed sparse attention integration with support for multiple sparsity patterns (dense, fixed, bigbird, bslongformer, variable) and gradient checkpointing. The combination of LayerDrop with sparse attention enables efficient training of very deep BERT models on long sequences.

Usage

Use this module when training deep BERT models where additional regularization through stochastic layer dropping is desired, particularly for large models that may be prone to overfitting or when training with limited data.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/bing_bert/nvidia/modelingpreln_layerdrop.py
• Lines: 1-1662

Signature

class BertConfig(object)
class LinearActivation(Module)
class BertEmbeddings(nn.Module)
class BertSelfAttention(nn.Module)
class BertSelfOutput(nn.Module)
class BertAttention(nn.Module)
class BertIntermediate(nn.Module)
class BertOutput(nn.Module)
class BertLayer(nn.Module)
class BertEncoder(nn.Module)
class BertPooler(nn.Module)
class BertPredictionHeadTransform(nn.Module)
class BertLMPredictionHead(nn.Module)
class BertOnlyMLMHead(nn.Module)
class BertOnlyNSPHead(nn.Module)
class BertPreTrainingHeads(nn.Module)
class BertPreTrainedModel(nn.Module)
class BertModel(BertPreTrainedModel)
class BertForPreTrainingPreLN(BertPreTrainedModel)
class BertForMaskedLM(BertPreTrainedModel)
class BertForNextSentencePrediction(BertPreTrainedModel)
class BertForSequenceClassification(BertPreTrainedModel)
class BertForMultipleChoice(BertPreTrainedModel)
class BertForTokenClassification(BertPreTrainedModel)
class BertForQuestionAnswering(BertPreTrainedModel)

Import

from nvidia.modelingpreln_layerdrop import (
    BertConfig, BertModel, BertForPreTrainingPreLN,
    BertForSequenceClassification, BertForQuestionAnswering,
    get_deepspeed_config, get_sparse_attention_config
)

I/O Contract

Inputs

Name	Type	Required	Description
config	BertConfig	Yes	Configuration object with model hyperparameters and LayerDrop probability
input_ids	Tensor	Yes	Token IDs of shape (batch_size, seq_length)
token_type_ids	Tensor	No	Segment IDs for sentence A vs sentence B distinction
attention_mask	Tensor	No	Mask indicating which positions to attend to
args	Namespace	No	Arguments containing deepspeed_config and sparse attention flags

Outputs

Name	Type	Description
sequence_output	Tensor	Hidden states from the last encoder layer of shape (batch_size, seq_length, hidden_size)
pooled_output	Tensor	Pooled [CLS] representation of shape (batch_size, hidden_size)
prediction_scores	Tensor	MLM logits of shape (batch_size, seq_length, vocab_size) for pretraining
seq_relationship_score	Tensor	NSP logits of shape (batch_size, 2) for pretraining

Usage Examples

from nvidia.modelingpreln_layerdrop import BertConfig, BertForPreTrainingPreLN

config = BertConfig(
    vocab_size_or_config_json_file=30522,
    hidden_size=1024,
    num_hidden_layers=24,
    num_attention_heads=16,
    intermediate_size=4096
)

model = BertForPreTrainingPreLN(config)
# During training, layers are randomly dropped
model.train()
prediction_scores, seq_relationship_score = model(
    input_ids=input_ids,
    token_type_ids=segment_ids,
    attention_mask=input_mask
)