Implementation:Microsoft DeepSpeedExamples BingBert Turing Models

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	Model Architecture, Multi-Task Learning
Last Updated	2026-02-07 12:00 GMT

Overview

A collection of BERT-based model heads for pretraining, classification, and regression tasks, along with a multi-task orchestrator that manages model initialization, checkpointing, and device placement.

Description

models.py defines three task-specific BERT model heads and a BertMultiTask orchestrator class. BertPretrainingLoss wraps a BERT encoder with BertPreTrainingHeads to compute combined masked language model and next sentence prediction losses using CrossEntropyLoss. BertClassificationLoss adds a dropout layer and linear classifier head on top of the pooled BERT output, using BCEWithLogitsLoss for binary/multi-label classification. BertRegressionLoss similarly adds a linear head but uses MSELoss for regression tasks.

All three model heads extend PreTrainedBertModel and follow a consistent pattern: during training (when labels are provided), they compute and return the loss; during inference (no labels), they return the raw predictions or logits. The classification head supports configurable num_labels for multi-label settings, while the regression head always outputs a single value.

The BertMultiTask class serves as a high-level model manager that handles initialization from either a pretrained BERT checkpoint (using BertModel.from_pretrained) or from scratch (using BertForPreTrainingPreLN with optional progressive layer dropping). It provides convenience methods for save/load checkpoint operations, device placement, eval/train mode switching, and batch movement to the target device.

Usage

Use BertPretrainingLoss for BERT pretraining with masked LM and next sentence prediction. Use BertClassificationLoss and BertRegressionLoss for downstream fine-tuning tasks. Use BertMultiTask as the top-level model wrapper in the Bing BERT Turing training pipeline, which handles model lifecycle management and supports both pretrained and from-scratch initialization.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/bing_bert/turing/models.py
• Lines: 1-163

Signature

class BertPretrainingLoss(PreTrainedBertModel):
    def __init__(self, bert_encoder, config):
    def forward(self, input_ids, token_type_ids=None, attention_mask=None,
                masked_lm_labels=None, next_sentence_label=None):

class BertClassificationLoss(PreTrainedBertModel):
    def __init__(self, bert_encoder, config, num_labels: int = 1):
    def forward(self, input_ids, token_type_ids=None, attention_mask=None,
                labels=None):

class BertRegressionLoss(PreTrainedBertModel):
    def __init__(self, bert_encoder, config):
    def forward(self, input_ids, token_type_ids=None, attention_mask=None,
                labels=None):

class BertMultiTask:
    def __init__(self, args):
    def set_device(self, device):
    def save(self, filename: str):
    def load(self, model_state_dict: str):
    def move_batch(self, batch: TorchTuple, non_blocking=False):
    def eval(self):
    def train(self):
    def save_bert(self, filename: str):
    def to(self, device):
    def half(self):

Import

from turing.models import BertPretrainingLoss, BertClassificationLoss, BertRegressionLoss, BertMultiTask

I/O Contract

Inputs

Name	Type	Required	Description
bert_encoder	BertModel	Yes	Pretrained or initialized BERT encoder module (for task heads)
config	BertConfig	Yes	BERT model configuration specifying hidden size, vocab size, dropout, etc.
num_labels	int	No	Number of output labels for classification. Default: 1
input_ids	Tensor	Yes	Token IDs of shape (batch_size, seq_length)
token_type_ids	Tensor	No	Segment IDs for sentence pair tasks
attention_mask	Tensor	No	Attention mask indicating padding tokens
masked_lm_labels	Tensor	No	Labels for masked LM task; -1 for non-masked positions
next_sentence_label	Tensor	No	Labels for next sentence prediction (0 or 1)
labels	Tensor	No	Labels for classification or regression tasks
args	Namespace	Yes	Training arguments for BertMultiTask (config, use_pretrain, tokenizer, progressive_layer_drop, local_rank)

Outputs

Name	Type	Description
total_loss	Tensor	Combined masked LM and NSP loss (BertPretrainingLoss, training mode)
prediction_scores, seq_relationship_score	tuple(Tensor, Tensor)	Raw predictions (BertPretrainingLoss, inference mode)
loss	Tensor	Classification or regression loss (task heads, training mode)
scores/logits	Tensor	Raw output scores or logits (task heads, inference mode)

Usage Examples

from turing.models import BertPretrainingLoss, BertMultiTask

# Using BertMultiTask for pretraining
model_manager = BertMultiTask(args)
model_manager.to(torch.device("cuda"))
model_manager.train()

# Forward pass with pretraining loss
loss = pretraining_model(
    input_ids=input_ids,
    token_type_ids=segment_ids,
    attention_mask=attention_mask,
    masked_lm_labels=masked_lm_labels,
    next_sentence_label=next_sentence_labels
)
loss.backward()

# Save checkpoint
model_manager.save("checkpoint.pt")

# Load checkpoint
model_manager.load("checkpoint.pt")