Upload BertForSequenceClassification

bert_layers.py

@@ -0,0 +1,1101 @@
+# Copyright 2022 MosaicML Examples authors
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018-2021, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2022, Tri Dao.
+
+"""Implements Mosaic BERT, with an eye towards the Hugging Face API.
+
+Mosaic BERT improves performance over Hugging Face BERT through the following:
+
+1. ALiBi. This architectural change removes positional embeddings and instead encodes positional
+information through attention biases based on query-key position distance. It improves the effectiveness
+of training with shorter sequence lengths by enabling extrapolation to longer sequences.
+
+2. Gated Linear Units (GLU). This architectural change replaces the FFN component of the BERT layer
+to improve overall expressiveness, providing better convergence properties.
+
+3. Flash Attention. The Mosaic BERT's self-attention layer makes use of Flash Attention, which dramatically
+improves the speed of self-attention. Our implementation utilizes a bleeding edge implementation that
+supports attention biases, which allows us to use Flash Attention with ALiBi.
+
+4. Unpadding. Padding is often used to simplify batching across sequences of different lengths. Standard BERT
+implementations waste computation on padded tokens. Mosaic BERT internally unpads to reduce unnecessary computation
+and improve speed. It does this without changing how the user interfaces with the model, thereby
+preserving the simple API of standard implementations.
+
+
+Currently, Mosaic BERT is available for masked language modeling :class:`BertForMaskedLM` and sequence
+classification :class:`BertForSequenceClassification`. We aim to expand this catalogue in future releases.
+
+See :file:`./mosaic_bert.py` for utilities to simplify working with Mosaic BERT in Composer, and for example usage
+of the core Mosaic BERT classes.
+"""
+
+import copy
+import logging
+import math
+import os
+import sys
+import warnings
+from typing import List, Optional, Tuple, Union
+from .configuration_bert import BertConfig
+# Add folder root to path to allow us to use relative imports regardless of what directory the script is run from
+sys.path.append(os.path.dirname(os.path.realpath(__file__)))
+
+from .bert_padding import (index_first_axis,
+                                            index_put_first_axis, pad_input,
+                                            unpad_input, unpad_input_only)
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from einops import rearrange
+from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (MaskedLMOutput,
+                                           SequenceClassifierOutput)
+from transformers.models.bert.modeling_bert import BertPreTrainedModel
+logger = logging.getLogger(__name__)
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        RMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+
+        # convert into half-precision if necessary
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            hidden_states = hidden_states.to(self.weight.dtype)
+
+        return self.weight * hidden_states
+    
+class RotaryEmbedding(torch.nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+        super().__init__()
+        self.inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
+        self.max_seq_len_cached = max_position_embeddings
+        t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32)
+        self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32)
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
+        if seq_len > self.max_seq_len_cached:
+            self.max_seq_len_cached = seq_len
+            t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32)
+            freqs = torch.outer(t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32).to(x.device)
+            self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32).to(x.device)
+        elif self.cos_cached.device != x.device:
+            self.cos_cached = self.cos_cached.to(x.device)
+            self.sin_cached = self.sin_cached.to(x.device)  
+        return (
+            self.cos_cached[:, :, :seq_len, ...],
+            self.sin_cached[:, :, :seq_len, ...],
+        )
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos_, sin_):
+    #cos = cos_.squeeze(1).squeeze(0)  # [seq_len, dim]
+    #sin = sin_.squeeze(1).squeeze(0)  # [seq_len, dim]
+    cos = torch.repeat_interleave(cos_[:, :, None, :], q.shape[0], 0).squeeze(1)
+    sin = torch.repeat_interleave(sin_[:, :, None, :], q.shape[0], 0).squeeze(1)
+    #position_ids = torch.Tensor([list(range(q.shape[2]))]*q.shape[0]).int().to(q.device)
+    #cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
+    #sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
+    q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
+    k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
+    return q_embed.to(q.dtype), k_embed.to(k.dtype)
+
+class BertEmbeddings(nn.Module):
+    """Construct the embeddings for words, ignoring position.
+
+    There are no positional embeddings since we use ALiBi and token_type
+    embeddings.
+
+    This module is modeled after the Hugging Face BERT's
+    :class:`~transformers.model.bert.modeling_bert.BertEmbeddings`, but is
+    modified as part of Mosaic BERT's ALiBi implementation. The key change is
+    that position embeddings are removed. Position information instead comes
+    from attention biases that scale linearly with the position distance
+    between query and key tokens.
+
+    This module ignores the `position_ids` input to the `forward` method.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size,
+                                            config.hidden_size,
+                                            padding_idx=config.pad_token_id)
+        # ALiBi doesn't use position embeddings
+        self.token_type_embeddings = nn.Embedding(config.type_vocab_size,
+                                                  config.hidden_size)
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model
+        # variable name and be able to load any TensorFlow checkpoint file
+        self.norm = RMSNorm(config.hidden_size,
+                                      eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.register_buffer('token_type_ids',
+                             torch.zeros(config.max_position_embeddings,
+                                         dtype=torch.long),
+                             persistent=False)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values_length: int = 0,
+    ) -> torch.Tensor:
+        if (input_ids is not None) == (inputs_embeds is not None):
+            raise ValueError('Must specify either input_ids or input_embeds!')
+        if input_ids is not None:
+            input_shape = input_ids.size()
+        else:
+            assert inputs_embeds is not None  # just for type checking
+            input_shape = inputs_embeds.size()[:-1]
+
+        seq_length = input_shape[1]
+
+        if position_ids is None:
+            # great! ALiBi
+            pass
+
+        # Setting the token_type_ids to the registered buffer in constructor
+        # where it is all zeros, which usually occurs when it's auto-generated;
+        # registered buffer helps users when tracing the model without passing
+        # token_type_ids, solves issue #5664
+        if token_type_ids is None:
+            if hasattr(self, 'token_type_ids'):
+                assert isinstance(self.token_type_ids, torch.LongTensor)
+                buffered_token_type_ids = self.token_type_ids[:, :seq_length]
+                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(
+                    input_shape[0], seq_length)
+                token_type_ids = buffered_token_type_ids_expanded  # type: ignore
+            else:
+                token_type_ids = torch.zeros(input_shape,  # type: ignore
+                                             dtype=torch.long,
+                                             device=self.word_embeddings.device) # type: ignore  # yapf: disable
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+        embeddings = inputs_embeds + token_type_embeddings
+        # no position embeddings! ALiBi
+        embeddings = self.norm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class BertUnpadSelfAttention(nn.Module):
+    """Performs multi-headed self attention on a batch of unpadded sequences.
+
+    If Triton is installed, this module uses Flash Attention to greatly improve throughput.
+    The Flash Attention implementation used in Mosaic BERT supports arbitrary attention biases (which
+    we use to implement ALiBi), but does not support attention dropout. If either Triton is not installed
+    or `config.attention_probs_dropout_prob > 0`, the implementation will default to a
+    math-equivalent pytorch version, which is much slower.
+
+    See `forward` method for additional detail.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
+                config, 'embedding_size'):
+            raise ValueError(
+                f'The hidden size ({config.hidden_size}) is not a multiple of the number of attention '
+                f'heads ({config.num_attention_heads})')
+        
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size /
+                                       config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+        self.p_dropout = config.attention_probs_dropout_prob
+        self.Wqkv = nn.Linear(self.all_head_size, 3 * config.hidden_size)
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rotary_emb = RotaryEmbedding(self.attention_head_size, max_position_embeddings=self.max_position_embeddings)
+        # Warn if defaulting to pytorch because of import issues
+
+
+    def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor,
+                max_seqlen_in_batch: int, indices: torch.Tensor,
+                attn_mask: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+        """Perform self-attention.
+
+        If dropout is zero, then we can use the Triton kernel, so we do that. However, if not, we send through a standard PyTorch
+        implementation of self-attention.
+
+        The arguments are unpadded, and our implementations of attention require padded arguments,
+        so we first call `pad_input`. Once we compute attention, we re-unpad our outputs for the other layers.
+        The pad/unpad operations add overhead, but not sending pad tokens through ffs saves compute.
+        It is possible to write an unpadded implementation of attention (in Triton and PyTorch), which we will eventually do.
+
+        Args:
+            hidden_states: (total_nnz, dim)
+            cu_seqlens: (batch + 1,)
+            max_seqlen_in_batch: int
+            indices: (total_nnz,)
+            attn_mask: (batch, max_seqlen_in_batch)
+            bias: (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
+
+        Returns:
+            attention: (total_nnz, dim)
+        """
+        qkv = self.Wqkv(hidden_states)
+        qkv = pad_input(
+            qkv, indices, cu_seqlens.shape[0] - 1,
+            max_seqlen_in_batch)  # batch, max_seqlen_in_batch, thd
+        qkv = rearrange(qkv,
+                        'b s (t h d) -> b s t h d',
+                        t=3,
+                        h=self.num_attention_heads)
+        # if we have nonzero attention dropout (e.g. during fine-tuning) or no Triton, compute attention in PyTorch
+        q = qkv[:, :, 0, :, :].transpose(1, 2)
+        k = qkv[:, :, 1, :, :].transpose(1, 2)
+        v = qkv[:, :, 2, :, :].transpose(1, 2)
+        kv_seq_len = k.shape[-2]
+
+        cos, sin = self.rotary_emb(v, seq_len=kv_seq_len)
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+        #q = q.transpose(1, 2)
+        k = k.permute(0, 1, 3, 2)
+        #v = v.transpose(1, 2)
+        # q = qkv[:, :, 0, :, :].permute(0, 2, 1, 3)  # b h s d
+        # k = qkv[:, :, 1, :, :].permute(0, 2, 3, 1)  # b h d s
+        # v = qkv[:, :, 2, :, :].permute(0, 2, 1, 3)  # b h s d
+
+        attention_scores = torch.matmul(q, k) / math.sqrt(
+            self.attention_head_size)
+        attention_scores = attention_scores + bias
+        attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+        attention_probs = self.dropout(attention_probs)
+        attention = torch.matmul(attention_probs, v).permute(0, 2, 1,
+                                                             3)  # b s h d
+        
+        # attn_mask is 1 for attend and 0 for don't
+        attention = unpad_input_only(
+            attention,
+            torch.squeeze(attn_mask) == 1)
+        return rearrange(attention, 'nnz h d -> nnz (h d)')
+
+
+# Copy of transformer's library BertSelfOutput that will not be caught by surgery methods looking for HF BERT modules.
+class BertSelfOutput(nn.Module):
+    """Computes the output of the attention layer.
+
+    This module is modeled after the Hugging Face BERT's
+    :class:`~transformers.model.bert.modeling_bert.BertSelfOutput`.
+    The implementation is identical. Rather than use the original module
+    directly, we re-implement it here so that Mosaic BERT's modules will not
+    be affected by any Composer surgery algorithm that modifies Hugging Face
+    BERT modules.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.norm = RMSNorm(config.hidden_size,
+                                      eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor,
+                input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.norm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertUnpadAttention(nn.Module):
+    """Chains attention, Dropout, and LayerNorm for Mosaic BERT."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.self = BertUnpadSelfAttention(config)
+        self.output = BertSelfOutput(config)
+
+    def forward(
+        self,
+        input_tensor: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_s: int,
+        subset_idx: Optional[torch.Tensor] = None,
+        indices: Optional[torch.Tensor] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward pass for scaled self-attention without padding.
+
+        Arguments:
+            input_tensor: (total_nnz, dim)
+            cu_seqlens: (batch + 1,)
+            max_s: int
+            subset_idx: () set of indices whose values we care about at the end of the layer
+                        (e.g., the masked tokens, if this is the final layer).
+            indices: None or (total_nnz,)
+            attn_mask: None or (batch, max_seqlen_in_batch)
+            bias: None or (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
+        """
+        self_output = self.self(input_tensor, cu_seqlens, max_s, indices,
+                                attn_mask, bias)
+        if subset_idx is not None:
+            return self.output(
+                index_first_axis(self_output, subset_idx),
+                index_first_axis(input_tensor, subset_idx))
+        else:
+            return self.output(self_output, input_tensor)
+
+class MLP(nn.Module):
+    def __init__(
+            self,
+            config
+    ):
+        super().__init__()
+        self.config = config
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+        self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states):
+        residual_connection = hidden_states
+        hidden_states = self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
+        hidden_states = self.norm(hidden_states + residual_connection)
+        return hidden_states
+    
+# class BertGatedLinearUnitMLP(nn.Module):
+#     """Applies the FFN at the end of each Mosaic BERT layer.
+
+#     Compared to the default BERT architecture, this block replaces :class:`~transformers.model.bert.modeling_bert.BertIntermediate`
+#     and :class:`~transformers.model.bert.modeling_bert.SelfOutput` with a single module that has similar functionality, but
+#     introduces Gated Linear Units.
+
+#     Note: Mosaic BERT adds parameters in order to implement Gated Linear Units. To keep parameter count consistent with that of a
+#     standard Hugging Face BERT, scale down `config.intermediate_size` by 2/3. For example, a Mosaic BERT constructed with
+#     `config.intermediate_size=2048` will have the same parameter footprint as its Hugging Face BERT counterpart constructed
+#     with the `config.intermediate_size=3072`.
+#     However, in most cases it will not be necessary to adjust `config.intermediate_size` since, despite the increased
+#     parameter size, Mosaic BERT typically offers a net higher throughput than a Hugging Face BERT built from the same `config`.
+#     """
+
+#     def __init__(self, config):
+#         super().__init__()
+#         self.config = config
+#         self.gated_layers = nn.Linear(config.hidden_size,
+#                                       config.intermediate_size * 2,
+#                                       bias=False)
+#         self.act = ACT2FN[config.hidden_act]#nn.GELU(approximate='none')
+#         self.wo = nn.Linear(config.intermediate_size, config.hidden_size)
+#         self.dropout = nn.Dropout(config.hidden_dropout_prob)
+#         self.norm = RMSNorm(config.hidden_size,
+#                                       eps=config.layer_norm_eps)
+
+#     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+#         """Compute new hidden states from current hidden states.
+
+#         Args:
+#             hidden_states (torch.Tensor): The (unpadded) hidden states from
+#                 the attention layer [nnz, dim].
+#         """
+#         residual_connection = hidden_states
+#         # compute the activation
+#         hidden_states = self.gated_layers(hidden_states)
+#         gated = hidden_states[:, :self.config.intermediate_size]
+#         non_gated = hidden_states[:, self.config.intermediate_size:]
+#         hidden_states = self.act(gated) * non_gated
+#         hidden_states = self.dropout(hidden_states)
+#         # multiply by the second matrix
+#         hidden_states = self.wo(hidden_states)
+#         # add the residual connection and post-LN
+#         hidden_states = self.norm(hidden_states + residual_connection)
+#         return hidden_states
+
+
+class BertLayer(nn.Module):
+    """Composes the Mosaic BERT attention and FFN blocks into a single layer."""
+
+    def __init__(self, config):
+        super(BertLayer, self).__init__()
+        self.attention = BertUnpadAttention(config)
+        self.mlp = MLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        seqlen: int,
+        subset_idx: Optional[torch.Tensor] = None,
+        indices: Optional[torch.Tensor] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward pass for a BERT layer, including both attention and MLP.
+
+        Args:
+            hidden_states: (total_nnz, dim)
+            cu_seqlens: (batch + 1,)
+            seqlen: int
+            subset_idx: () set of indices whose values we care about at the end of the layer
+                        (e.g., the masked tokens, if this is the final layer).
+            indices: None or (total_nnz,)
+            attn_mask: None or (batch, max_seqlen_in_batch)
+            bias: None or (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
+        """
+        attention_output = self.attention(hidden_states, cu_seqlens, seqlen,
+                                          subset_idx, indices, attn_mask, bias)
+        layer_output = self.mlp(attention_output)
+        return layer_output
+
+
+class BertEncoder(nn.Module):
+    """A stack of BERT layers providing the backbone of Mosaic BERT.
+
+    This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.modeling_bert.BertEncoder`,
+    but with substantial modifications to implement unpadding and ALiBi.
+
+    Compared to the analogous Hugging Face BERT module, this module handles unpadding to reduce unnecessary computation
+    at padded tokens, and pre-computes attention biases to implement ALiBi.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        layer = BertLayer(config)
+        self.layer = nn.ModuleList(
+            [copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])
+
+        self.num_attention_heads = config.num_attention_heads
+
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        output_all_encoded_layers: Optional[bool] = True,
+        subset_mask: Optional[torch.Tensor] = None,
+    ) -> List[torch.Tensor]:
+
+        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+        extended_attention_mask = extended_attention_mask.to(
+            dtype=next(self.parameters()).dtype)  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+
+        attention_mask_bool = attention_mask.bool()
+        batch, seqlen = hidden_states.shape[:2]
+        # Unpad inputs and mask. It will remove tokens that are padded.
+        # Assume ntokens is total number of tokens (padded and non-padded)
+        # and ntokens_unpad is total number of non-padded tokens.
+        # Then unpadding performs the following compression of the inputs:
+        # hidden_states[ntokens,hidden] -> hidden_states[ntokens_unpad,hidden]
+        hidden_states, indices, cu_seqlens, _ = unpad_input(
+            hidden_states, attention_mask_bool)
+
+        attn_bias = extended_attention_mask[:, :, :seqlen, :seqlen]
+        all_encoder_layers = []
+        if subset_mask is None:
+            for layer_module in self.layer:
+                hidden_states = layer_module(hidden_states,
+                                             cu_seqlens,
+                                             seqlen,
+                                             None,
+                                             indices,
+                                             attn_mask=attention_mask,
+                                             bias=attn_bias)
+                if output_all_encoded_layers:
+                    all_encoder_layers.append(hidden_states)
+            # Pad inputs and mask. It will insert back zero-padded tokens.
+            # Assume ntokens is total number of tokens (padded and non-padded)
+            # and ntokens_unpad is total number of non-padded tokens.
+            # Then padding performs the following de-compression:
+            #     hidden_states[ntokens_unpad,hidden] -> hidden_states[ntokens,hidden]
+            hidden_states = pad_input(
+                hidden_states, indices, batch, seqlen)
+        else:
+            for i in range(len(self.layer) - 1):
+                layer_module = self.layer[i]
+                hidden_states = layer_module(hidden_states,
+                                             cu_seqlens,
+                                             seqlen,
+                                             None,
+                                             indices,
+                                             attn_mask=attention_mask,
+                                             bias=attn_bias)
+                if output_all_encoded_layers:
+                    all_encoder_layers.append(hidden_states)
+            subset_idx = torch.nonzero(subset_mask[attention_mask_bool],
+                                       as_tuple=False).flatten()
+            hidden_states = self.layer[-1](hidden_states,
+                                           cu_seqlens,
+                                           seqlen,
+                                           subset_idx=subset_idx,
+                                           indices=indices,
+                                           attn_mask=attention_mask,
+                                           bias=attn_bias)
+
+        if not output_all_encoded_layers:
+            all_encoder_layers.append(hidden_states)
+        return all_encoder_layers
+
+
+class BertPooler(nn.Module):
+
+    def __init__(self, config):
+        super(BertPooler, self).__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                pool: Optional[bool] = True) -> torch.Tensor:
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0] if pool else hidden_states
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class BertPredictionHeadTransform(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        if isinstance(config.hidden_act, str):
+            self.transform_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.norm = RMSNorm(config.hidden_size, eps=1e-12)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.norm(hidden_states)
+        return hidden_states
+
+
+class BertModel(BertPreTrainedModel):
+    """Overall BERT model.
+
+    Args:
+        config: a BertConfig class instance with the configuration to build a new model
+
+    Inputs:
+        `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
+            with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
+            `extract_features.py`, `run_classifier.py` and `run_squad.py`)
+        `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
+            types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
+            a `sentence B` token (see BERT paper for more details).
+        `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
+            selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
+            input sequence length in the current batch. It's the mask that we typically use for attention when
+            a batch has varying length sentences.
+        `output_all_encoded_layers`: boolean which controls the content of the `encoded_layers` output as described below. Default: `True`.
+
+    Outputs: Tuple of (encoded_layers, pooled_output)
+        `encoded_layers`: controlled by `output_all_encoded_layers` argument:
+            - `output_all_encoded_layers=True`: outputs a list of the full sequences of encoded-hidden-states at the end
+                of each attention block (i.e. 12 full sequences for BERT-base, 24 for BERT-large), each
+                encoded-hidden-state is a torch.FloatTensor of size [batch_size, sequence_length, hidden_size],
+            - `output_all_encoded_layers=False`: outputs only the full sequence of hidden-states corresponding
+                to the last attention block of shape [batch_size, sequence_length, hidden_size],
+        `pooled_output`: a torch.FloatTensor of size [batch_size, hidden_size] which is the output of a
+            classifier pretrained on top of the hidden state associated to the first character of the
+            input (`CLS`) to train on the Next-Sentence task (see BERT's paper).
+
+    Example usage:
+    ```python
+    # Already been converted into WordPiece token ids
+    input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
+    input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
+    token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
+    config = modeling.BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
+        num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
+    model = BertModel(config=config)
+    all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask)
+    ```
+    """
+
+    def __init__(self, config, add_pooling_layer=True):
+        super(BertModel, self).__init__(config)
+        self.embeddings = BertEmbeddings(config)
+        self.encoder = BertEncoder(config)
+        self.pooler = BertPooler(config) if add_pooling_layer else None
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        token_type_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_all_encoded_layers: Optional[bool] = False,
+        masked_tokens_mask: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> Tuple[Union[List[torch.Tensor], torch.Tensor], Optional[torch.Tensor]]:
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+
+        embedding_output = self.embeddings(input_ids, token_type_ids,
+                                           position_ids)
+
+        subset_mask = []
+        first_col_mask = []
+
+        if masked_tokens_mask is None:
+            subset_mask = None
+        else:
+            first_col_mask = torch.zeros_like(masked_tokens_mask)
+            first_col_mask[:, 0] = True
+            subset_mask = masked_tokens_mask | first_col_mask
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask,
+            output_all_encoded_layers=output_all_encoded_layers,
+            subset_mask=subset_mask)
+
+        if masked_tokens_mask is None:
+            sequence_output = encoder_outputs[-1]
+            pooled_output = self.pooler(
+                sequence_output) if self.pooler is not None else None
+        else:
+            # TD [2022-03-01]: the indexing here is very tricky.
+            attention_mask_bool = attention_mask.bool()
+            subset_idx = subset_mask[attention_mask_bool]  # type: ignore
+            sequence_output = encoder_outputs[-1][
+                masked_tokens_mask[attention_mask_bool][subset_idx]]
+            if self.pooler is not None:
+                pool_input = encoder_outputs[-1][
+                    first_col_mask[attention_mask_bool][subset_idx]]
+                pooled_output = self.pooler(pool_input, pool=False)
+            else:
+                pooled_output = None
+
+        if not output_all_encoded_layers:
+            encoder_outputs = sequence_output
+
+        if self.pooler is not None:
+            return encoder_outputs, pooled_output
+
+        return encoder_outputs, None
+
+
+###################
+# Bert Heads
+###################
+class BertLMPredictionHead(nn.Module):
+
+    def __init__(self, config, bert_model_embedding_weights):
+        super().__init__()
+        self.transform = BertPredictionHeadTransform(config)
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.weight = nn.Parameter(torch.empty((bert_model_embedding_weights.size(0), bert_model_embedding_weights.size(1))))
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        self.first_flag = True
+        
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.transform(hidden_states)
+        if self.training:
+            norm_weight = nn.functional.normalize(self.weight)
+            self.first_flag = True
+        elif self.first_flag:
+            self.first_flag = False
+            self.weight.data = nn.functional.normalize(self.weight)
+            norm_weight = self.weight
+        else:
+            norm_weight = self.weight
+        return nn.functional.linear(hidden_states, norm_weight)
+
+
+class BertOnlyMLMHead(nn.Module):
+
+    def __init__(self, config, bert_model_embedding_weights):
+        super().__init__()
+        self.predictions = BertLMPredictionHead(config,
+                                                bert_model_embedding_weights)
+
+    def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
+        prediction_scores = self.predictions(sequence_output)
+        return prediction_scores
+
+
+class BertOnlyNSPHead(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.seq_relationship = nn.Linear(config.hidden_size, 2)
+
+    def forward(self, pooled_output: torch.Tensor) -> torch.Tensor:
+        seq_relationship_score = self.seq_relationship(pooled_output)
+        return seq_relationship_score
+
+
+#####################
+# Various Bert models
+#####################
+
+
+class BertForPreTraining(BertPreTrainedModel):
+    #TBD: Coming in Future Commit
+    pass
+
+
+class BertLMHeadModel(BertPreTrainedModel):
+    #TBD: Coming in Future Commit
+    pass
+
+
+class BertForMaskedLM(BertPreTrainedModel):
+    config_class = BertConfig
+    def __init__(self, config):
+        super().__init__(config)
+
+        if config.is_decoder:
+            warnings.warn(
+                'If you want to use `BertForMaskedLM` make sure `config.is_decoder=False` for '
+                'bi-directional self-attention.')
+        self.config = config
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.cls = BertOnlyMLMHead(config,
+                                   self.bert.embeddings.word_embeddings.weight)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @classmethod
+    def from_composer(cls,
+                      pretrained_checkpoint,
+                      state_dict=None,
+                      cache_dir=None,
+                      from_tf=False,
+                      config=None,
+                      *inputs,
+                      **kwargs):
+        """Load from pre-trained."""
+        model = cls(config, *inputs, **kwargs)
+        if from_tf:
+            raise ValueError(
+                'Mosaic BERT does not support loading TensorFlow weights.')
+
+        state_dict = torch.load(pretrained_checkpoint)
+        # If the state_dict was saved after wrapping with `composer.HuggingFaceModel`, it takes on the `model` prefix
+        consume_prefix_in_state_dict_if_present(state_dict, prefix='model.')
+        missing_keys, unexpected_keys = model.load_state_dict(state_dict,
+                                                              strict=False)
+
+        if len(missing_keys) > 0:
+            logger.warning(
+                f"Found these missing keys in the checkpoint: {', '.join(missing_keys)}"
+            )
+        if len(unexpected_keys) > 0:
+            logger.warning(
+                f"Found these unexpected keys in the checkpoint: {', '.join(unexpected_keys)}"
+            )
+
+        return model
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.weight
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.weight = new_embeddings
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
+        # labels should be a `torch.LongTensor` of shape
+        # `(batch_size, sequence_length)`. These are used for computing the
+        #  masked language modeling loss.
+        #
+        # Indices should be in `[-100, 0, ..., config.vocab_size]` (see
+        # `input_ids` docstring) Tokens with indices set to `-100` are ignored
+        # (masked), the loss is only computed for the tokens with labels in `[0,
+        # ..., config.vocab_size]`
+        #
+        # Prediction scores are only computed for masked tokens and the (bs,
+        # seqlen) dimensions are flattened
+        if (input_ids is not None) == (inputs_embeds is not None):
+            raise ValueError('Must specify either input_ids or input_embeds!')
+
+        if labels is None:
+            masked_tokens_mask = None
+        else:
+            masked_tokens_mask = labels > 0
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            masked_tokens_mask=masked_tokens_mask,
+        )
+
+        sequence_output = outputs[0]
+        prediction_scores = self.cls(sequence_output)
+
+        loss = None
+        if labels is not None:
+            # Compute loss
+            loss_fct = nn.CrossEntropyLoss()
+            softmax_normalizer = prediction_scores.max(-1).values ** 2
+            z_loss_weight = 0.2
+            z_loss = z_loss_weight * softmax_normalizer.mean()
+            # Enable model parallelism
+            masked_token_idx = torch.nonzero(labels.flatten() > 0,
+                                             as_tuple=False).flatten()
+            
+            loss = loss_fct(prediction_scores,
+                            labels.flatten()[masked_token_idx]) + z_loss
+            assert input_ids is not None, 'Coding error; please open an issue'
+            batch, seqlen = input_ids.shape[:2]
+            prediction_scores = rearrange(
+                index_put_first_axis(
+                    prediction_scores, masked_token_idx, batch * seqlen),
+                '(b s) d -> b s d',
+                b=batch)
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MaskedLMOutput(
+            loss=loss,
+            logits=prediction_scores,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids: torch.Tensor,
+                                      attention_mask: torch.Tensor,
+                                      **model_kwargs):
+        input_shape = input_ids.shape
+        effective_batch_size = input_shape[0]
+
+        #  add a dummy token
+        if self.config.pad_token_id is None:
+            raise ValueError('The PAD token should be defined for generation')
+
+        attention_mask = torch.cat([
+            attention_mask,
+            attention_mask.new_zeros((attention_mask.shape[0], 1))
+        ],
+                                   dim=-1)
+        dummy_token = torch.full((effective_batch_size, 1),
+                                 self.config.pad_token_id,
+                                 dtype=torch.long,
+                                 device=input_ids.device)
+        input_ids = torch.cat([input_ids, dummy_token], dim=1)
+
+        return {'input_ids': input_ids, 'attention_mask': attention_mask}
+
+
+class BertForNextSentencePrediction(BertPreTrainedModel):
+    #TBD: Push in future commit
+    pass
+
+
+class BertForSequenceClassification(BertPreTrainedModel):
+    """Bert Model transformer with a sequence classification/regression head.
+
+    This head is just a linear layer on top of the pooled output. Used for,
+    e.g., GLUE tasks.
+    """
+    config_class = BertConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+
+        self.bert = BertModel(config)
+        classifier_dropout = (config.classifier_dropout
+                              if config.classifier_dropout is not None else
+                              config.hidden_dropout_prob)
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @classmethod
+    def from_composer(cls,
+                      pretrained_checkpoint,
+                      state_dict=None,
+                      cache_dir=None,
+                      from_tf=False,
+                      config=None,
+                      *inputs,
+                      **kwargs):
+        """Load from pre-trained."""
+        model = cls(config, *inputs, **kwargs)
+        if from_tf:
+            raise ValueError(
+                'Mosaic BERT does not support loading TensorFlow weights.')
+
+        state_dict = torch.load(pretrained_checkpoint)
+        # If the state_dict was saved after wrapping with `composer.HuggingFaceModel`, it takes on the `model` prefix
+        consume_prefix_in_state_dict_if_present(state_dict, prefix='model.')
+        missing_keys, unexpected_keys = model.load_state_dict(state_dict,
+                                                              strict=False)
+
+        if len(missing_keys) > 0:
+            logger.warning(
+                f"Found these missing keys in the checkpoint: {', '.join(missing_keys)}"
+            )
+        if len(unexpected_keys) > 0:
+            logger.warning(
+                f"Found these unexpected keys in the checkpoint: {', '.join(unexpected_keys)}"
+            )
+
+        return model
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
+        # labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+        # Labels for computing the sequence classification/regression loss.
+        # Indices should be in `[0, ..., config.num_labels - 1]`.
+        # If `config.num_labels == 1` a regression loss is computed
+        # (mean-square loss). If `config.num_labels > 1` a classification loss
+        # is computed (cross-entropy).
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        pooled_output = outputs[1]
+
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            # Compute loss
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = 'regression'
+                elif self.num_labels > 1 and (labels.dtype == torch.long or
+                                              labels.dtype == torch.int):
+                    self.config.problem_type = 'single_label_classification'
+                else:
+                    self.config.problem_type = 'multi_label_classification'
+
+            if self.config.problem_type == 'regression':
+                loss_fct = nn.MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == 'single_label_classification':
+                loss_fct = nn.CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels),
+                                labels.view(-1))
+            elif self.config.problem_type == 'multi_label_classification':
+                loss_fct = nn.BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=None,
+            attentions=None,
+        )
+
+
+class BertForMultipleChoice(BertPreTrainedModel):
+    #TBD: Push in future commit
+    pass
+
+
+class BertForTokenClassification(BertPreTrainedModel):
+    #TBD: Push in future commit
+    pass
+
+
+class BertForQuestionAnswering(BertPreTrainedModel):
+    """Bert Model with a span classification head.
+
+    This is used for extractive question-answering tasks like SQuAD (a linear
+    layers on top of the hidden states' output to compute `span start logits`
+    and `span end logits`).
+    """
+    #TBD: Push in future commit

bert_padding.py

@@ -0,0 +1,159 @@
+# Copyright 2022 MosaicML Examples authors
+# SPDX-License-Identifier: Apache-2.0
+
+# Adapted from https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/bert_padding.py
+# Which was adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py
+
+"""Helper functions for padding and unpadding batches.
+
+These functions are used extensively throughout the Mosaic BERT implementation
+in `bert_layers.py`.
+"""
+
+from typing import Tuple, cast
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+
+class IndexFirstAxis(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, input: torch.Tensor,
+                indices: torch.Tensor) -> torch.Tensor:
+        """Get just the values of `input` which are at `indices`.
+
+        Arguments:
+            ctx: the autograd context object
+            input: (b, ...) 2+ dimensional tensor
+            indices: (num_idx) 1D tensor
+        """
+        ctx.save_for_backward(indices)
+        assert input.ndim >= 2
+        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[
+            1:]  # type: ignore
+        second_dim = other_shape.numel(
+        )  # product of sizes of all but first dimension
+        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
+        return torch.gather(
+            rearrange(input, 'b ... -> b (...)'),  # (b, ...) -> (b, second_dim)
+            0,
+            repeat(indices, 'z -> z d',
+                   d=second_dim)  # (indices,) -> (indices, second_dim)
+        ).reshape(-1, *other_shape)  # (num_idx, ...)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
+        indices, = ctx.saved_tensors
+        assert grad_output.ndim >= 2
+        other_shape = grad_output.shape[1:]
+        grad_output = rearrange(grad_output, 'b ... -> b (...)')
+        grad_input = torch.zeros([ctx.first_axis_dim, grad_output.shape[1]],
+                                 device=grad_output.device,
+                                 dtype=grad_output.dtype)
+        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
+        # grad_input[indices] = grad_output
+        grad_input.scatter_(0,
+                            repeat(indices, 'z -> z d', d=grad_output.shape[1]),
+                            grad_output)
+        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
+
+
+index_first_axis = IndexFirstAxis.apply
+
+
+class IndexPutFirstAxis(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, values: torch.Tensor, indices: torch.Tensor,
+                first_axis_dim) -> torch.Tensor:
+        ctx.save_for_backward(indices)
+        assert indices.ndim == 1
+        assert values.ndim >= 2
+        output = torch.zeros(first_axis_dim,
+                             *values.shape[1:],
+                             device=values.device,
+                             dtype=values.dtype)
+        output[indices] = values
+        return output
+
+    @staticmethod
+    def backward(ctx,
+                 grad_output: torch.Tensor) -> Tuple[torch.Tensor, None, None]:
+        indices, = ctx.saved_tensors
+        grad_values = grad_output[indices]
+        return grad_values, None, None
+
+
+index_put_first_axis = IndexPutFirstAxis.apply
+
+
+def unpad_input(
+    hidden_states: torch.Tensor,
+    attention_mask: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+    """Remove padding from input sequences.
+
+    Arguments:
+        hidden_states: (batch, seqlen, ...)
+        attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
+
+    Returns:
+        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
+        indices: (total_nnz)
+        cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states.
+        max_seqlen_in_batch: int ()
+    """
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = int(seqlens_in_batch.max().item())
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32),
+                       (1, 0))
+    # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
+    # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
+    # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
+    # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
+    # so we write custom forward and backward to make it a bit faster.
+    hidden_states = cast(
+        torch.Tensor,
+        index_first_axis(rearrange(hidden_states, 'b s ... -> (b s) ...'),
+                         indices))
+    return hidden_states, indices, cu_seqlens, max_seqlen_in_batch
+
+
+def unpad_input_only(
+    hidden_states: torch.Tensor,
+    attention_mask: torch.Tensor,
+) -> torch.Tensor:
+    """Like unpad_input, but only return the unpadded first tensor.
+
+    Save a small amount of overhead.
+
+    Arguments:
+        hidden_states: (batch, seqlen, ...)
+        attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
+
+    Returns:
+        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
+    """
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    return index_first_axis(rearrange(hidden_states, 'b s ... -> (b s) ...'),
+                            indices)
+
+
+def pad_input(hidden_states: torch.Tensor, indices: torch.Tensor, batch: int,
+              seqlen: int) -> torch.Tensor:
+    """Add padding to sequences.
+
+    Arguments:
+        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
+        indices: (total_nnz)
+        batch: int batch_size
+        seqlen: int max sequence length
+
+    Returns:
+        hidden_states: (batch, seqlen, ...)
+    """
+    output = index_put_first_axis(hidden_states, indices, batch * seqlen)
+    return rearrange(output, '(b s) ... -> b s ...', b=batch)

config.json

@@ -0,0 +1,114 @@
+{
+  "_name_or_path": "output_dir_XTBert",
+  "alibi_starting_size": 512,
+  "architectures": [
+    "BertForSequenceClassification"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_bert.BertConfig",
+    "AutoModelForMaskedLM": "xiaotinghe/XTBert--bert_layers.BertForMaskedLM",
+    "AutoModelForSequenceClassification": "bert_layers.BertForSequenceClassification"
+  },
+  "bos_token_id": 0,
+  "classifier_dropout": null,
+  "directionality": "bidi",
+  "eos_token_id": 2,
+  "gradient_checkpointing": false,
+  "hidden_act": "silu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "id2label": {
+    "0": "academic disciplines",
+    "1": "business",
+    "2": "code",
+    "3": "communication",
+    "4": "culture",
+    "5": "economy",
+    "6": "education",
+    "7": "energy",
+    "8": "engineering",
+    "9": "entertainment",
+    "10": "food and drink",
+    "11": "geography",
+    "12": "government",
+    "13": "history",
+    "14": "human behavior",
+    "15": "humanities",
+    "16": "information",
+    "17": "internet",
+    "18": "knowledge",
+    "19": "language",
+    "20": "law",
+    "21": "life health",
+    "22": "mass media",
+    "23": "mathematics",
+    "24": "military",
+    "25": "nature",
+    "26": "people",
+    "27": "philosophy",
+    "28": "politics",
+    "29": "religion",
+    "30": "science",
+    "31": "society",
+    "32": "sports",
+    "33": "time"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 2048,
+  "label2id": {
+    "academic disciplines": 0,
+    "business": 1,
+    "code": 2,
+    "communication": 3,
+    "culture": 4,
+    "economy": 5,
+    "education": 6,
+    "energy": 7,
+    "engineering": 8,
+    "entertainment": 9,
+    "food and drink": 10,
+    "geography": 11,
+    "government": 12,
+    "history": 13,
+    "human behavior": 14,
+    "humanities": 15,
+    "information": 16,
+    "internet": 17,
+    "knowledge": 18,
+    "language": 19,
+    "law": 20,
+    "life health": 21,
+    "mass media": 22,
+    "mathematics": 23,
+    "military": 24,
+    "nature": 25,
+    "people": 26,
+    "philosophy": 27,
+    "politics": 28,
+    "religion": 29,
+    "science": 30,
+    "society": 31,
+    "sports": 32,
+    "time": 33
+  },
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 4096,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "output_past": true,
+  "pad_token_id": 1,
+  "pooler_fc_size": 768,
+  "pooler_num_attention_heads": 12,
+  "pooler_num_fc_layers": 3,
+  "pooler_size_per_head": 128,
+  "pooler_type": "first_token_transform",
+  "position_embedding_type": "absolute",
+  "problem_type": "single_label_classification",
+  "torch_dtype": "float32",
+  "transformers_version": "4.33.2",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 39984
+}

configuration_bert.py

@@ -0,0 +1,26 @@
+# Copyright 2022 MosaicML Examples authors
+# SPDX-License-Identifier: Apache-2.0
+
+from transformers import BertConfig as TransformersBertConfig
+
+
+class BertConfig(TransformersBertConfig):
+
+    def __init__(
+        self,
+        alibi_starting_size: int = 512,
+        attention_probs_dropout_prob: float = 0.0,
+        **kwargs,
+    ):
+        """Configuration class for MosaicBert.
+
+        Args:
+            alibi_starting_size (int): Use `alibi_starting_size` to determine how large of an alibi tensor to
+                create when initializing the model. You should be able to ignore this parameter in most cases.
+                Defaults to 512.
+            attention_probs_dropout_prob (float): By default, turn off attention dropout in Mosaic BERT
+                (otherwise, Flash Attention will be off by default). Defaults to 0.0.
+        """
+        super().__init__(
+            attention_probs_dropout_prob=attention_probs_dropout_prob, **kwargs)
+        self.alibi_starting_size = alibi_starting_size

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:715241907dbdcbd5783080f5e62ef3fda5985b3883b051016d2154f0b71843a5
+size 465304470