added in file

README.md

@@ -1,33 +1,3 @@
 ---
-language:
-- en
-pipeline_tag: fill-mask
----
-
-## How to run:
-Install these requirements
-```
-pip install flash_attn
-pip install transformers==4.45.2 # (probably works with newer/older but tested with >=4.45.2)
-```
-Then you can load with 
-```
-from transformers import AutoModel, AutoTokenizer
-import torch
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
-
-model = AutoModel.from_pretrained("ModernBERT/bert24-base-v2-2ep-decay_100B-0.08-lr", trust_remote_code=True)
-model = model.to(device)
-tokenizer = AutoTokenizer.from_pretrained("ModernBERT/bert24-base-v2-2ep-decay_100B-0.08-lr", trust_remote_code=True)
-
-# test it out and encode some text
-text = "This is a test sentence"
-inputs = tokenizer(text, return_tensors="pt")
-inputs = {k: v.to(device) for k, v in inputs.items()}
-
-outputs = model(**inputs)
-last_hidden_states = outputs.last_hidden_state
-print(last_hidden_states.shape)
-```
+license: apache-2.0
+---

__init__.py

@@ -0,0 +1,74 @@
+from .attention import (
+    BertAlibiUnpadAttention,
+    BertAlibiUnpadSelfAttention,
+    BertSelfOutput,
+    FlexBertPaddedAttention,
+    FlexBertUnpadAttention,
+)
+from .embeddings import (
+    BertAlibiEmbeddings,
+    FlexBertAbsoluteEmbeddings,
+    FlexBertSansPositionEmbeddings,
+)
+from .layers import (
+    BertAlibiEncoder,
+    BertAlibiLayer,
+    BertResidualGLU,
+    FlexBertPaddedPreNormLayer,
+    FlexBertPaddedPostNormLayer,
+    FlexBertUnpadPostNormLayer,
+    FlexBertUnpadPreNormLayer,
+)
+from .modeling_flexbert import (
+    BertLMPredictionHead,
+    BertModel,
+    BertForMaskedLM,
+    BertForSequenceClassification,
+    BertForMultipleChoice,
+    BertOnlyMLMHead,
+    BertOnlyNSPHead,
+    BertPooler,
+    BertPredictionHeadTransform,
+    FlexBertModel,
+    FlexBertForMaskedLM,
+    FlexBertForSequenceClassification,
+    FlexBertForMultipleChoice,
+    FlexBertForCasualLM,
+)
+from .bert_padding import(
+    IndexFirstAxis,
+    IndexPutFirstAxis
+)
+
+__all__ = [
+    "BertAlibiEmbeddings",
+    "BertAlibiEncoder",
+    "BertForMaskedLM",
+    "BertForSequenceClassification",
+    "BertForMultipleChoice",
+    "BertResidualGLU",
+    "BertAlibiLayer",
+    "BertLMPredictionHead",
+    "BertModel",
+    "BertOnlyMLMHead",
+    "BertOnlyNSPHead",
+    "BertPooler",
+    "BertPredictionHeadTransform",
+    "BertSelfOutput",
+    "BertAlibiUnpadAttention",
+    "BertAlibiUnpadSelfAttention",
+    "FlexBertPaddedAttention",
+    "FlexBertUnpadAttention",
+    "FlexBertAbsoluteEmbeddings",
+    "FlexBertSansPositionEmbeddings",
+    "FlexBertPaddedPreNormLayer",
+    "FlexBertPaddedPostNormLayer",
+    "FlexBertUnpadPostNormLayer",
+    "FlexBertUnpadPreNormLayer",
+    "FlexBertModel",
+    "FlexBertForMaskedLM",
+    "FlexBertForSequenceClassification",
+    "FlexBertForMultipleChoice",
+    "IndexFirstAxis",
+    "IndexPutFirstAxis"
+]

activation.py

@@ -0,0 +1,60 @@
+# Copyright 2024 **AUTHORS_TODO**
+# License: Apache-2.0
+
+# Copyright 2020 The HuggingFace Team.
+# License: Apache-2.0
+
+from collections import OrderedDict
+from typing import Union
+import torch.nn as nn
+from .configuration_bert import FlexBertConfig
+
+
+class ClassInstantier(OrderedDict):
+    def __getitem__(self, key):
+        content = super().__getitem__(key)
+        cls, kwargs = content if isinstance(content, tuple) else (content, {})
+        return cls(**kwargs)
+
+
+ACT2CLS = {
+    "celu": nn.CELU,
+    "elu": nn.ELU,
+    "gelu": nn.GELU,
+    "gelu_tanh": (nn.GELU, {"approximate": "tanh"}),
+    "hardtanh": nn.Hardtanh,
+    "hardsigmoid": nn.Hardsigmoid,
+    "hardshrink": nn.Hardshrink,
+    "hardswish": nn.Hardswish,
+    "leaky_relu": nn.LeakyReLU,
+    "logsigmoid": nn.LogSigmoid,
+    "mish": nn.Mish,
+    "prelu": nn.PReLU,
+    "relu": nn.ReLU,
+    "relu6": nn.ReLU6,
+    "rrelu": nn.RReLU,
+    "selu": nn.SELU,
+    "sigmoid": nn.Sigmoid,
+    "silu": nn.SiLU,
+    "softmin": nn.Softmin,
+    "softplus": nn.Softplus,
+    "softshrink": nn.Softshrink,
+    "softsign": nn.Softsign,
+    "swish": nn.SiLU,
+    "tanh": nn.Tanh,
+    "tanhshrink": nn.Tanhshrink,
+    "threshold": nn.Threshold,
+}
+ACT2FN = ClassInstantier(ACT2CLS)
+
+
+def get_act_fn(config: Union[FlexBertConfig, str]) -> nn.Module:
+    try:
+        if isinstance(config, str):
+            return ACT2FN[config]
+        return ACT2FN[config.hidden_act]
+    except KeyError:
+        if isinstance(config, str):
+            raise ValueError(f"Invalid activation function type: {config}, must be one of {ACT2FN.keys()}.")
+        else:
+            raise ValueError(f"Invalid activation function type: {config.hidden_act=}, must be one of {ACT2FN.keys()}.")

attention.py

@@ -0,0 +1,1601 @@
+# Copyright 2024 **AUTHORS_TODO**
+# License: Apache-2.0
+
+# Copyright 2022 MosaicML Examples authors
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright 2023 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) 2023, Tri Dao.
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import warnings
+from typing import Optional
+import importlib.metadata
+import logging
+import math
+
+import bert_padding
+from .configuration_bert import FlexBertConfig, maybe_add_padding
+from .normalization import get_norm_layer
+from .initialization import ModuleType, init_weights
+import src.utils  # noqa: F401
+
+IMPL_USE_FLASH3 = False
+IMPL_USE_FLASH2 = False
+try:
+    from flash_attn_interface import flash_attn_varlen_func
+
+    IMPL_USE_FLASH3 = True
+except ImportError:
+    pass
+# Import Flash Attention 2, which supports ALiBi https://github.com/Dao-AILab/flash-attention
+try:
+    from flash_attn import flash_attn_varlen_qkvpacked_func, flash_attn_qkvpacked_func  # type: ignore
+
+    installed_version = importlib.metadata.version("flash_attn")  # type: ignore
+    if installed_version < "2.5.7":
+        raise ImportError("newer version of flash_attn required (>= 2.5.7)")
+    IMPL_USE_FLASH2 = True
+except ImportError:
+    pass
+
+try:
+    from flash_attn.layers.rotary import RotaryEmbedding  # type: ignore
+    from .rotary import UnpaddedRotaryEmbedding  # type: ignore
+
+except ImportError:
+    RotaryEmbedding = None
+    UnpaddedRotaryEmbedding = None
+
+logger = logging.getLogger(__name__)
+
+
+class BertAlibiUnpadSelfAttention(nn.Module):
+    """Performs multi-headed self attention on a batch of unpadded sequences.
+
+    If Flash Attention 2 is installed, this module uses Flash Attention to greatly improve throughput.
+    The Flash Attention implementation used in MosaicBERT supports arbitrary attention biases (which
+    we use to implement ALiBi). If either Flash Attention 2 is not installed the implementation will
+    default to a math-equivalent pytorch version, which is much slower.
+
+    See `forward` method for additional details.
+    """
+
+    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.is_casual = config.casual_mask
+        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.deterministic_fa2 = getattr(config, "deterministic_fa2", False)
+
+        # Warn if defaulting to pytorch because of import issues
+        if not IMPL_USE_FLASH2:
+            warnings.warn(
+                "Unable to import flash_attn; defaulting MosaicBERT attention implementation to "
+                "vanilla PyTorch (this will reduce throughput when using this model)."
+            )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+        indices: torch.Tensor,
+        attn_mask: torch.Tensor,
+        bias: torch.Tensor,
+        slopes: torch.Tensor,
+    ) -> torch.Tensor:
+        """Perform self-attention.
+
+        There are two attention implementations: vanilla attention with ALiBi, and Flash Attention 2 with ALiBi
+
+        The arguments are unpadded. The vanilla implementation of attention requires padded arguments while the
+        Flash Attention implementation does not. If using vanilla 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.
+
+        Args:
+            hidden_states: (total_nnz, dim)
+            cu_seqlens: (batch + 1,)
+            max_seqlen: int
+            indices: (total_nnz,)
+            attn_mask: (batch, max_seqlen)
+            bias: (batch, heads, max_seqlen, max_seqlen)
+            slopes: (heads) or (batch, heads)
+
+        Returns:
+            attention: (total_nnz, dim)
+        """
+        bs, dim = hidden_states.shape
+        qkv = self.Wqkv(hidden_states)
+
+        # Option 1: Flash Attention with ALiBi
+        if IMPL_USE_FLASH2:
+            qkv = qkv.view(-1, 3, self.num_attention_heads, self.attention_head_size)
+            assert 1 <= len(slopes.shape) <= 2, f"{slopes=}"
+            assert slopes.shape[-1] == self.num_attention_heads, f"{slopes=}"
+
+            convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)
+            if convert_dtype:
+                # FA2 implementation only supports fp16 and bf16
+                # If FA2 is supported, bfloat16 must be supported
+                # as of FA2 2.4.2. (Turing GPUs not supported)
+                orig_dtype = qkv.dtype
+                qkv = qkv.to(torch.bfloat16)
+
+                attention = flash_attn_varlen_qkvpacked_func(
+                    qkv,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    alibi_slopes=slopes,
+                    casual=self.is_casual
+                )
+                attention = attention.to(orig_dtype)  # type: ignore
+            else:
+                attention = flash_attn_varlen_qkvpacked_func(
+                    qkv,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    alibi_slopes=slopes,
+                    casual = self.is_casual
+                )
+        else:
+            assert not self.is_casual, f"Casual mask not implemented here yet"
+            qkv = bert_padding.pad_input(qkv, indices, cu_seqlens.shape[0] - 1, max_seqlen)  # batch, max_seqlen, thd
+            unpad_bs, *_ = qkv.shape
+            qkv = qkv.view(unpad_bs, -1, 3, self.num_attention_heads, self.attention_head_size)
+            # if we have nonzero attention dropout (e.g. during fine-tuning) or no Triton, compute attention in PyTorch
+            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
+
+            attention = bert_padding.unpad_input_only(attention, torch.squeeze(attn_mask) == 1)
+
+        return attention.view(bs, dim)
+
+
+# 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.LayerNorm = get_norm_layer(config)
+        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.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertAlibiUnpadAttention(nn.Module):
+    """Chains attention, Dropout, and LayerNorm for Mosaic BERT."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.self = BertAlibiUnpadSelfAttention(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,
+        slopes: 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)
+            bias: None or (batch, heads, max_seqlen, max_seqlen)
+            slopes: None or (batch, heads) or (heads,)
+        """
+        assert (bias is None) == (slopes is None), f"{bias=}, {slopes=}"
+        self_output = self.self(input_tensor, cu_seqlens, max_s, indices, attn_mask, bias, slopes)
+        if subset_idx is not None:
+            return self.output(
+                bert_padding.index_first_axis(self_output, subset_idx),
+                bert_padding.index_first_axis(input_tensor, subset_idx),
+            )
+        else:
+            return self.output(self_output, input_tensor)
+
+
+class FlexBertAttentionBase(nn.Module):
+    """A FlexBERT attention base class for type hints."""
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_id = layer_id
+
+    def _init_weights(self, reset_params: bool = False):
+        raise NotImplementedError("This is a base class and should not be used directly.")
+
+    def forward(self, hidden_states: torch.Tensor, attn_mask: torch.Tensor, **kwargs) -> torch.Tensor:
+        raise NotImplementedError("This is a base class and should not be used directly.")
+
+    def extra_repr(self) -> str:
+        repr = ""
+        if hasattr(self, "num_attention_heads"):
+            repr += f"num_attention_heads={self.num_attention_heads}"
+        if hasattr(self, "attn_head_size"):
+            repr += f", attn_head_size={self.attn_head_size}"
+        if hasattr(self, "sliding_window"):
+            repr += f", sliding_window={self.sliding_window if self.sliding_window != (-1, -1) else 'False'}"
+        if hasattr(self, "use_fa2"):
+            repr += f", use_fa2={self.use_fa2}"
+        if hasattr(self, "deterministic_fa2"):
+            repr += f", deterministic_fa2={self.deterministic_fa2}"
+        return repr
+
+
+class FlexBertUnpadAttention(FlexBertAttentionBase):
+    """Performs multi-headed self attention on a batch of unpadded sequences.
+
+    If Flash Attention 2 is installed, this module uses Flash Attention to improve throughput.
+    If Flash Attention 2 is not installed, the implementation will use PyTorch's SDPA kernel,
+    which requires padding and unpadding inputs, adding some overhead.
+
+    See `forward` method for additional detail.
+    """
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        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.is_casual = config.casual_mask
+        self.num_attention_heads = config.num_attention_heads
+        self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attn_head_size
+        self.p_dropout = config.attention_probs_dropout_prob
+        self.Wqkv = nn.Linear(config.hidden_size, 3 * self.all_head_size, bias=config.attn_qkv_bias)
+        self.Wo = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_out_bias)
+        self.out_drop = (
+            nn.Dropout(config.attn_out_dropout_prob) if config.attn_out_dropout_prob > 0.0 else nn.Identity()
+        )
+        self.use_fa2 = config.use_fa2
+        self.deterministic_fa2 = config.deterministic_fa2
+        self.use_sdpa_attn_mask = config.use_sdpa_attn_mask
+
+        if config.global_attn_every_n_layers > 0:
+            if config.sliding_window == -1:
+                raise ValueError("global_attn_every_n_layers` requires `sliding_window` to be set")
+            if layer_id % config.global_attn_every_n_layers != 0:
+                self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+            else:
+                self.sliding_window = (-1, -1)
+        else:
+            self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+
+        # Warn if defaulting to pytorch because of import issues
+        if not IMPL_USE_FLASH2 and self.use_fa2:
+            logger.warn_once(
+                "Unable to import flash_attn; defaulting FlexBERT attention implementation to PyTorch's"
+                " SDPA kernel. This requires padding and unpadding inputs, which will add some overhead."
+            )
+            self.use_fa2 = False
+        if not self.use_fa2:
+            if not self.use_sdpa_attn_mask:
+                logger.warn_once(
+                    "SDPA attention is being used without an attention mask. Including padding in the "
+                    " attention calculation may cause differences from the Flash Attention implementation."
+                )
+            else:
+                logger.warn_once(
+                    "SDPA attention with an attention mask doesn't use the Flash Attention kernel and will"
+                    " use more memory during the backward pass. Use the FA2 backend for linear memory scaling"
+                    " with sequence length."
+                )
+            if self.sliding_window[0] > 0:
+                raise ValueError("Sliding window is not implemented for the PyTorch SDPA path. Use the FA2 backend.")
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(
+            self.config,
+            self.Wqkv,
+            layer_dim=self.config.hidden_size,
+            layer_id=None,
+            type_of_module=ModuleType.in_module,
+        )
+        init_weights(
+            self.config,
+            self.Wo,
+            layer_dim=self.config.hidden_size,
+            layer_id=self.layer_id,
+            type_of_module=ModuleType.out_module,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+        indices: torch.Tensor,
+        attn_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Perform self-attention.
+
+        There are two attention implementations supported: PyTorch's SDPA attention and Flash Attention 2.
+
+        The arguments are unpadded. The SDPA implementation of attention requires padded arguments while the
+        Flash Attention implementation does not. If using SDPA 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.
+
+        Args:
+            hidden_states: (total_nnz, dim)
+            cu_seqlens: (batch + 1,)
+            max_seqlen: int
+            indices: (total_nnz,)
+            attn_mask: (batch, max_seqlen)
+
+        Returns:
+            attention: (total_nnz, dim)
+        """
+        bs, dim = hidden_states.shape
+        qkv = self.Wqkv(hidden_states)
+
+        if self.use_fa2:
+            qkv = qkv.view(-1, 3, self.num_attention_heads, self.attn_head_size)
+
+            convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)
+            if convert_dtype:
+                # FA2 implementation only supports fp16 and bf16. If FA2 is supported,
+                # bfloat16 must be supported as of FA2 2.5.7. (Turing GPUs not supported)
+                orig_dtype = qkv.dtype
+                qkv = qkv.to(torch.bfloat16)
+
+                attn = flash_attn_varlen_qkvpacked_func(
+                    qkv,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual
+                )
+                attn = attn.to(orig_dtype)  # type: ignore
+            else:
+                attn = flash_attn_varlen_qkvpacked_func(
+                    qkv,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual
+                )
+            attn = attn.view(bs, dim)
+        else:
+            assert not self.is_casual, f"Casual mask not implemented here yet"
+            qkv = bert_padding.pad_input(qkv, indices, cu_seqlens.shape[0] - 1, max_seqlen)  # batch, max_seqlen, thd
+            unpad_bs, seqlen, _ = qkv.shape
+
+            qkv = qkv.view(unpad_bs, -1, 3, self.num_attention_heads, self.attn_head_size)
+            q, k, v = qkv.transpose(3, 1).unbind(dim=2)  # b h s d
+            attn = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                dropout_p=self.p_dropout,
+                attn_mask=attn_mask[:, None, None, :seqlen].to(torch.bool).expand(unpad_bs, 1, seqlen, seqlen)
+                if self.use_sdpa_attn_mask
+                else None,
+            )
+            attn = attn.transpose(1, 2).view(unpad_bs, -1, dim)  # b s h d
+            attn = bert_padding.unpad_input_only(attn, torch.squeeze(attn_mask) == 1)
+
+        return self.out_drop(self.Wo(attn))
+
+
+class FlexBertUnpadParallelAttention(FlexBertAttentionBase):
+    """Computes the output of the multi-headed self parallel attention on a batch of unpadded sequences
+
+    If Flash Attention 2 is installed, this module uses Flash Attention to improve throughput.
+    If Flash Attention 2 is not installed, the implementation will use PyTorch's SDPA kernel,
+    which requires padding and unpadding inputs, adding some overhead.
+
+    See `forward` method for additional detail.
+    """
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        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.is_casual = config.casual_mask
+        self.num_attention_heads = config.num_attention_heads
+        self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.hidden_size = config.hidden_size
+        self.p_dropout = config.attention_probs_dropout_prob
+        self.Wo = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_out_bias)
+        self.out_drop = (
+            nn.Dropout(config.attn_out_dropout_prob) if config.attn_out_dropout_prob > 0.0 else nn.Identity()
+        )
+        self.use_fa2 = config.use_fa2
+        self.deterministic_fa2 = config.deterministic_fa2
+        self.use_sdpa_attn_mask = config.use_sdpa_attn_mask
+
+        if config.global_attn_every_n_layers > 0:
+            if config.sliding_window == -1:
+                raise ValueError("global_attn_every_n_layers` requires `sliding_window` to be set")
+            if layer_id % config.global_attn_every_n_layers != 0:
+                self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+            else:
+                self.sliding_window = (-1, -1)
+        else:
+            self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+
+        # Warn if defaulting to pytorch because of import issues
+        if not IMPL_USE_FLASH2 and self.use_fa2:
+            logger.warn_once(
+                "Unable to import flash_attn; defaulting FlexBERT attention implementation to PyTorch's"
+                " SDPA kernel. This requires padding and unpadding inputs, which will add some overhead."
+            )
+            self.use_fa2 = False
+        if not self.use_fa2:
+            if not self.use_sdpa_attn_mask:
+                logger.warn_once(
+                    "SDPA attention is being used without an attention mask. Including padding in the "
+                    " attention calculation may cause differences from the Flash Attention implementation."
+                )
+            else:
+                logger.warn_once(
+                    "SDPA attention with an attention mask doesn't use the Flash Attention kernel and will"
+                    " use more memory during the backward pass. Use the FA2 backend for linear memory scaling"
+                    " with sequence length."
+                )
+            if self.sliding_window[0] > 0:
+                raise ValueError("Sliding window is not implemented for the PyTorch SDPA path. Use the FA2 backend.")
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(
+            self.config,
+            self.Wo,
+            layer_dim=self.config.hidden_size,
+            layer_id=self.layer_id,
+            type_of_module=ModuleType.out_module,
+        )
+
+    def forward(
+        self,
+        qkv: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+        indices: torch.Tensor,
+        attn_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Perform self-attention.
+
+        There are two attention implementations supported: PyTorch's SDPA attention and Flash Attention 2.
+
+        The arguments are unpadded. The SDPA implementation of attention requires padded arguments while the
+        Flash Attention implementation does not. If using SDPA 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.
+
+        Args:
+            qkv: (total_nnz, 3 * dim)
+            cu_seqlens: (batch + 1,)
+            max_seqlen: int
+            indices: (total_nnz,)
+            attn_mask: (batch, max_seqlen)
+
+        Returns:
+            attention: (total_nnz, dim)
+        """
+        bs = qkv.shape[0]
+        dim = self.hidden_size
+        if self.use_fa2:
+            qkv = qkv.view(-1, 3, self.num_attention_heads, self.attn_head_size)
+
+            convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)
+            if convert_dtype:
+                # FA2 implementation only supports fp16 and bf16. If FA2 is supported,
+                # bfloat16 must be supported as of FA2 2.5.7. (Turing GPUs not supported)
+                orig_dtype = qkv.dtype
+                qkv = qkv.to(torch.bfloat16)
+
+                attn = flash_attn_varlen_qkvpacked_func(
+                    qkv,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual
+                )
+                attn = attn.to(orig_dtype)  # type: ignore
+            else:
+                attn = flash_attn_varlen_qkvpacked_func(
+                    qkv,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual
+                )
+            attn = attn.view(bs, dim)
+        else:
+            assert not self.is_casual, f"Casual mask not implemented here yet"
+            qkv = bert_padding.pad_input(qkv, indices, cu_seqlens.shape[0] - 1, max_seqlen)  # batch, max_seqlen, thd
+            unpad_bs, seqlen, _ = qkv.shape
+
+            qkv = qkv.view(unpad_bs, -1, 3, self.num_attention_heads, self.attn_head_size)
+            q, k, v = qkv.transpose(3, 1).unbind(dim=2)  # b h s d
+            attn = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                dropout_p=self.p_dropout,
+                attn_mask=attn_mask[:, None, None, :seqlen].to(torch.bool).expand(unpad_bs, 1, seqlen, seqlen)
+                if self.use_sdpa_attn_mask
+                else None,
+            )
+            attn = attn.transpose(1, 2).view(unpad_bs, -1, dim)  # b s h d
+            attn = bert_padding.unpad_input_only(attn, torch.squeeze(attn_mask) == 1)
+
+        return self.out_drop(self.Wo(attn.view(bs, dim)))
+
+
+class FlexBertPaddedAttention(FlexBertAttentionBase):
+    """Performs multi-headed self attention on a batch of padded sequences.
+
+    This module supports two attention implementations:
+    1. Flash Attention 2 (if installed), which improves throughput.
+    2. PyTorch's scaled_dot_product_attention.
+
+    See `forward` method for additional detail.
+    """
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        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.is_casual = config.casual_mask
+        self.num_attention_heads = config.num_attention_heads
+        self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attn_head_size
+        self.p_dropout = config.attention_probs_dropout_prob
+        self.Wqkv = nn.Linear(config.hidden_size, 3 * self.all_head_size, bias=config.attn_qkv_bias)
+        self.Wo = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_out_bias)
+        self.out_drop = (
+            nn.Dropout(config.attn_out_dropout_prob) if config.attn_out_dropout_prob > 0.0 else nn.Identity()
+        )
+        self.use_fa2 = config.use_fa2
+        self.deterministic_fa2 = config.deterministic_fa2
+        self.use_sdpa_attn_mask = config.use_sdpa_attn_mask
+
+        if config.global_attn_every_n_layers > 0:
+            if config.sliding_window == -1:
+                raise ValueError("global_attn_every_n_layers` requires `sliding_window` to be set")
+            if layer_id % config.global_attn_every_n_layers != 0:
+                self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+            else:
+                self.sliding_window = (-1, -1)
+        else:
+            self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+
+        if not IMPL_USE_FLASH2 and self.use_fa2:
+            self.use_fa2 = False
+        if self.use_fa2 and self.use_sdpa_attn_mask:
+            logger.warn_once(
+                "Flash Attention 2 does not support attention masks. Use unpadded attention "
+                "the equivalent functionality of masking out padding tokens."
+            )
+        if not self.use_fa2 and self.sliding_window[0] > 0:
+            raise ValueError("Sliding window is not implemented for the PyTorch SDPA path. Use the FA2 backend.")
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(
+            self.config,
+            self.Wqkv,
+            layer_dim=self.config.hidden_size,
+            layer_id=None,
+            type_of_module=ModuleType.in_module,
+        )
+        init_weights(
+            self.config,
+            self.Wo,
+            layer_dim=self.config.hidden_size,
+            layer_id=self.layer_id,
+            type_of_module=ModuleType.out_module,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Perform self-attention.
+
+        There are two attention implementations supported:
+        Flash Attention 2 and PyTorch's scaled_dot_product_attention.
+
+        Args:
+            hidden_states: (batch, seqlen, dim)
+            attn_mask: (batch, seqlen)
+
+        Returns:
+            attention: (batch, seqlen, dim)
+        """
+        bs, seqlen, dim = hidden_states.shape
+        qkv = self.Wqkv(hidden_states)
+
+        if self.use_fa2:
+            qkv = qkv.view(bs, seqlen, 3, self.num_attention_heads, self.attn_head_size)
+
+            convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)
+            if convert_dtype:
+                # FA2 implementation only supports fp16 and bf16. If FA2 is supported,
+                # bfloat16 must be supported as of FA2 2.5.7. (Turing GPUs not supported)
+                orig_dtype = qkv.dtype
+                qkv = qkv.to(torch.bfloat16)
+
+                attn = flash_attn_qkvpacked_func(
+                    qkv,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual
+                )
+                attn = attn.to(orig_dtype)  # type: ignore
+            else:
+                attn = flash_attn_qkvpacked_func(
+                    qkv,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual
+                )
+        else:
+            assert not self.is_casual, f"Casual mask not implemented here yet"
+            qkv = qkv.view(bs, seqlen, 3, self.num_attention_heads, self.attn_head_size)
+
+            q, k, v = qkv.transpose(3, 1).unbind(dim=2)
+            attn = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                dropout_p=self.p_dropout,
+                attn_mask=attn_mask[:, None, None, :seqlen].to(torch.bool).expand(bs, 1, seqlen, seqlen)
+                if self.use_sdpa_attn_mask
+                else None,
+            ).transpose(1, 2)
+
+        attn = attn.view(bs, seqlen, dim)
+        return self.out_drop(self.Wo(attn))
+
+
+class FlexBertUnpadRopeAttention(FlexBertAttentionBase):
+    """Performs multi-headed self attention on a batch of unpadded sequences.
+
+    If Flash Attention 2 is installed, this module uses Flash Attention to improve throughput.
+    If Flash Attention 2 is not installed, the implementation will use PyTorch's SDPA kernel,
+    which requires padding and unpadding inputs, adding some overhead.
+
+    See `forward` method for additional details.
+    """
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        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.is_casual = config.casual_mask
+        self.num_attention_heads = config.num_attention_heads
+        self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attn_head_size
+        self.p_dropout = config.attention_probs_dropout_prob
+        self.Wqkv = nn.Linear(config.hidden_size, 3 * self.all_head_size, bias=config.attn_qkv_bias)
+        self.Wo = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_out_bias)
+        self.out_drop = (
+            nn.Dropout(config.attn_out_dropout_prob) if config.attn_out_dropout_prob > 0.0 else nn.Identity()
+        )
+
+        if config.global_attn_every_n_layers > 0:
+            if config.sliding_window == -1:
+                raise ValueError("global_attn_every_n_layers` requires `sliding_window` to be set")
+            if layer_id % config.global_attn_every_n_layers != 0:
+                self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+            else:
+                self.sliding_window = (-1, -1)
+        else:
+            self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+
+        if config.rotary_emb_dim is None:
+            config.rotary_emb_dim = self.attn_head_size
+
+        rotary_base = config.rotary_emb_base
+        rotary_dim = config.rotary_emb_dim
+        if self.sliding_window != (-1, -1):
+            if config.local_attn_rotary_emb_base != -1:
+                rotary_base = config.local_attn_rotary_emb_base
+            if config.local_attn_rotary_emb_dim is not None:
+                rotary_dim = config.local_attn_rotary_emb_dim
+
+        assert UnpaddedRotaryEmbedding is not None, "rotary_emb is not installed"
+        self.rotary_emb = UnpaddedRotaryEmbedding(
+            dim=rotary_dim,
+            base=rotary_base,
+            scale_base=config.rotary_emb_scale_base,  # If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554).
+            interleaved=config.rotary_emb_interleaved,
+        )
+
+        self.use_fa2 = config.use_fa2
+        # flash attention 3 only supports global attention
+        self.use_fa3 = config.use_fa2 and self.sliding_window == (-1, -1) and IMPL_USE_FLASH3
+        self.deterministic_fa2 = config.deterministic_fa2
+        self.use_sdpa_attn_mask = config.use_sdpa_attn_mask
+
+        # Warn if defaulting to pytorch because of import issues
+        if not IMPL_USE_FLASH2 and self.use_fa2:
+            logger.warn_once(
+                "Unable to import flash_attn; defaulting FlexBERT attention implementation to PyTorch's"
+                " SDPA kernel. This requires padding and unpadding inputs, which will add some overhead."
+            )
+            self.use_fa2 = False
+        if not self.use_fa2:
+            if not self.use_sdpa_attn_mask:
+                logger.warn_once(
+                    "SDPA attention is being used without an attention mask. Including padding in the "
+                    " attention calculation may cause differences from the Flash Attention implementation."
+                )
+            else:
+                logger.warn_once(
+                    "SDPA attention with an attention mask doesn't use the Flash Attention kernel and will"
+                    " use more memory during the backward pass. Use the FA2 backend for linear memory scaling"
+                    " with sequence length."
+                )
+            if self.sliding_window[0] > 0:
+                raise ValueError("Sliding window is not implemented for the PyTorch SDPA path. Use the FA2 backend.")
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(
+            self.config,
+            self.Wqkv,
+            layer_dim=self.config.hidden_size,
+            layer_id=None,
+            type_of_module=ModuleType.in_module,
+        )
+        init_weights(
+            self.config,
+            self.Wo,
+            layer_dim=self.config.hidden_size,
+            layer_id=self.layer_id,
+            type_of_module=ModuleType.out_module,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+        indices: torch.Tensor,
+        attn_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Perform self-attention.
+
+        There are two attention implementations supported: PyTorch's SDPA attention and Flash Attention 2.
+
+        The arguments are unpadded. The SDPA implementation of attention requires padded arguments while the
+        Flash Attention implementation does not. If using SDPA 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.
+
+        Args:
+            hidden_states: (total_nnz, dim)
+            cu_seqlens: (batch + 1,)
+            max_seqlen: int
+            indices: (total_nnz,)
+            attn_mask: (batch, max_seqlen)
+
+        Returns:
+            attention: (total_nnz, dim)
+        """
+        bs, dim = hidden_states.shape
+        qkv = self.Wqkv(hidden_states)
+
+        # only needed for inference when we have KV cache
+        seqlen_offset = 0
+
+        # (total_seqlen, 3, nheads, headdim)
+        qkv = qkv.view(-1, 3, self.num_attention_heads, self.attn_head_size)
+        qkv = self.rotary_emb(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, seqlen_offset=seqlen_offset)
+
+        if self.use_fa3:
+            convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)
+            if convert_dtype:
+                # FA2 implementation only supports fp16 and bf16. If FA2 is supported,
+                # bfloat16 must be supported as of FA2 2.5.7. (Turing GPUs not supported)
+                orig_dtype = qkv.dtype
+                qkv = qkv.to(torch.bfloat16)
+                q, k, v = qkv.view(-1, 3, self.num_attention_heads, self.attn_head_size).unbind(dim=1)
+
+                attn, _ = flash_attn_varlen_func(
+                    q=q,
+                    k=k,
+                    v=v,
+                    cu_seqlens_q=cu_seqlens,
+                    cu_seqlens_k=cu_seqlens,
+                    max_seqlen_q=max_seqlen,
+                    max_seqlen_k=max_seqlen,
+                    deterministic=self.deterministic_fa2,
+                    causal=self.is_casual,
+                )
+                attn = attn.to(orig_dtype)  # type: ignore
+            else:
+                q, k, v = qkv.view(-1, 3, self.num_attention_heads, self.attn_head_size).unbind(dim=1)
+                attn, _ = flash_attn_varlen_func(
+                    q=q,
+                    k=k,
+                    v=v,
+                    cu_seqlens_q=cu_seqlens,
+                    cu_seqlens_k=cu_seqlens,
+                    max_seqlen_q=max_seqlen,
+                    max_seqlen_k=max_seqlen,
+                    deterministic=self.deterministic_fa2,
+                    causal=self.is_casual,
+                )
+            attn = attn.view(bs, dim)
+        elif self.use_fa2:
+            convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)
+            if convert_dtype:
+                # FA2 implementation only supports fp16 and bf16. If FA2 is supported,
+                # bfloat16 must be supported as of FA2 2.5.7. (Turing GPUs not supported)
+                orig_dtype = qkv.dtype
+                qkv = qkv.to(torch.bfloat16)
+
+                attn = flash_attn_varlen_qkvpacked_func(
+                    qkv,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    causal=self.is_casual,
+                )
+                attn = attn.to(orig_dtype)  # type: ignore
+            else:
+                attn = flash_attn_varlen_qkvpacked_func(
+                    qkv,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    causal=self.is_casual,
+                )
+            attn = attn.view(bs, dim)
+        else:
+            assert not self.is_casual, f"Casual mask not implemented here yet"
+            qkv = bert_padding.pad_input(
+                qkv, indices, cu_seqlens.shape[0] - 1, attn_mask.shape[-1]
+            )  # batch, max_seqlen, thd
+            unpad_bs, seqlen, *_ = qkv.shape
+
+            q, k, v = qkv.transpose(3, 1).unbind(dim=2)  # b h s d
+            attn = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                dropout_p=self.p_dropout,
+                attn_mask=attn_mask[:, None, None, :seqlen].to(torch.bool).expand(unpad_bs, 1, seqlen, seqlen)
+                if self.use_sdpa_attn_mask
+                else None,
+            )
+            attn = attn.transpose(1, 2).view(unpad_bs, -1, dim)  # b s h d
+            attn = bert_padding.unpad_input_only(attn, torch.squeeze(attn_mask) == 1)
+
+        return self.out_drop(self.Wo(attn))
+
+
+class FlexBertPaddedRopeAttention(FlexBertAttentionBase):
+    """Performs multi-headed self attention on a batch of padded sequences.
+
+    This module supports two attention implementations:
+    1. Flash Attention 2 (if installed), which improves throughput.
+    2. PyTorch's scaled_dot_product_attention.
+
+    See `forward` method for additional details.
+    """
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        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.is_casual = config.casual_mask
+        self.num_attention_heads = config.num_attention_heads
+        self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attn_head_size
+        self.p_dropout = config.attention_probs_dropout_prob
+        self.Wqkv = nn.Linear(config.hidden_size, 3 * self.all_head_size, bias=config.attn_qkv_bias)
+        self.Wo = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_out_bias)
+        self.out_drop = (
+            nn.Dropout(config.attn_out_dropout_prob) if config.attn_out_dropout_prob > 0.0 else nn.Identity()
+        )
+
+        self.use_fa2 = config.use_fa2
+        self.deterministic_fa2 = config.deterministic_fa2
+        self.use_sdpa_attn_mask = config.use_sdpa_attn_mask
+
+        if config.global_attn_every_n_layers > 0:
+            if config.sliding_window == -1:
+                raise ValueError("global_attn_every_n_layers` requires `sliding_window` to be set")
+            if layer_id % config.global_attn_every_n_layers != 0:
+                self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+            else:
+                self.sliding_window = (-1, -1)
+        else:
+            self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+
+        if config.rotary_emb_dim is None:
+            config.rotary_emb_dim = self.attn_head_size
+
+        rotary_base = config.rotary_emb_base
+        rotary_dim = config.rotary_emb_dim
+        if self.sliding_window != (-1, -1):
+            if config.local_attn_rotary_emb_base != -1:
+                rotary_base = config.local_attn_rotary_emb_base
+            if config.local_attn_rotary_emb_dim is not None:
+                rotary_dim = config.local_attn_rotary_emb_dim
+
+        assert RotaryEmbedding is not None, "rotary_emb is not installed"
+        self.rotary_emb = RotaryEmbedding(
+            dim=rotary_dim,
+            base=rotary_base,
+            scale_base=config.rotary_emb_scale_base,  # If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554).
+            interleaved=config.rotary_emb_interleaved,
+        )
+
+        if not IMPL_USE_FLASH2 and self.use_fa2:
+            self.use_fa2 = False
+        if self.use_fa2 and self.use_sdpa_attn_mask:
+            logger.warn_once(
+                "Flash Attention 2 does not support attention masks. Use unpadded attention "
+                "the equivalent functionality of masking out padding tokens."
+            )
+        if not self.use_fa2 and self.sliding_window[0] > 0:
+            raise ValueError("Sliding window is not implemented for the PyTorch SDPA path. Use the FA2 backend.")
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(
+            self.config,
+            self.Wqkv,
+            layer_dim=self.config.hidden_size,
+            layer_id=None,
+            type_of_module=ModuleType.in_module,
+        )
+        init_weights(
+            self.config,
+            self.Wo,
+            layer_dim=self.config.hidden_size,
+            layer_id=self.layer_id,
+            type_of_module=ModuleType.out_module,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Perform self-attention.
+
+        There are two attention implementations supported:
+        Flash Attention 2 and PyTorch's scaled_dot_product_attention.
+
+        Args:
+            hidden_states: (batch, seqlen, dim)
+            attn_mask: (batch, seqlen)
+
+        Returns:
+            attention: (batch, seqlen, dim)
+        """
+        bs, seqlen, dim = hidden_states.shape
+        qkv = self.Wqkv(hidden_states)
+
+        seqlen_offset = 0
+
+        # Reshape to (batch, seqlen, 3, nheads, headdim)
+        qkv = qkv.view(bs, seqlen, 3, self.num_attention_heads, self.attn_head_size)
+
+        if IMPL_USE_FLASH2:
+            # Apply RoPE
+            qkv = self.rotary_emb(qkv, seqlen_offset=seqlen_offset, max_seqlen=None)
+
+            convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)
+            if convert_dtype:
+                # FA2 implementation only supports fp16 and bf16. If FA2 is supported,
+                # bfloat16 must be supported as of FA2 2.5.7. (Turing GPUs not supported)
+                orig_dtype = qkv.dtype
+                qkv = qkv.to(torch.bfloat16)
+
+                attn = flash_attn_qkvpacked_func(
+                    qkv,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual,
+                )
+                attn = attn.to(orig_dtype)  # type: ignore
+            else:
+                attn = flash_attn_qkvpacked_func(
+                    qkv,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual
+                )
+        else:
+            assert not self.is_casual, f"Casual mask not implemented here yet"
+            qkv = self.rotary_emb(qkv, seqlen_offset=seqlen_offset, max_seqlen=None)
+            q, k, v = qkv.transpose(3, 1).unbind(dim=2)
+            attn = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                dropout_p=self.p_dropout,
+                attn_mask=attn_mask[:, None, None, :seqlen].to(torch.bool).expand(bs, 1, seqlen, seqlen)
+                if self.use_sdpa_attn_mask
+                else None,
+            ).transpose(1, 2)
+
+        attn = attn.view(bs, seqlen, dim)
+        return self.out_drop(self.Wo(attn))
+
+
+class FlexBertUnpadRopeParallelAttention(FlexBertAttentionBase):
+    """Performs multi-headed self attention on a batch of unpadded sequences.
+
+    If Flash Attention 2 is installed, this module uses Flash Attention to improve throughput.
+    If Flash Attention 2 is not installed, the implementation will use PyTorch's SDPA kernel,
+    which requires padding and unpadding inputs, adding some overhead.
+
+    See `forward` method for additional details.
+    """
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        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.is_casual = config.casual_mask
+        self.num_attention_heads = config.num_attention_heads
+        self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.hidden_size = config.hidden_size
+        self.p_dropout = config.attention_probs_dropout_prob
+        self.Wo = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_out_bias)
+        self.out_drop = (
+            nn.Dropout(config.attn_out_dropout_prob) if config.attn_out_dropout_prob > 0.0 else nn.Identity()
+        )
+
+        if config.global_attn_every_n_layers > 0:
+            if config.sliding_window == -1:
+                raise ValueError("global_attn_every_n_layers` requires `sliding_window` to be set")
+            if layer_id % config.global_attn_every_n_layers != 0:
+                self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+            else:
+                self.sliding_window = (-1, -1)
+        else:
+            self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+
+        if config.rotary_emb_dim is None:
+            config.rotary_emb_dim = self.attn_head_size
+
+        rotary_base = config.rotary_emb_base
+        rotary_dim = config.rotary_emb_dim
+        if self.sliding_window != (-1, -1):
+            if config.local_attn_rotary_emb_base != -1:
+                rotary_base = config.local_attn_rotary_emb_base
+            if config.local_attn_rotary_emb_dim is not None:
+                rotary_dim = config.local_attn_rotary_emb_dim
+
+        assert UnpaddedRotaryEmbedding is not None, "rotary_emb is not installed"
+        self.rotary_emb = UnpaddedRotaryEmbedding(
+            dim=rotary_dim,
+            base=rotary_base,
+            scale_base=config.rotary_emb_scale_base,  # If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554).
+            interleaved=config.rotary_emb_interleaved,
+        )
+
+        self.use_fa2 = config.use_fa2
+        self.deterministic_fa2 = config.deterministic_fa2
+        self.use_sdpa_attn_mask = config.use_sdpa_attn_mask
+
+        # Warn if defaulting to pytorch because of import issues
+        if not IMPL_USE_FLASH2 and self.use_fa2:
+            logger.warn_once(
+                "Unable to import flash_attn; defaulting FlexBERT attention implementation to PyTorch's"
+                " SDPA kernel. This requires padding and unpadding inputs, which will add some overhead."
+            )
+            self.use_fa2 = False
+        if not self.use_fa2:
+            if not self.use_sdpa_attn_mask:
+                logger.warn_once(
+                    "SDPA attention is being used without an attention mask. Including padding in the "
+                    " attention calculation may cause differences from the Flash Attention implementation."
+                )
+            else:
+                logger.warn_once(
+                    "SDPA attention with an attention mask doesn't use the Flash Attention kernel and will"
+                    " use more memory during the backward pass. Use the FA2 backend for linear memory scaling"
+                    " with sequence length."
+                )
+            if self.sliding_window[0] > 0:
+                raise ValueError("Sliding window is not implemented for the PyTorch SDPA path. Use the FA2 backend.")
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(
+            self.config,
+            self.Wo,
+            layer_dim=self.config.hidden_size,
+            layer_id=self.layer_id,
+            type_of_module=ModuleType.out_module,
+        )
+
+    def forward(
+        self,
+        qkv: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+        indices: torch.Tensor,
+        attn_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Perform self-attention.
+
+        There are two attention implementations supported: PyTorch's SDPA attention and Flash Attention 2.
+
+        The arguments are unpadded. The SDPA implementation of attention requires padded arguments while the
+        Flash Attention implementation does not. If using SDPA 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.
+
+        Args:
+            qkv: (total_nnz, 3 * dim)
+            cu_seqlens: (batch + 1,)
+            max_seqlen: int
+            indices: (total_nnz,)
+            attn_mask: (batch, max_seqlen)
+
+        Returns:
+            attention: (total_nnz, dim)
+        """
+        bs = qkv.shape[0]
+        dim = self.hidden_size
+
+        # only needed for inference when we have KV cache
+        seqlen_offset = 0
+
+        # (total_seqlen, 3, nheads, headdim)
+        qkv = qkv.view(-1, 3, self.num_attention_heads, self.attn_head_size)
+        qkv = self.rotary_emb(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, seqlen_offset=seqlen_offset)
+
+        if self.use_fa2:
+            convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)
+            if convert_dtype:
+                # FA2 implementation only supports fp16 and bf16. If FA2 is supported,
+                # bfloat16 must be supported as of FA2 2.5.7. (Turing GPUs not supported)
+                orig_dtype = qkv.dtype
+                qkv = qkv.to(torch.bfloat16)
+
+                attn = flash_attn_varlen_qkvpacked_func(
+                    qkv,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual,
+                )
+                attn = attn.to(orig_dtype)  # type: ignore
+            else:
+                attn = flash_attn_varlen_qkvpacked_func(
+                    qkv,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual,
+                )
+            attn = attn.view(bs, dim)
+        else:
+            assert not self.is_casual, f"Casual mask not implemented here yet"
+            qkv = bert_padding.pad_input(
+                qkv, indices, cu_seqlens.shape[0] - 1, attn_mask.shape[-1]
+            )  # batch, max_seqlen, thd
+            unpad_bs, seqlen, *_ = qkv.shape
+
+            q, k, v = qkv.transpose(3, 1).unbind(dim=2)  # b h s d
+            attn = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                dropout_p=self.p_dropout,
+                attn_mask=attn_mask[:, None, None, :seqlen].to(torch.bool).expand(unpad_bs, 1, seqlen, seqlen)
+                if self.use_sdpa_attn_mask
+                else None,
+            )
+            attn = attn.transpose(1, 2).view(unpad_bs, -1, dim)  # b s h d
+            attn = bert_padding.unpad_input_only(attn, torch.squeeze(attn_mask) == 1)
+
+        return self.out_drop(self.Wo(attn))
+
+
+class FlexBertPaddedRopeParallelAttention(FlexBertAttentionBase):
+    """Performs multi-headed self attention on a batch of padded sequences.
+
+    This module supports two attention implementations:
+    1. Flash Attention 2 (if installed), which improves throughput.
+    2. PyTorch's scaled_dot_product_attention.
+
+    See `forward` method for additional details.
+    """
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        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.is_casual = config.casual_mask
+        self.num_attention_heads = config.num_attention_heads
+        self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.hidden_size = config.hidden_size
+        self.p_dropout = config.attention_probs_dropout_prob
+        self.Wo = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_out_bias)
+        self.out_drop = (
+            nn.Dropout(config.attn_out_dropout_prob) if config.attn_out_dropout_prob > 0.0 else nn.Identity()
+        )
+
+        self.use_fa2 = config.use_fa2
+        self.deterministic_fa2 = config.deterministic_fa2
+        self.use_sdpa_attn_mask = config.use_sdpa_attn_mask
+        if not IMPL_USE_FLASH2 and self.use_fa2:
+            self.use_fa2 = False
+
+        if config.global_attn_every_n_layers > 0:
+            if config.sliding_window == -1:
+                raise ValueError("global_attn_every_n_layers` requires `sliding_window` to be set")
+            if layer_id % config.global_attn_every_n_layers != 0:
+                self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+            else:
+                self.sliding_window = (-1, -1)
+        else:
+            self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+
+        if config.rotary_emb_dim is None:
+            config.rotary_emb_dim = self.attn_head_size
+
+        rotary_base = config.rotary_emb_base
+        rotary_dim = config.rotary_emb_dim
+        if self.sliding_window != (-1, -1):
+            if config.local_attn_rotary_emb_base != -1:
+                rotary_base = config.local_attn_rotary_emb_base
+            if config.local_attn_rotary_emb_dim is not None:
+                rotary_dim = config.local_attn_rotary_emb_dim
+
+        assert RotaryEmbedding is not None, "rotary_emb is not installed"
+        self.rotary_emb = RotaryEmbedding(
+            dim=rotary_dim,
+            base=rotary_base,
+            scale_base=config.rotary_emb_scale_base,  # If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554).
+            interleaved=config.rotary_emb_interleaved,
+        )
+
+        if not IMPL_USE_FLASH2 and self.use_fa2:
+            self.use_fa2 = False
+        if self.use_fa2 and self.use_sdpa_attn_mask:
+            logger.warn_once(
+                "Flash Attention 2 does not support attention masks. Use unpadded attention "
+                "the equivalent functionality of masking out padding tokens."
+            )
+        if not self.use_fa2 and self.sliding_window[0] > 0:
+            raise ValueError("Sliding window is not implemented for the PyTorch SDPA path. Use the FA2 backend.")
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(
+            self.config,
+            self.Wo,
+            layer_dim=self.config.hidden_size,
+            layer_id=self.layer_id,
+            type_of_module=ModuleType.out_module,
+        )
+
+    def forward(
+        self,
+        qkv: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Perform self-attention.
+
+        There are two attention implementations supported:
+        Flash Attention 2 and PyTorch's scaled_dot_product_attention.
+
+        Args:
+            qkv: (batch, seqlen, 3 * dim)
+            attn_mask: (batch, seqlen)
+
+        Returns:
+            attention: (batch, seqlen, dim)
+        """
+        bs, seqlen, _ = qkv.shape
+        dim = self.hidden_size
+
+        seqlen_offset = 0
+
+        # Reshape to (batch, seqlen, 3, nheads, headdim)
+        qkv = qkv.view(bs, seqlen, 3, self.num_attention_heads, self.attn_head_size)
+
+        if self.use_fa2:
+            # Apply RoPE
+            qkv = self.rotary_emb(qkv, seqlen_offset=seqlen_offset, max_seqlen=None)
+
+            convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)
+            if convert_dtype:
+                # FA2 implementation only supports fp16 and bf16. If FA2 is supported,
+                # bfloat16 must be supported as of FA2 2.5.7. (Turing GPUs not supported)
+                orig_dtype = qkv.dtype
+                qkv = qkv.to(torch.bfloat16)
+
+                attn = flash_attn_qkvpacked_func(
+                    qkv,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual
+                )
+                attn = attn.to(orig_dtype)  # type: ignore
+            else:
+                attn = flash_attn_qkvpacked_func(
+                    qkv,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual
+                )
+        else:
+            assert not self.is_casual, f"Casual mask not implemented here yet"
+            qkv = self.rotary_emb(qkv, seqlen_offset=seqlen_offset, max_seqlen=None)
+            q, k, v = qkv.transpose(3, 1).unbind(dim=2)
+            attn = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                dropout_p=self.p_dropout,
+                attn_mask=attn_mask[:, None, None, :seqlen].to(torch.bool).expand(bs, 1, seqlen, seqlen)
+                if self.use_sdpa_attn_mask
+                else None,
+            ).transpose(1, 2)
+
+        attn = attn.view(bs, seqlen, dim)
+        return self.out_drop(self.Wo(attn))
+
+
+class FlexBertPaddedParallelAttention(FlexBertAttentionBase):
+    """Performs multi-headed self attention on a batch of padded sequences.
+
+    This module supports two attention implementations:
+    1. Flash Attention 2 (if installed), which improves throughput.
+    2. PyTorch's scaled_dot_product_attention.
+
+    See `forward` method for additional detail.
+    """
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        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.is_casual = config.casual_mask
+        self.num_attention_heads = config.num_attention_heads
+        self.attn_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.hidden_size = config.hidden_size
+        self.p_dropout = config.attention_probs_dropout_prob
+        self.Wo = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_out_bias)
+        self.out_drop = (
+            nn.Dropout(config.attn_out_dropout_prob) if config.attn_out_dropout_prob > 0.0 else nn.Identity()
+        )
+        self.use_fa2 = config.use_fa2
+        self.deterministic_fa2 = config.deterministic_fa2
+        self.use_sdpa_attn_mask = config.use_sdpa_attn_mask
+
+        if config.global_attn_every_n_layers > 0:
+            if config.sliding_window == -1:
+                raise ValueError("global_attn_every_n_layers` requires `sliding_window` to be set")
+            if layer_id % config.global_attn_every_n_layers != 0:
+                self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+            else:
+                self.sliding_window = (-1, -1)
+        else:
+            self.sliding_window = (config.sliding_window // 2, config.sliding_window // 2)
+
+        if not IMPL_USE_FLASH2 and self.use_fa2:
+            self.use_fa2 = False
+        if self.use_fa2 and self.use_sdpa_attn_mask:
+            logger.warn_once(
+                "Flash Attention 2 does not support attention masks. Use unpadded attention "
+                "the equivalent functionality of masking out padding tokens."
+            )
+        if not self.use_fa2 and self.sliding_window[0] > 0:
+            raise ValueError("Sliding window is not implemented for the PyTorch SDPA path. Use the FA2 backend.")
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(
+            self.config,
+            self.Wo,
+            layer_dim=self.config.hidden_size,
+            layer_id=self.layer_id,
+            type_of_module=ModuleType.out_module,
+        )
+
+    def forward(
+        self,
+        qkv: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Perform self-attention.
+
+        There are two attention implementations supported:
+        Flash Attention 2 and PyTorch's scaled_dot_product_attention.
+
+        Args:
+            qkv: (batch, seqlen, 3 * dim)
+            attn_mask: (batch, seqlen)
+
+        Returns:
+            attention: (batch, seqlen, dim)
+        """
+        bs, seqlen, _ = qkv.shape
+        dim = self.hidden_size
+
+        if self.use_fa2:
+            qkv = qkv.view(bs, seqlen, 3, self.num_attention_heads, self.attn_head_size)
+
+            convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)
+            if convert_dtype:
+                # FA2 implementation only supports fp16 and bf16. If FA2 is supported,
+                # bfloat16 must be supported as of FA2 2.5.7. (Turing GPUs not supported)
+                orig_dtype = qkv.dtype
+                qkv = qkv.to(torch.bfloat16)
+
+                attn = flash_attn_qkvpacked_func(
+                    qkv,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual
+                )
+                attn = attn.to(orig_dtype)  # type: ignore
+            else:
+                attn = flash_attn_qkvpacked_func(
+                    qkv,
+                    dropout_p=self.p_dropout,
+                    deterministic=self.deterministic_fa2,
+                    window_size=self.sliding_window,
+                    casual=self.is_casual
+                )
+        else:
+            assert not self.is_casual, f"Casual attention mask not yet implemented here"
+            qkv = qkv.view(bs, seqlen, 3, self.num_attention_heads, self.attn_head_size)
+            q, k, v = qkv.transpose(3, 1).unbind(dim=2)  # b h s d
+            attn = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                dropout_p=self.p_dropout,
+                attn_mask=attn_mask[:, None, None, :seqlen].to(torch.bool).expand(bs, 1, seqlen, seqlen)
+                if self.use_sdpa_attn_mask
+                else None,
+            ).transpose(1, 2)
+
+        attn = attn.view(bs, seqlen, dim)
+        return self.out_drop(self.Wo(attn))
+
+
+ATTN2CLS = {
+    "unpadded_base": FlexBertUnpadAttention,
+    "padded_base": FlexBertPaddedAttention,
+    "unpadded_parallel": FlexBertUnpadParallelAttention,
+    "padded_parallel": FlexBertPaddedParallelAttention,
+    "unpadded_rope": FlexBertUnpadRopeAttention,
+    "padded_rope": FlexBertPaddedRopeAttention,
+    "unpadded_rope_parallel": FlexBertUnpadRopeParallelAttention,
+    "padded_rope_parallel": FlexBertPaddedRopeParallelAttention,
+}
+
+
+def get_attention_layer(config: FlexBertConfig, layer_id: Optional[int] = None) -> FlexBertAttentionBase:
+    try:
+        attention_layer = (
+            config.initial_attention_layer
+            if layer_id < config.num_initial_layers and getattr(config, "initial_attention_layer", None) is not None
+            else config.attention_layer
+        )
+        return ATTN2CLS[maybe_add_padding(config, attention_layer)](config, layer_id=layer_id)
+    except KeyError:
+        if layer_id < config.num_initial_layers and getattr(config, "initial_attention_layer", None) is not None:
+            raise ValueError(
+                f"Invalid attention layer type: {config.initial_attention_layer=}, must be one of {ATTN2CLS.keys()}."
+                f"{config.padding=} will be automatically prepended to `config.attention_layer` if unspecified."
+            )
+        else:
+            raise ValueError(
+                f"Invalid attention layer type: {config.attention_layer=}, must be one of {ATTN2CLS.keys()}. "
+                f"{config.padding=} will be automatically prepended to `config.attention_layer` if unspecified."
+            )

bert_padding.py

@@ -0,0 +1,141 @@
+# 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()
+    rearranged = rearrange(hidden_states, "b s ... -> (b s) ...")
+    return index_first_axis(rearranged, indices)  # type: ignore
+
+
+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)  # type: ignore

config.json

@@ -1,11 +1,13 @@
 {
   "allow_embedding_resizing": true,
   "architectures": [
+    "FlexBertModel",
     "FlexBertForCasualLM"
   ],
   "auto_map": {                                                                                                                                                                                                               
     "AutoConfig": "orionweller/FlexGPT--configuration_bert.FlexBertConfig",
-    "AutoModel": "lightonai/FlexGPT--modeling_flexbert.FlexBertForCasualLM"                                                                                                                                                                                                                  
+    "AutoModel": "lightonai/FlexGPT--modeling_flexbert.FlexBertModel",
+    "AutoModelForCasualLM": "lightonai/FlexGPT--modeling_flexbert.FlexBertForCasualLM",                                                                                                                                                                                                                 
   },
   "attention_layer": "rope",
   "attention_probs_dropout_prob": 0.0,

configuration_bert.py

@@ -0,0 +1,272 @@
+# Copyright 2022 MosaicML Examples authors
+# SPDX-License-Identifier: Apache-2.0
+
+import warnings
+
+from transformers import BertConfig as TransformersBertConfig
+
+
+class BertConfig(TransformersBertConfig):
+    def __init__(
+        self,
+        alibi_starting_size: int = 512,
+        normalization: str = "layernorm",
+        attention_probs_dropout_prob: float = 0.0,
+        head_pred_act: str = "gelu",
+        deterministic_fa2: bool = False,
+        allow_embedding_resizing: bool = False,
+        **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 MosaicBERT
+                Note that the custom Triton Flash Attention with ALiBi implementation does not support droput.
+                However, Flash Attention 2 supports ALiBi and dropout https://github.com/Dao-AILab/flash-attention
+            embed_dropout_prob (float): Dropout probability for the embedding layer.
+            attn_out_dropout_prob (float): Dropout probability for the attention output layer.
+            mlp_dropout_prob (float): Dropout probability for the MLP layer.
+            allow_embedding_resizing (bool): Embeddings will be automatically resized when they are smaller than the tokenizer vocab size.
+        """
+        super().__init__(attention_probs_dropout_prob=attention_probs_dropout_prob, **kwargs)
+        self.alibi_starting_size = alibi_starting_size
+        self.normalization = normalization
+        self.head_pred_act = head_pred_act
+        self.deterministic_fa2 = deterministic_fa2
+        self.allow_embedding_resizing = allow_embedding_resizing
+
+
+class FlexBertConfig(TransformersBertConfig):
+    def __init__(
+        self,
+        attention_layer: str = "base",
+        attention_probs_dropout_prob: float = 0.0,
+        attn_out_bias: bool = False,
+        attn_out_dropout_prob: float = 0.0,
+        attn_qkv_bias: bool = False,
+        bert_layer: str = "prenorm",
+        decoder_bias: bool = True,
+        embed_dropout_prob: float = 0.0,
+        embed_norm: bool = True,
+        final_norm: bool = False,
+        embedding_layer: str = "absolute_pos",
+        encoder_layer: str = "base",
+        loss_function: str = "cross_entropy",
+        loss_kwargs: dict = {},
+        mlp_dropout_prob: float = 0.0,
+        mlp_in_bias: bool = False,
+        mlp_layer: str = "mlp",
+        mlp_out_bias: bool = False,
+        norm_kwargs: dict = {},
+        normalization: str = "rmsnorm",
+        padding: str = "unpadded",
+        head_class_act: str = "silu",
+        head_class_bias: bool = False,
+        head_class_dropout: float = 0.0,
+        head_class_norm: str = False,
+        head_pred_act: str = "silu",
+        head_pred_bias: bool = False,
+        head_pred_dropout: float = 0.0,
+        head_pred_norm: bool = True,
+        pooling_type: str = "cls",
+        rotary_emb_dim: int | None = None,
+        rotary_emb_base: float = 10000.0,
+        rotary_emb_scale_base=None,
+        rotary_emb_interleaved: bool = False,
+        use_fa2: bool = True,
+        use_sdpa_attn_mask: bool = False,
+        allow_embedding_resizing: bool = False,
+        init_method: str = "default",
+        init_std: float = 0.02,
+        init_cutoff_factor: float = 2.0,
+        init_small_embedding: bool = False,
+        initial_attention_layer: str | None = None,
+        initial_bert_layer: str | None = None,
+        initial_mlp_layer: str | None = None,
+        num_initial_layers: int = 1,
+        skip_first_prenorm: bool = False,
+        deterministic_fa2: bool = False,
+        sliding_window: int = -1,
+        global_attn_every_n_layers: int = -1,
+        local_attn_rotary_emb_base: float = -1,
+        local_attn_rotary_emb_dim: int | None = None,
+        unpad_embeddings: bool = False,
+        pad_logits: bool = False,
+        compile_model: bool = False,
+        masked_prediction: bool = False,
+        casual_mask: bool = False,
+        **kwargs,
+    ):
+        """
+        Args:
+            attention_layer (str): Attention layer type.
+            attention_probs_dropout_prob (float): Dropout probability for attention probabilities.
+            attn_out_bias (bool): use bias in attention output projection.
+            attn_out_dropout_prob (float): Dropout probability for attention output.
+            attn_qkv_bias (bool): use bias for query, key, value linear layer(s).
+            bert_layer (str): BERT layer type.
+            decoder_bias (bool): use bias in decoder linear layer.
+            embed_dropout_prob (float): Dropout probability for embeddings.
+            embed_norm (bool): Normalize embedding output.
+            final_norm (bool): Add normalization after the final encoder layer and before head.
+            embedding_layer (str): Embedding layer type.
+            encoder_layer (str): Encoder layer type.
+            loss_function (str): Loss function to use.
+            loss_kwargs (dict): Keyword arguments for loss function.
+            mlp_dropout_prob (float): Dropout probability for MLP layers.
+            mlp_in_bias (bool): Use bias in MLP input linear layer.
+            mlp_layer (str): MLP layer type.
+            mlp_out_bias (bool): Use bias in MLP output linear layer.
+            norm_kwargs (dict): Keyword arguments for normalization layers.
+            normalization (str): Normalization type.
+            padding (str): Unpad inputs. Best with `use_fa2=True`.
+            head_class_act (str): Activation function for classification head.
+            head_class_bias (bool): Use bias in classification head linear layer(s).
+            head_class_dropout (float): Dropout probability for classification head.
+            head_class_norm (str): Normalization type for classification head.
+            head_pred_act (str): Activation function for prediction head.
+            head_pred_bias (bool): Use bias in prediction head linear layer(s).
+            head_pred_dropout (float): Dropout probability for prediction head.
+            head_pred_norm (bool): Normalize prediction head output.
+            pooling_type (str): Pooling type.
+            rotary_emb_dim (int | None): Rotary embedding dimension.
+            rotary_emb_base (float): Rotary embedding base.
+            rotary_emb_scale_base (float): Rotary embedding scale base.
+            rotary_emb_interleaved (bool): Use interleaved rotary embeddings.
+            use_fa2 (bool): Use FlashAttention2. Requires flash_attn package.
+            use_sdpa_attn_mask (bool): Pass a mask to SDPA. This will prevent SDPA from using the PyTorch FA2 kernel.
+            allow_embedding_resizing (bool): Embeddings will be automatically resized when they are smaller than the tokenizer vocab size.
+            init_method (str): Model layers initialization method.
+            init_std (float): Standard deviation for initialization. Used for normal and full_megatron init.
+            init_cutoff_factor (float): Cutoff factor for initialization. Used for normal and full_megatron init.
+            init_small_embedding (bool): Initialize embeddings with RWKV small init.
+            initial_attention_layer (str | None): Replace first `num_initial_layers` attention_layer instance with this layer.
+            initial_bert_layer (str | None): Replace first `num_initial_layers` bert_layer instance with this layer.
+            initial_mlp_layer (str | None): Replace first `num_initial_layers` mlp_layer instance with this layer.
+            num_initial_layers (int): Number of initial layers to set via `initial_attention_layer`, `initial_bert_layer`, and `initial_mlp_layer`.
+            skip_first_prenorm (bool): Skip pre-normalization for the first bert layer. Requires `embed_norm=True`.
+            deterministic_fa2 (bool): Use Flash Attention 2 deterministic mode. This is slower then the default non-deterministic mode.
+            sliding_window (int): Use sliding window attention with window size `n`. -1 to disable. Window size split between the left and right context. Only supports FA2.
+            global_attn_every_n_layers (int): Use global attention every `n` layers and sliding window for the rest. -1 to disable.
+            local_attn_rotary_emb_base (float): Rotary embedding base for local attention. -1 to disable and use `rotary_emb_base` for all layers.
+            local_attn_rotary_emb_dim (int | None): Rotary embedding dimension for local attention. None to disable and use `rotary_emb_dim` for all layers.
+            unpad_embeddings (bool): Unpad inputs before the embedding layer.
+            pad_logits (bool): Pad logits after the calculating the loss.
+            compile_model (bool): Compile the subset of the model which can be compiled.
+            masked_prediction (bool): Use only pass the masked tokens throught the final MLM layers
+            casual_mask (bool): Use a casual mask, defaulting to false.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(attention_probs_dropout_prob=attention_probs_dropout_prob, **kwargs)
+        self.attention_layer = attention_layer
+        self.attn_out_bias = attn_out_bias
+        self.attn_out_dropout_prob = attn_out_dropout_prob
+        self.attn_qkv_bias = attn_qkv_bias
+        self.bert_layer = bert_layer
+        self.decoder_bias = decoder_bias
+        self.embed_dropout_prob = embed_dropout_prob
+        self.embed_norm = embed_norm
+        self.final_norm = final_norm
+        self.embedding_layer = embedding_layer
+        self.encoder_layer = encoder_layer
+        self.loss_function = loss_function
+        self.loss_kwargs = loss_kwargs
+        self.mlp_dropout_prob = mlp_dropout_prob
+        self.mlp_in_bias = mlp_in_bias
+        self.mlp_layer = mlp_layer
+        self.mlp_out_bias = mlp_out_bias
+        self.norm_kwargs = norm_kwargs
+        self.normalization = normalization
+        self.padding = padding
+        self.head_class_act = head_class_act
+        self.head_class_bias = head_class_bias
+        self.head_class_dropout = head_class_dropout
+        self.head_class_norm = head_class_norm
+        self.head_pred_act = head_pred_act
+        self.head_pred_bias = head_pred_bias
+        self.head_pred_dropout = head_pred_dropout
+        self.head_pred_norm = head_pred_norm
+        self.pooling_type = pooling_type
+        self.rotary_emb_dim = rotary_emb_dim
+        self.rotary_emb_base = rotary_emb_base
+        self.rotary_emb_scale_base = rotary_emb_scale_base
+        self.rotary_emb_interleaved = rotary_emb_interleaved
+        self.use_fa2 = use_fa2
+        self.use_sdpa_attn_mask = use_sdpa_attn_mask
+        self.allow_embedding_resizing = allow_embedding_resizing
+        self.init_method = init_method
+        self.init_std = init_std
+        self.init_cutoff_factor = init_cutoff_factor
+        self.init_small_embedding = init_small_embedding
+        self.initial_attention_layer = initial_attention_layer
+        self.initial_bert_layer = initial_bert_layer
+        self.initial_mlp_layer = initial_mlp_layer
+        self.num_initial_layers = num_initial_layers
+        self.skip_first_prenorm = skip_first_prenorm
+        self.deterministic_fa2 = deterministic_fa2
+        self.sliding_window = sliding_window
+        self.global_attn_every_n_layers = global_attn_every_n_layers
+        self.local_attn_rotary_emb_base = local_attn_rotary_emb_base
+        self.local_attn_rotary_emb_dim = local_attn_rotary_emb_dim
+        self.unpad_embeddings = unpad_embeddings
+        self.pad_logits = pad_logits
+        self.compile_model = compile_model
+        self.masked_prediction = masked_prediction
+        self.casual_mask = casual_mask
+
+        if loss_kwargs.get("return_z_loss", False):
+            if loss_function != "fa_cross_entropy":
+                raise ValueError("loss_function must be 'fa_cross_entropy' when return_z_loss is True")
+            if loss_kwargs.get("lse_square_scale", 0) <= 0:
+                raise ValueError(
+                    "lse_square_scale must be passed to `loss_kwargs` and must be greater than 0 for z_loss"
+                )
+        if loss_kwargs.get("inplace_backward", False):
+            self.loss_kwargs["inplace_backward"] = False
+            warnings.warn("`inplace_backward=True` will cause incorrect metrics. Automatically setting to False.")
+
+        if global_attn_every_n_layers > 0 and (self.num_hidden_layers - 1) % global_attn_every_n_layers != 0:
+            raise ValueError(
+                f"{global_attn_every_n_layers=} must be a divisor of one less than {self.num_hidden_layers=}"
+            )
+
+        if self.sliding_window != -1:
+            if not self.use_fa2:
+                raise ValueError("Sliding window attention is only supported with FlashAttention2")
+            if self.sliding_window % 2 != 0 and self.sliding_window % 64 != 0:
+                raise ValueError(
+                    f"Sliding window must be an even number and divisible by 64: {self.sliding_window=} {self.sliding_window % 64} {self.sliding_window % 2}"
+                )
+        else:
+            if self.global_attn_every_n_layers != -1:
+                raise ValueError("global_attn_every_n_layers must be -1 when sliding_window is disabled")
+            if self.local_attn_rotary_emb_base != -1:
+                raise ValueError("local_attn_rotary_emb_base must be -1 when sliding_window is disabled")
+            if self.local_attn_rotary_emb_dim is not None:
+                raise ValueError("local_attn_rotary_emb_dim must be None when sliding_window is disabled")
+
+        if self.unpad_embeddings and self.padding != "unpadded":
+            warnings.warn(
+                "`unpad_embeddings=True` requires `padding='unpadded'`. Automatically setting `padding='unpadded'`."
+            )
+            self.padding = "unpadded"
+        if self.pad_logits and not self.unpad_embeddings:
+            raise ValueError("`pad_logits=True` requires `unpad_embeddings=True`")
+        if self.unpad_embeddings and self.embedding_layer == "absolute_pos":
+            raise ValueError(f"{self.unpad_embeddings=} is incompatible with {self.embedding_layer=}")
+
+
+PADDING = ["unpadded", "padded"]
+
+
+def maybe_add_padding(config: FlexBertConfig, config_option: str) -> str:
+    if config.padding not in PADDING:
+        raise ValueError(f"Invalid padding type: {config.padding}, must be one of {PADDING}")
+
+    if not any(config_option.startswith(pad + "_") for pad in PADDING):
+        config_option = f"{config.padding}_{config_option}"
+
+    return config_option

embeddings.py

@@ -0,0 +1,218 @@
+# Copyright 2024 **AUTHORS_TODO**
+# License: Apache-2.0
+
+# Copyright 2022 MosaicML Examples authors
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright 2023 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) 2023, Tri Dao.
+
+
+import torch
+import torch.nn as nn
+from typing import Optional
+
+from .configuration_bert import FlexBertConfig
+from .normalization import get_norm_layer
+from .initialization import ModuleType, init_weights
+
+
+class BertAlibiEmbeddings(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
+        if getattr(config, "token_type_embeddings", True):
+            self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+            self.use_token_type_embeddings = True
+        else:
+            self.use_token_type_embeddings = False
+
+        self.LayerNorm = get_norm_layer(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        if self.use_token_type_embeddings:
+            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 self.use_token_type_embeddings and token_type_ids is None:
+            if hasattr(self, "token_type_ids"):
+                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
+            else:
+                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=input_ids.device)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+
+        if self.use_token_type_embeddings:
+            token_type_embeddings = self.token_type_embeddings(token_type_ids)
+            embeddings = inputs_embeds + token_type_embeddings
+        else:
+            embeddings = inputs_embeds
+
+        # no position embeddings! ALiBi
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class FlexBertEmbeddingsBase(nn.Module):
+    """A FlexBERT embeddings base class for type hints."""
+
+    def __init__(self, config: FlexBertConfig):
+        super().__init__()
+        self.config = config
+
+    def _init_weights(self, reset_params: bool = False):
+        raise NotImplementedError("This is a base class and should not be used directly.")
+
+    def reset_parameters(self):
+        self._init_weights(reset_params=True)
+
+    def forward(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
+        raise NotImplementedError("This is a base class and should not be used directly.")
+
+
+class FlexBertAbsoluteEmbeddings(FlexBertEmbeddingsBase):
+    """Construct the embeddings with absolute positional embeddings."""
+
+    def __init__(self, config: FlexBertConfig):
+        super().__init__(config)
+        self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+
+        self.norm = get_norm_layer(config) if config.embed_norm else nn.Identity()
+        self.drop = nn.Dropout(config.embed_dropout_prob) if config.embed_dropout_prob > 0.0 else nn.Identity()
+
+        self.register_buffer(
+            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+        )
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(self.config, self.tok_embeddings, type_of_module=ModuleType.emb)
+        init_weights(self.config, self.position_embeddings, type_of_module=ModuleType.emb)
+
+        if reset_params:
+            if self.config.embed_norm:
+                self.norm.reset_parameters()  # type: ignore
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        position_ids: Optional[torch.LongTensor] = None,
+    ) -> torch.Tensor:
+        if position_ids is None:
+            position_ids = self.position_ids[:, 0 : input_ids.shape[1]]
+
+        embeddings = self.tok_embeddings(input_ids)
+        position_embeddings = self.position_embeddings(position_ids)
+
+        embeddings = self.norm(embeddings + position_embeddings)
+        return self.drop(embeddings)
+
+
+class FlexBertCompiledSansPositionEmbeddings(FlexBertEmbeddingsBase):
+    """Construct the embeddings from token embeddings without any positional embeddings."""
+
+    def __init__(self, config: FlexBertConfig):
+        super().__init__(config)
+        self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+
+        self.norm = get_norm_layer(config, compiled_norm=config.compile_model) if config.embed_norm else nn.Identity()
+        self.drop = nn.Dropout(config.embed_dropout_prob) if config.embed_dropout_prob > 0.0 else nn.Identity()
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(self.config, self.tok_embeddings, type_of_module=ModuleType.emb)
+
+        if reset_params:
+            if self.config.embed_norm:
+                self.norm.reset_parameters()  # type: ignore
+
+    @torch.compile(dynamic=True)
+    def forward(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
+        return self.drop(self.norm(self.tok_embeddings(input_ids)))
+
+
+class FlexBertSansPositionEmbeddings(FlexBertEmbeddingsBase):
+    """Construct the embeddings from token embeddings without any positional embeddings."""
+
+    def __init__(self, config: FlexBertConfig):
+        super().__init__(config)
+        self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+
+        self.norm = get_norm_layer(config) if config.embed_norm else nn.Identity()
+        self.drop = nn.Dropout(config.embed_dropout_prob) if config.embed_dropout_prob > 0.0 else nn.Identity()
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(self.config, self.tok_embeddings, type_of_module=ModuleType.emb)
+
+        if reset_params:
+            if self.config.embed_norm:
+                self.norm.reset_parameters()  # type: ignore
+
+    def forward(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
+        return self.drop(self.norm(self.tok_embeddings(input_ids)))
+
+
+EBB2CLS = {
+    "absolute_pos": FlexBertAbsoluteEmbeddings,
+    "sans_pos": FlexBertSansPositionEmbeddings,
+}
+
+
+def get_embedding_layer(config: FlexBertConfig) -> FlexBertEmbeddingsBase:
+    try:
+        if config.compile_model and config.embedding_layer == "sans_pos":
+            return FlexBertCompiledSansPositionEmbeddings(config)
+        elif config.compile_model:
+            raise ValueError(f"{config.compile_model=} only supports sans_pos embeddings.")
+        return EBB2CLS[config.embedding_layer](config)
+    except KeyError:
+        raise ValueError(f"Invalid embeddings layer type: {config.embedding_layer=}, must be one of {EBB2CLS.keys()}.")

initialization.py

@@ -0,0 +1,551 @@
+# Copyright 2022 MosaicML Examples authors
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright 2023 OLMo Authors
+# License: Apache-2.0
+
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+# License: Apache-2.0
+
+import math
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .utils import StrEnum
+
+from .configuration_bert import FlexBertConfig
+from .normalization import RMSNorm
+
+__all__ = ["init_weights", "ModuleType", "InitFnType"]
+
+
+class InitFnType(StrEnum):
+    mitchell = "mitchell"
+    """
+    The strategy suggested to us by Mitchell Wortsman from UW.
+    This uses a truncated normal distribution with an adaptive standard deviation that depends
+    on the size of the weights as well as the depth of the layer.
+    """
+
+    normal = "normal"
+    """
+    All weights are initialized from the same normal distribution.
+    """
+
+    default = "default"
+    """
+    All weights are initialized with the default HuggingFace Bert method. Set init_std=0.02 to match.
+    """
+
+    kaiming_normal = "kaiming_normal"
+    """
+    All weights are initialized with the Kaiming method from a normal distribution.
+    Note this currently won't work with FSDP.
+    """
+
+    fan_in = "fan_in"
+    """
+    "Fan-in variance scaling", i.e. normal with a standard deviation of ``1/sqrt(d_in)`` where ``d_in``
+    is the input dimensionality of the kernel.
+    """
+
+    full_megatron = "full_megatron"
+    """
+    This is what metaseq calls "full megatron init". It is the init used for Llama 2.
+    """
+
+
+class ModuleType(StrEnum):
+    in_module = "in"
+    out_module = "out"
+    emb = "emb"
+    final_out = "final_out"
+
+
+def init_weights(
+    config: FlexBertConfig,
+    module: Union[nn.Linear, nn.Embedding],
+    layer_dim: Optional[int] = None,
+    layer_id: Optional[int] = None,
+    std_factor: float = 1.0,
+    type_of_module: Optional[ModuleType] = None,
+) -> None:
+    """
+    Initialize weights of a linear or embedding module.
+
+    :param config: The model config.
+    :param module: The linear or embedding submodule to initialize.
+    :param layer_dim: The effective input dimensionality of the weights. This could be smaller than the actual dimensions
+        for fused layers.
+    :param layer_id: When set, the standard deviation for the "mitchell" method will be adjusted by
+        ``1 / sqrt(2 * (layer_id + 1))``.
+    """
+    if config.init_method == InitFnType.full_megatron and config.init_small_embedding:
+        raise ValueError("Cannot use 'small_embedding_init' with 'full_megatron' init.")
+
+    layer_dim = layer_dim if layer_dim is not None else config.hidden_size
+    if config.init_method == InitFnType.normal:
+        std = config.init_std * std_factor
+        if config.init_cutoff_factor is not None:
+            cutoff_value = config.init_cutoff_factor * std
+            nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-cutoff_value, b=cutoff_value)
+        else:
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+    elif config.init_method == InitFnType.mitchell:
+        std = std_factor / math.sqrt(layer_dim)
+        if layer_id is not None:
+            std = std / math.sqrt(2 * (layer_id + 1))
+        nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-3 * std, b=3 * std)
+    elif config.init_method == InitFnType.kaiming_normal:
+        nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
+    elif config.init_method == InitFnType.fan_in:
+        std = std_factor / math.sqrt(layer_dim)
+        nn.init.normal_(module.weight, mean=0.0, std=std)
+    elif config.init_method == InitFnType.full_megatron:
+        if type_of_module is None:
+            raise RuntimeError(f"When using the {InitFnType.full_megatron} init, every module must have a type.")
+
+        cutoff_factor = config.init_cutoff_factor
+        if cutoff_factor is None:
+            cutoff_factor = 3
+
+        if type_of_module == ModuleType.in_module:
+            # for att_proj (same as QKV), ff_proj
+            std = config.init_std
+        elif type_of_module == ModuleType.out_module:
+            # for attn_out, ff_out
+            std = config.init_std / math.sqrt(2.0 * config.num_hidden_layers)
+        elif type_of_module == ModuleType.emb:
+            # positional embeddings (wpe)
+            # token embeddings (wte)
+            std = config.init_std
+        elif type_of_module == ModuleType.final_out:
+            # final output (ff_out)
+            std = config.hidden_size**-0.5
+        else:
+            raise RuntimeError(f"Unknown module type '{type_of_module}'")
+
+        nn.init.trunc_normal_(
+            module.weight,
+            mean=0.0,
+            std=std,
+            a=-cutoff_factor * std,
+            b=cutoff_factor * std,
+        )
+    elif config.init_method == InitFnType.default:
+        # default hugging face bert initialization
+        # normalization layers already init to ones and zeros
+        if isinstance(module, nn.Linear):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=config.init_std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=config.init_std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+    else:
+        raise NotImplementedError(config.init_method)
+
+    if isinstance(module, nn.Linear):
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+        if config.init_method == InitFnType.normal and getattr(module, "_is_residual", False):
+            with torch.no_grad():
+                module.weight.div_(math.sqrt(2 * config.num_hidden_layers))
+
+    if isinstance(module, nn.Embedding) and config.init_small_embedding:
+        nn.init.uniform_(module.weight, a=-1e-4, b=1e-4)
+
+
+class TileMode(StrEnum):
+    center_weights = "center_weights"
+    tile_weights_from_edge = "tile_weights_from_edge"
+    tile_weights_from_middle = "tile_weights_from_middle"
+
+
+def tile_weight(
+    pretrained_weights: torch.Tensor,
+    new_weights: torch.Tensor,
+    mode: Union[str, TileMode] = TileMode.tile_weights_from_middle,
+) -> torch.Tensor:
+    """
+    Tile or center an input tensor to a larger desired size. Works for both 2D and 1D tensors.
+
+    Args:
+    pretrained_weights (torch.Tensor): The input tensor to be tiled or centered (1D or 2D).
+    new_weights (torch.Tensor): The tensor with the desired size.
+    mode (Union[str, TileMode]): 'center_weights', 'tile_weights_from_edge', or 'tile_weights_from_middle'
+
+    Returns:
+    torch.Tensor: The resulting tensor of the desired size.
+    """
+    assert pretrained_weights.dim() in (1, 2), "Input tensor must be 1-dimensional or 2-dimensional"
+    if isinstance(mode, str):
+        mode = TileMode(mode)
+
+    pretrained_weights = pretrained_weights.clone()
+
+    if pretrained_weights.dim() == 1:
+        return _tile_1d(pretrained_weights, new_weights, mode)
+    else:
+        return _tile_2d(pretrained_weights, new_weights, mode)
+
+
+def _tile_1d(pretrained_weights: torch.Tensor, new_weights: torch.Tensor, mode: TileMode) -> torch.Tensor:
+    assert pretrained_weights.dim() == 1, "Input tensor must be 1-dimensional"
+    input_size = pretrained_weights.shape[0]
+    new_size = new_weights.shape[0]
+    assert new_size >= input_size, "Desired size must be greater than or equal to input size"
+
+    if mode == TileMode.center_weights:
+        offset = (new_size - input_size) // 2
+        new_weights[offset : offset + input_size] = pretrained_weights
+        return new_weights.clone()
+    elif mode == TileMode.tile_weights_from_edge:
+        repeat_count = (new_size + input_size - 1) // input_size
+        tiled_tensor = pretrained_weights.repeat(repeat_count)
+        return tiled_tensor[:new_size].clone()
+    elif mode == TileMode.tile_weights_from_middle:
+        # Calculate offsets to center the original tensor
+        offset = (new_size - input_size) // 2
+
+        # Create a new tensor with the desired size
+        result = torch.zeros(new_size, dtype=pretrained_weights.dtype, device=pretrained_weights.device)
+
+        # Place the original tensor in the center
+        result[offset : offset + input_size] = pretrained_weights
+
+        # Tile the left and right sides
+        for i in range(offset):
+            result[offset - 1 - i] = pretrained_weights[input_size - 1 - (i % input_size)]
+        for i in range(offset + input_size, new_size):
+            result[i] = pretrained_weights[(i - offset) % input_size]
+        return result.clone()
+
+
+def _tile_2d(pretrained_weights: torch.Tensor, new_weights: torch.Tensor, mode: TileMode) -> torch.Tensor:
+    assert pretrained_weights.dim() == 2, "Input tensor must be 2-dimensional"
+    input_height, input_width = pretrained_weights.shape
+    new_height, new_width = new_weights.shape
+    assert new_height >= input_height, "Desired height must be greater than or equal to input height"
+    assert new_width >= input_width, "Desired width must be greater than or equal to input width"
+
+    if mode == TileMode.center_weights:
+        height_offset = (new_height - input_height) // 2
+        width_offset = (new_width - input_width) // 2
+        new_weights[height_offset : height_offset + input_height, width_offset : width_offset + input_width] = pretrained_weights  # fmt: skip
+        return new_weights.clone()
+    elif mode == TileMode.tile_weights_from_edge:
+        repeat_height = (new_height + input_height - 1) // input_height
+        repeat_width = (new_width + input_width - 1) // input_width
+        tiled_tensor = pretrained_weights.repeat(repeat_height, repeat_width)
+        return tiled_tensor[:new_height, :new_width].clone()
+    elif mode == TileMode.tile_weights_from_middle:
+        # Calculate offsets to center the original tensor
+        height_offset = (new_height - input_height) // 2
+        width_offset = (new_width - input_width) // 2
+
+        # Create a new tensor with the desired width and input height
+        horizontal_tiled = torch.zeros(
+            input_height, new_width, dtype=pretrained_weights.dtype, device=pretrained_weights.device
+        )
+
+        # Place the original tensor in the center horizontally
+        horizontal_tiled[:, width_offset : width_offset + input_width] = pretrained_weights
+
+        # Tile the left and right sides
+        for i in range(width_offset):
+            horizontal_tiled[:, i] = horizontal_tiled[
+                :, width_offset + input_width - 1 - (width_offset - i - 1) % input_width
+            ]
+        for i in range(width_offset + input_width, new_width):
+            horizontal_tiled[:, i] = horizontal_tiled[:, width_offset + (i - width_offset) % input_width]
+
+        # Now tile vertically
+        result = torch.zeros(new_height, new_width, dtype=pretrained_weights.dtype, device=pretrained_weights.device)
+        result[height_offset : height_offset + input_height, :] = horizontal_tiled
+
+        # Tile top
+        for i in range(height_offset):
+            row_to_copy = (input_height - 1) - (i % input_height)
+            result[height_offset - 1 - i, :] = horizontal_tiled[row_to_copy, :]
+
+        # Tile bottom
+        for i in range(height_offset + input_height, new_height):
+            row_to_copy = (i - height_offset) % input_height
+            result[i, :] = horizontal_tiled[row_to_copy, :]
+        return result.clone()
+
+
+def tile_fused_qkv(
+    pretrained_qkv_weight: torch.Tensor,
+    new_qkv_weight: torch.Tensor,
+    mode: Union[str, TileMode] = TileMode.tile_weights_from_middle,
+):
+    """
+    Tile the weights of a fused pretrained QKV layer to a new, larger QKV dimension.
+
+    Args:
+        pretrained_qkv_weight (torch.Tensor): The original fused QKV layer
+        new_qkv_weight (torch.Tensor): The new fused QKV layer with larger linear_dim
+        mode (Union[str, TileMode]): The tiling mode to use
+    Returns:
+        torch.Tensor: The new fused QKV layer with tiled weights
+    """
+    # Split QKV, assume new_q, new_k, new_v are the same shape
+    pretrained_q, pretrained_k, pretrained_v = pretrained_qkv_weight.chunk(3, dim=0)
+    new_q, new_k, new_v = new_qkv_weight.chunk(3, dim=0)
+
+    # Tile Q, K, V separately
+    new_q = tile_weight(pretrained_q, new_q, mode=mode)
+    new_k = tile_weight(pretrained_k, new_k, mode=mode)
+    new_v = tile_weight(pretrained_v, new_v, mode=mode)
+
+    # Concatenate tiled Q, K, V
+    return torch.cat([new_q, new_k, new_v], dim=0)
+
+
+def tile_fused_glu(
+    pretrained_glu_weight: torch.Tensor,
+    new_glu_weight: torch.Tensor,
+    mode: Union[str, TileMode] = TileMode.tile_weights_from_middle,
+):
+    """
+    Tile the weights of a fused pretrained GLU layer to a new, larger GLU dimension.
+
+    Args:
+        pretrained_glu_weight (torch.Tensor): The original fused GLU layer
+        new_glu_weight (torch.Tensor): The new fused GLU layer with larger linear_dim
+        mode (Union[str, TileMode]): The tiling mode to use
+    Returns:
+        torch.Tensor: The new fused GLU layer with tiled weights
+    """
+    # Split GLU, assume new_glu_wi, new_glu_wg are the same shape
+    pretrained_glu_wi, pretrained_glu_wg = pretrained_glu_weight.chunk(2, dim=0)
+    new_glu_wi, new_glu_wg = new_glu_weight.chunk(2, dim=0)
+
+    # Tile GLU separately
+    new_glu_wi = tile_weight(pretrained_glu_wi, new_glu_wi, mode=mode)
+    new_glu_wg = tile_weight(pretrained_glu_wg, new_glu_wg, mode=mode)
+
+    # Concatenate tiled GLU
+    return torch.cat([new_glu_wi, new_glu_wg], dim=0)
+
+
+def tile_fused_qkvff(
+    pretrained_qkvff_weight: torch.Tensor,
+    new_qkvff_weight: torch.Tensor,
+    pretrained_attn_size: int,
+    pretrained_mlp_size: int,
+    new_attn_size: int,
+    new_mlp_size: int,
+    is_glu: bool = False,
+    mode: Union[str, TileMode] = TileMode.tile_weights_from_middle,
+):
+    """
+    Tile the weights of a fused pretrained QKVFF layer to a new, larger QKVFF dimension.
+
+    Args:
+        pretrained_qkvff_weight (torch.Tensor): The original fused QKVFF layer
+        new_qkvff_weight (torch.Tensor): The new fused QKVFF layer with larger linear_dim
+        pretrained_attn_size (int): The attention size of the pretrained fused QKVFF layer
+        pretrained_mlp_size (int): The mlp size of the pretrained fused QKVFF layer
+        new_attn_size (int): The attention size of the new fused QKVFF layer
+        new_mlp_size (int): The mlp size of the new fused QKVFF layer
+        is_glu (bool): Whether the QKVFF layer is a GLU layer
+        mode (Union[str, TileMode]): The tiling mode to use
+    Returns:
+        torch.Tensor: The new fused QKVFF layer with tiled weights
+    """
+    # Split QKVFF
+    pretrained_qkv, pretrained_ff = pretrained_qkvff_weight.split([pretrained_attn_size, pretrained_mlp_size], dim=0)
+    new_qkv, new_ff = new_qkvff_weight.split([new_attn_size, new_mlp_size], dim=0)
+
+    # Tile QKVFF separately
+    new_qkv = tile_fused_qkv(pretrained_qkv, new_qkv, mode=mode)
+    if is_glu:
+        new_ff = tile_fused_glu(pretrained_ff, new_ff, mode=mode)
+    else:
+        new_ff = tile_weight(pretrained_ff, new_ff, mode=mode)
+
+    # Concatenate tiled QKVFF
+    return torch.cat([new_qkv, new_ff], dim=0)
+
+
+class TileLinear(StrEnum):
+    wqkv = "wqkv"
+    glu = "glu"
+    wqkvff = "wqkvff"
+    default = "default"
+
+
+def tile_linear(
+    pretrained_linear: nn.Linear,
+    new_linear: nn.Linear,
+    linear_type: Union[str, TileLinear] = TileLinear.default,
+    mode: Union[str, TileMode] = TileMode.tile_weights_from_middle,
+    pretrained_attn_size: Optional[int] = None,
+    pretrained_mlp_size: Optional[int] = None,
+    new_attn_size: Optional[int] = None,
+    new_mlp_size: Optional[int] = None,
+    wqkvff_is_glu: Optional[bool] = None,
+    bias_only: Optional[bool] = False,
+):
+    """
+    Tile the weights of a linear layer to a new, larger linear dimension.
+
+    Args:
+        pretrained_linear (nn.Linear): The original linear layer
+        new_linear (nn.Linear): The new linear layer with larger linear_dim
+        linear_type (Union[str, TileLinear]): The type of linear layer to tile
+        mode (Union[str, TileMode]): The tiling mode to use
+        pretrained_attn_size (int): The attention size of the pretrained linear layer. Only used if linear_type is wqkvff.
+        pretrained_mlp_size (int): The mlp size of the pretrained linear layer. Only used if linear_type is wqkvff.
+        new_attn_size (int): The attention size of the new linear layer. Only used if linear_type is wqkvff.
+        new_mlp_size (int): The mlp size of the new linear layer. Only used if linear_type is wqkvff.
+        wqkvff_is_glu (bool): Whether the wqkvff layer is a GLU layer. Only used if linear_type is wqkvff.
+        bias_only (bool): Whether to only tile the bias. Only used if tiling weight tied decoder.
+    """
+    if isinstance(linear_type, str):
+        linear_type = TileLinear(linear_type)
+    if isinstance(mode, str):
+        mode = TileMode(mode)
+
+    with torch.no_grad():
+        if linear_type == TileLinear.wqkv:
+            if not bias_only:
+                new_linear.weight = nn.Parameter(
+                    tile_fused_qkv(pretrained_linear.weight, new_linear.weight, mode=mode),
+                    requires_grad=new_linear.weight.requires_grad,
+                )
+            if pretrained_linear.bias is not None:
+                new_linear.bias = nn.Parameter(
+                    tile_fused_qkv(pretrained_linear.bias, new_linear.bias, mode=mode),
+                    requires_grad=new_linear.bias.requires_grad,
+                )
+        elif linear_type == TileLinear.glu:
+            if not bias_only:
+                new_linear.weight = nn.Parameter(
+                    tile_fused_glu(pretrained_linear.weight, new_linear.weight, mode=mode),
+                    requires_grad=new_linear.weight.requires_grad,
+                )
+            if pretrained_linear.bias is not None:
+                new_linear.bias = nn.Parameter(
+                    tile_fused_glu(pretrained_linear.bias, new_linear.bias, mode=mode),
+                    requires_grad=new_linear.bias.requires_grad,
+                )
+        elif linear_type == TileLinear.wqkvff:
+            if not bias_only:
+                new_linear.weight = nn.Parameter(
+                    tile_fused_qkvff(
+                        pretrained_linear.weight,
+                        new_linear.weight,
+                        pretrained_attn_size,
+                        pretrained_mlp_size,
+                        new_attn_size,
+                        new_mlp_size,
+                        wqkvff_is_glu,
+                        mode=mode,
+                    ),
+                    requires_grad=new_linear.weight.requires_grad,
+                )
+            if pretrained_linear.bias is not None:
+                new_linear.bias = nn.Parameter(
+                    tile_fused_qkvff(
+                        pretrained_linear.bias,
+                        new_linear.bias,
+                        pretrained_attn_size,
+                        pretrained_mlp_size,
+                        new_attn_size,
+                        new_mlp_size,
+                        wqkvff_is_glu,
+                        mode=mode,
+                    ),
+                    requires_grad=new_linear.bias.requires_grad,
+                )
+        else:
+            if not bias_only:
+                new_linear.weight = nn.Parameter(
+                    tile_weight(pretrained_linear.weight, new_linear.weight, mode=mode),
+                    requires_grad=new_linear.weight.requires_grad,
+                )
+            if pretrained_linear.bias is not None:
+                new_linear.bias = nn.Parameter(
+                    tile_weight(pretrained_linear.bias, new_linear.bias, mode=mode),
+                    requires_grad=new_linear.bias.requires_grad,
+                )
+
+
+def tile_norm(
+    pretrained_norm: Union[nn.LayerNorm, RMSNorm, nn.Identity],
+    new_norm: Union[nn.LayerNorm, RMSNorm, nn.Identity],
+    mode: Union[str, TileMode] = TileMode.tile_weights_from_middle,
+):
+    """
+    Tile the weights of a pretrained norm layer to a new, larger layer norm dimension.
+
+    Args:
+        pretrained_norm (Union[nn.LayerNorm, RMSNorm, nn.Identity]): The original norm layer
+        new_norm (Union[nn.LayerNorm, RMSNorm, nn.Identity]): The new norm layer with larger layer norm dimension
+        mode (Union[str, TileMode]): The Phi-style weight tiling mode to use
+    """
+    if isinstance(pretrained_norm, nn.Identity):
+        return
+    if isinstance(mode, str):
+        mode = TileMode(mode)
+
+    with torch.no_grad():
+        new_norm.weight.data = nn.Parameter(
+            tile_weight(pretrained_norm.weight, new_norm.weight, mode=mode),
+            requires_grad=new_norm.weight.requires_grad,
+        )
+        if hasattr(pretrained_norm, "bias") and pretrained_norm.bias is not None:
+            new_norm.bias.data = nn.Parameter(
+                tile_weight(pretrained_norm.bias, new_norm.bias, mode=mode),
+                requires_grad=new_norm.bias.requires_grad,
+            )
+
+
+def tile_embedding(
+    pretrained_embedding: nn.Embedding,
+    new_embedding: nn.Embedding,
+    mode: Union[str, TileMode] = TileMode.tile_weights_from_middle,
+) -> nn.Embedding:
+    """
+    Tile the weights of an embedding layer to a new, larger embedding dimension.
+
+    Args:
+    pretrained_embedding (nn.Embedding): The original embedding layer
+    new_embedding (nn.Embedding): The new embedding layer with larger embedding_dim
+    tile_mode (Union[str, TileMode]): The Phi-style weight tiling mode to use
+
+    Returns:
+    nn.Embedding: The new embedding layer with tiled weights
+    """
+    with torch.no_grad():
+        # Ensure vocabulary size remains the same
+        if pretrained_embedding.num_embeddings != new_embedding.num_embeddings:
+            raise ValueError("Vocabulary size (num_embeddings) must remain constant")
+
+        # Ensure new embedding dimension is larger
+        if new_embedding.embedding_dim <= pretrained_embedding.embedding_dim:
+            raise ValueError("New embedding_dim must be larger than the old embedding_dim")
+
+        # Tile the weights
+        new_embedding.weight.data = nn.Parameter(
+            tile_weight(pretrained_embedding.weight, new_embedding.weight, mode=mode),
+            requires_grad=new_embedding.weight.requires_grad,
+        )
+
+        # Handle padding_idx if it exists
+        if pretrained_embedding.padding_idx is not None:
+            if new_embedding.padding_idx is None:
+                new_embedding.padding_idx = pretrained_embedding.padding_idx
+            else:
+                assert new_embedding.padding_idx == pretrained_embedding.padding_idx, "padding_idx must remain the same"

layers.py

@@ -0,0 +1,700 @@
+# Copyright 2024 **AUTHORS_TODO**
+# License: Apache-2.0
+
+# Copyright 2022 MosaicML Examples authors
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright 2023 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) 2023, Tri Dao.
+
+
+import copy
+import math
+import warnings
+from typing import Optional, Union, List
+
+import torch
+import torch.nn as nn
+
+from .bert_padding import unpad_input, pad_input
+
+from .activation import get_act_fn
+from .attention import FlexBertAttentionBase, BertAlibiUnpadAttention, get_attention_layer
+from .mlp import FlexBertMLPBase, BertResidualGLU, get_mlp_layer
+from .configuration_bert import FlexBertConfig, maybe_add_padding
+from .normalization import get_norm_layer
+from .initialization import ModuleType, init_weights
+
+
+class BertAlibiLayer(nn.Module):
+    """Composes the Mosaic BERT attention and FFN blocks into a single layer."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.attention = BertAlibiUnpadAttention(config)
+        self.mlp = BertResidualGLU(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,
+        slopes: 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)
+            slopes: None or (batch, heads) or (heads,)
+        """
+        assert (bias is None) == (slopes is None), f"{bias=}, {slopes=}"
+        attention_output = self.attention(
+            hidden_states, cu_seqlens, seqlen, subset_idx, indices, attn_mask, bias, slopes
+        )
+        layer_output = self.mlp(attention_output)
+        return layer_output
+
+
+class BertAlibiEncoder(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.BertAlibiEncoder`,
+    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 = BertAlibiLayer(config)
+        self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])
+
+        self.num_attention_heads = config.num_attention_heads
+
+        # The alibi mask will be dynamically expanded if it is too small for
+        # the input the model receives. But it generally helps to initialize it
+        # to a reasonably large size to help pre-allocate CUDA memory.
+        # The default `alibi_starting_size` is 512.
+        self._current_alibi_size = int(config.alibi_starting_size)
+        self.alibi = torch.zeros((1, self.num_attention_heads, self._current_alibi_size, self._current_alibi_size))
+        self.rebuild_alibi_tensor(size=config.alibi_starting_size)
+
+    def rebuild_alibi_tensor(self, size: int, device: Optional[Union[torch.device, str]] = None):
+        # Alibi
+        # Following https://github.com/ofirpress/attention_with_linear_biases/issues/5 (Implementation 1)
+        # In the causal case, you can exploit the fact that softmax is invariant to a uniform translation
+        # of the logits, which makes the math work out *after* applying causal masking. If no causal masking
+        # will be applied, it is necessary to construct the diagonal mask.
+        n_heads = self.num_attention_heads
+
+        def _get_alibi_head_slopes(n_heads: int) -> List[float]:
+            def get_slopes_power_of_2(n_heads: int) -> List[float]:
+                start = 2 ** (-(2 ** -(math.log2(n_heads) - 3)))
+                ratio = start
+                return [start * ratio**i for i in range(n_heads)]
+
+            # In the paper, they only train models that have 2^a heads for some a. This function
+            # has some good properties that only occur when the input is a power of 2. To
+            # maintain that even when the number of heads is not a power of 2, we use a
+            # workaround.
+            if math.log2(n_heads).is_integer():
+                return get_slopes_power_of_2(n_heads)
+
+            closest_power_of_2 = 2 ** math.floor(math.log2(n_heads))
+            slopes_a = get_slopes_power_of_2(closest_power_of_2)
+            slopes_b = _get_alibi_head_slopes(2 * closest_power_of_2)
+            slopes_b = slopes_b[0::2][: n_heads - closest_power_of_2]
+            return slopes_a + slopes_b
+
+        context_position = torch.arange(size, device=device)[:, None]
+        memory_position = torch.arange(size, device=device)[None, :]
+        relative_position = torch.abs(memory_position - context_position)
+        # [n_heads, max_token_length, max_token_length]
+        relative_position = relative_position.unsqueeze(0).expand(n_heads, -1, -1)
+        slopes = torch.Tensor(_get_alibi_head_slopes(n_heads)).to(device)
+        self.slopes = slopes
+        alibi = slopes.unsqueeze(1).unsqueeze(1) * -relative_position
+        # [1, n_heads, max_token_length, max_token_length]
+        alibi = alibi.unsqueeze(0)
+        assert alibi.shape == torch.Size([1, n_heads, size, size])
+
+        self._current_alibi_size = size
+        self.alibi = alibi
+
+    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)
+
+        # Add alibi matrix to extended_attention_mask
+        if self._current_alibi_size < seqlen:
+            # Rebuild the alibi tensor when needed
+            warnings.warn(f"Increasing alibi size from {self._current_alibi_size} to {seqlen}")
+            self.rebuild_alibi_tensor(size=seqlen, device=hidden_states.device)
+        elif self.alibi.device != hidden_states.device:
+            # Device catch-up
+            self.alibi = self.alibi.to(hidden_states.device)
+            self.slopes = self.slopes.to(hidden_states.device)  # type: ignore
+        alibi_bias = self.alibi[:, :, :seqlen, :seqlen]
+        attn_bias = extended_attention_mask[:, :, :seqlen, :seqlen]
+        alibi_attn_mask = attn_bias + alibi_bias
+
+        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=alibi_attn_mask,
+                    slopes=self.slopes,
+                )
+                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=alibi_attn_mask,
+                    slopes=self.slopes,
+                )
+                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=alibi_attn_mask,
+                slopes=self.slopes,
+            )
+
+        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 = get_act_fn(config.head_pred_act)
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.LayerNorm = get_norm_layer(config)
+
+    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.LayerNorm(hidden_states)
+        return hidden_states
+
+
+class FlexBertLayerBase(nn.Module):
+    """A FlexBERT Layer base class for type hints."""
+
+    attn: FlexBertAttentionBase
+    mlp: FlexBertMLPBase
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_id = layer_id
+
+    def _init_weights(self, reset_params: bool = False):
+        if hasattr(self, "attn"):
+            self.attn._init_weights(reset_params)
+        if hasattr(self, "mlp"):
+            self.mlp._init_weights(reset_params)
+
+    def reset_parameters(self):
+        self._init_weights(reset_params=True)
+
+    def forward(self, hidden_states: torch.Tensor, attn_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
+        raise NotImplementedError("This is a base class and should not be used directly.")
+
+
+class FlexBertCompileUnpadPreNormLayer(FlexBertLayerBase):
+    """Composes the FlexBERT attention and MLP blocks into a single layer using pre-normalization."""
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        if config.skip_first_prenorm and config.embed_norm and layer_id == 0:
+            self.attn_norm = nn.Identity()
+        else:
+            self.attn_norm = get_norm_layer(config)
+        self.attn = get_attention_layer(config, layer_id=layer_id)
+        self.mlp_norm = get_norm_layer(config, compiled_norm=config.compile_model)
+        self.mlp = get_mlp_layer(config, layer_id=layer_id)
+        self.compile_model = config.compile_model
+
+    def _init_weights(self, reset_params: bool = False):
+        super()._init_weights(reset_params)
+        if reset_params:
+            self.attn_norm.reset_parameters()
+            self.mlp_norm.reset_parameters()
+
+    @torch.compile(dynamic=True)
+    def compiled_mlp(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.mlp(self.mlp_norm(hidden_states))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+        indices: Optional[torch.Tensor] = None,
+        attn_mask: 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,)
+            max_seqlen: int
+            indices: None or (total_nnz,)
+            attn_mask: None or (batch, max_seqlen)
+        """
+        attn_out = hidden_states + self.attn(self.attn_norm(hidden_states), cu_seqlens, max_seqlen, indices, attn_mask)
+        return attn_out + self.compiled_mlp(attn_out)
+
+
+class FlexBertUnpadPreNormLayer(FlexBertLayerBase):
+    """Composes the FlexBERT attention and MLP blocks into a single layer using pre-normalization."""
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        if config.skip_first_prenorm and config.embed_norm and layer_id == 0:
+            self.attn_norm = nn.Identity()
+        else:
+            self.attn_norm = get_norm_layer(config)
+        self.attn = get_attention_layer(config, layer_id=layer_id)
+        self.mlp_norm = get_norm_layer(config)
+        self.mlp = get_mlp_layer(config, layer_id=layer_id)
+
+    def _init_weights(self, reset_params: bool = False):
+        super()._init_weights(reset_params)
+        if reset_params:
+            self.attn_norm.reset_parameters()
+            self.mlp_norm.reset_parameters()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+        indices: Optional[torch.Tensor] = None,
+        attn_mask: 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,)
+            max_seqlen: int
+            indices: None or (total_nnz,)
+            attn_mask: None or (batch, max_seqlen)
+        """
+        attn_out = hidden_states + self.attn(self.attn_norm(hidden_states), cu_seqlens, max_seqlen, indices, attn_mask)
+        return attn_out + self.mlp(self.mlp_norm(attn_out))
+
+
+class FlexBertUnpadParallelPreNormLayer(FlexBertLayerBase):
+    """Composes the FlexBERT parallel attention and MLP blocks into a single layer using pre-normalization."""
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        self.attn_size = config.hidden_size * 3
+        self.mlp_size = config.intermediate_size * 2
+        # Compute QKV and FF outputs at once
+        self.Wqkvff = nn.Linear(config.hidden_size, self.attn_size + self.mlp_size, bias=config.attn_qkv_bias)
+        if config.skip_first_prenorm and config.embed_norm and layer_id == 0:
+            self.norm = nn.Identity()
+        else:
+            self.norm = get_norm_layer(config)
+        self.attn = get_attention_layer(config, layer_id=layer_id)
+        self.mlp = get_mlp_layer(config, layer_id=layer_id)
+
+    def _init_weights(self, reset_params: bool = False):
+        super()._init_weights(reset_params)
+        if reset_params and hasattr(self.norm, "reset_parameters"):
+            self.norm.reset_parameters()
+
+        init_weights(
+            self.config,
+            self.Wqkvff,
+            layer_dim=self.config.hidden_size,
+            layer_id=None,
+            type_of_module=ModuleType.in_module,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+        indices: Optional[torch.Tensor] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward pass for a BERT layer, including both attention and MLP.
+
+        Args:
+            hidden_states: (total_nnz, dim)
+            attn_mask: None or (batch, max_seqlen)
+        """
+        # Compute QKV and FF outputs at once and split them
+        qkv, intermediate_ff = self.Wqkvff(self.norm(hidden_states)).split([self.attn_size, self.mlp_size], dim=1)
+        return hidden_states + self.attn(qkv, cu_seqlens, max_seqlen, indices, attn_mask) + self.mlp(intermediate_ff)
+
+
+class FlexBertPaddedPreNormLayer(FlexBertLayerBase):
+    """Composes the FlexBERT attention and MLP blocks into a single layer using pre-normalization."""
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        if config.skip_first_prenorm and config.embed_norm and layer_id == 0:
+            self.attn_norm = nn.Identity()
+        else:
+            self.attn_norm = get_norm_layer(config)
+        self.attn = get_attention_layer(config, layer_id=layer_id)
+        self.mlp_norm = get_norm_layer(config)
+        self.mlp = get_mlp_layer(config, layer_id=layer_id)
+
+    def _init_weights(self, reset_params: bool = False):
+        super()._init_weights(reset_params)
+        if reset_params:
+            self.attn_norm.reset_parameters()
+            self.mlp_norm.reset_parameters()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward pass for a BERT layer, including both attention and MLP.
+
+        Args:
+            hidden_states: (batch, max_seqlen, dim)
+            attn_mask: None or (batch, max_seqlen)
+        """
+        attn_out = hidden_states + self.attn(self.attn_norm(hidden_states), attn_mask)
+        return attn_out + self.mlp(self.mlp_norm(attn_out))
+
+
+class FlexBertPaddedParallelPreNormLayer(FlexBertLayerBase):
+    """Composes the FlexBERT attention and MLP blocks into a single layer using pre-normalization."""
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        self.attn_size = config.hidden_size * 3
+        self.mlp_size = config.intermediate_size * 2
+        # Compute QKV and FF outputs at once
+        self.Wqkvff = nn.Linear(config.hidden_size, self.attn_size + self.mlp_size, bias=config.attn_qkv_bias)
+        if config.skip_first_prenorm and config.embed_norm and layer_id == 0:
+            self.norm = nn.Identity()
+        else:
+            self.norm = get_norm_layer(config)
+        self.attn = get_attention_layer(config, layer_id=layer_id)
+        self.mlp = get_mlp_layer(config, layer_id=layer_id)
+
+    def _init_weights(self, reset_params: bool = False):
+        super()._init_weights(reset_params)
+        if reset_params:
+            self.norm.reset_parameters()
+
+        init_weights(
+            self.config,
+            self.Wqkvff,
+            layer_dim=self.config.hidden_size,
+            layer_id=None,
+            type_of_module=ModuleType.in_module,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward pass for a BERT layer, including both attention and MLP.
+
+        Args:
+            hidden_states: (batch, max_seqlen, dim)
+            attn_mask: None or (batch, max_seqlen)
+        """
+        # Compute QKV and FF outputs at once and split them
+        qkv, intermediate_ff = self.Wqkvff(self.norm(hidden_states)).split([self.attn_size, self.mlp_size], dim=2)
+        return hidden_states + self.attn(qkv, attn_mask) + self.mlp(intermediate_ff)
+
+
+class FlexBertUnpadPostNormLayer(FlexBertLayerBase):
+    """Composes the FlexBERT attention and MLP blocks into a single layer using post-normalization."""
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        self.attn = get_attention_layer(config, layer_id=layer_id)
+        self.attn_norm = get_norm_layer(config)
+        self.mlp = get_mlp_layer(config, layer_id=layer_id)
+        self.mlp_norm = get_norm_layer(config)
+
+    def _init_weights(self, reset_params: bool = False):
+        super()._init_weights(reset_params)
+        if reset_params:
+            self.attn_norm.reset_parameters()
+            self.mlp_norm.reset_parameters()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+        indices: Optional[torch.Tensor] = None,
+        attn_mask: 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,)
+            max_seqlen: int
+            indices: None or (total_nnz,)
+            attn_mask: None or (batch, max_seqlen)
+        """
+        attn_out = self.attn_norm(hidden_states + self.attn(hidden_states, cu_seqlens, max_seqlen, indices, attn_mask))
+        return self.mlp_norm(attn_out + self.mlp(attn_out))
+
+
+class FlexBertPaddedPostNormLayer(FlexBertLayerBase):
+    """Composes the FlexBERT attention and MLP blocks into a single layer using post-normalization."""
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        self.attn = get_attention_layer(config, layer_id=layer_id)
+        self.attn_norm = get_norm_layer(config)
+        self.mlp = get_mlp_layer(config, layer_id=layer_id)
+        self.mlp_norm = get_norm_layer(config)
+
+    def _init_weights(self, reset_params: bool = False):
+        super()._init_weights(reset_params)
+        if reset_params:
+            self.mlp_norm.reset_parameters()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward pass for a BERT layer, including both attention and MLP.
+
+        Args:
+            hidden_states: (batch, max_seqlen, dim)
+            attn_mask: None or (batch, max_seqlen)
+        """
+        attn_out = self.attn_norm(hidden_states + self.attn(hidden_states, attn_mask))
+        return self.mlp_norm(attn_out + self.mlp(attn_out))
+
+
+LAYER2CLS = {
+    "unpadded_prenorm": FlexBertUnpadPreNormLayer,
+    "unpadded_compile_prenorm": FlexBertCompileUnpadPreNormLayer,
+    "unpadded_parallel_prenorm": FlexBertUnpadParallelPreNormLayer,
+    "unpadded_postnorm": FlexBertUnpadPostNormLayer,
+    "padded_prenorm": FlexBertPaddedPreNormLayer,
+    "padded_parallel_prenorm": FlexBertPaddedParallelPreNormLayer,
+    "padded_postnorm": FlexBertPaddedPostNormLayer,
+}
+
+
+def get_bert_layer(config: FlexBertConfig, layer_id: Optional[int] = None) -> FlexBertLayerBase:
+    try:
+        bert_layer = (
+            config.initial_bert_layer
+            if layer_id < config.num_initial_layers and getattr(config, "initial_bert_layer", None) is not None
+            else config.bert_layer
+        )
+        bert_layer = maybe_add_padding(config, bert_layer)
+        if config.compile_model and bert_layer == "unpadded_prenorm":
+            bert_layer = "unpadded_compile_prenorm"
+        return LAYER2CLS[bert_layer](config, layer_id=layer_id)
+    except KeyError:
+        if layer_id < config.num_initial_layers and getattr(config, "initial_bert_layer", None) is not None:
+            raise ValueError(
+                f"Invalid BERT layer type: {config.initial_bert_layer=}, must be one of {LAYER2CLS.keys()}."
+                f"{config.padding=} will be automatically prepended to `config.bert_layer` if unspecified."
+            )
+        else:
+            raise ValueError(
+                f"Invalid BERT layer type: {config.bert_layer=}, must be one of {LAYER2CLS.keys()}. "
+                f"{config.padding=} will be automatically prepended to `config.bert_layer` if unspecified."
+            )
+
+
+class FlexBertEncoderBase(nn.Module):
+    """A FlexBERT base class for type hints."""
+
+    layers: nn.ModuleList
+
+    def _init_weights(self, reset_params: bool = False):
+        if hasattr(self, "layers"):
+            for layer in self.layers:
+                layer._init_weights(reset_params=reset_params)
+
+    def reset_parameters(self):
+        self._init_weights(reset_params=True)
+
+    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError("This is a base class and should not be used directly.")
+
+
+class FlexBertUnpadEncoder(FlexBertEncoderBase):
+    """A stack of BERT layers providing the backbone of FlexBERT.
+
+    This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.modeling_bert.BertAlibiEncoder`,
+    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: FlexBertConfig):
+        super().__init__()
+        self.layers = nn.ModuleList([get_bert_layer(config, layer_id=i) for i 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,
+        indices: Optional[torch.Tensor] = None,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+    ) -> torch.Tensor:
+        if indices is None and cu_seqlens is None and max_seqlen is None:
+            attention_mask_bool = attention_mask.bool()
+            batch, seqlen = hidden_states.shape[:2]
+            hidden_states, indices, cu_seqlens, max_seqlen = unpad_input(
+                hidden_states, attention_mask_bool
+            )
+
+            for layer_module in self.layers:
+                hidden_states = layer_module(
+                    hidden_states,
+                    cu_seqlens,
+                    max_seqlen,
+                    indices,
+                    attn_mask=attention_mask,
+                )
+
+            return pad_input(hidden_states, indices, batch, seqlen)
+        else:
+            for layer_module in self.layers:
+                hidden_states = layer_module(
+                    hidden_states,
+                    cu_seqlens,
+                    max_seqlen,
+                    indices,
+                    attn_mask=attention_mask,
+                )
+            return hidden_states
+
+
+class FlexBertPaddedEncoder(FlexBertEncoderBase):
+    """A stack of BERT layers providing the backbone of FlexBERT.
+
+    This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.modeling_bert.BertAlibiEncoder`,
+    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: FlexBertConfig):
+        super().__init__()
+        self.layers = nn.ModuleList([get_bert_layer(config, layer_id=i) for i 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, **kwargs) -> torch.Tensor:
+        for layer_module in self.layers:
+            hidden_states = layer_module(hidden_states, attn_mask=attention_mask)
+
+        return hidden_states
+
+
+ENC2CLS = {
+    "unpadded_base": FlexBertUnpadEncoder,
+    "padded_base": FlexBertPaddedEncoder,
+}
+
+
+def get_encoder_layer(config: FlexBertConfig) -> FlexBertEncoderBase:
+    try:
+        return ENC2CLS[maybe_add_padding(config, config.encoder_layer)](config)
+    except KeyError:
+        raise ValueError(
+            f"Invalid encoder layer type: {config.encoder_layer=}, must be one of {ENC2CLS.keys()}. "
+            f"{config.padding=} will be automatically prepended to `config.encoder_layer` if unspecified."
+        )

mlp.py

@@ -0,0 +1,214 @@
+# Copyright 2024 **AUTHORS_TODO**
+# License: Apache-2.0
+
+# Copyright 2022 MosaicML Examples authors
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright 2023 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) 2023, Tri Dao.
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from .configuration_bert import FlexBertConfig
+from .activation import get_act_fn
+from .normalization import get_norm_layer
+from .initialization import ModuleType, init_weights
+
+
+class BertResidualGLU(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 = get_act_fn(config.hidden_act)
+        self.wo = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.layernorm = get_norm_layer(config)
+
+    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.layernorm(hidden_states + residual_connection)
+        return hidden_states
+
+
+class FlexBertMLPBase(nn.Module):
+    """A FlexBERT MLP base class for type hints."""
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_id = layer_id
+
+    def _init_weights(self, reset_params: bool = False):
+        raise NotImplementedError("This is a base class and should not be used directly.")
+
+    def reset_parameters(self):
+        self._init_weights(reset_params=True)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError("This is a base class and should not be used directly.")
+
+
+class FlexBertMLP(FlexBertMLPBase):
+    """Applies the MLP at the end of each FlexBERT 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.
+    """
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        self.Wi = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_in_bias)
+        self.act = get_act_fn(config.hidden_act)
+        self.drop = nn.Dropout(config.mlp_dropout_prob) if config.mlp_dropout_prob > 0.0 else nn.Identity()
+        self.Wo = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_out_bias)
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(
+            self.config,
+            self.Wi,
+            layer_dim=self.config.hidden_size,
+            layer_id=None,
+            type_of_module=ModuleType.in_module,
+        )
+        init_weights(
+            self.config,
+            self.Wo,
+            layer_dim=self.config.intermediate_size,
+            layer_id=self.layer_id,
+            type_of_module=ModuleType.out_module,
+        )
+
+    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].
+        """
+        return self.Wo(self.drop(self.act(self.Wi(hidden_states))))
+
+
+class FlexBertGLU(FlexBertMLPBase):
+    """Applies the GLU at the end of each FlexBERT 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.
+    """
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        self.Wi = nn.Linear(config.hidden_size, int(config.intermediate_size) * 2, bias=config.mlp_in_bias)
+        self.act = get_act_fn(config.hidden_act)
+        self.drop = nn.Dropout(config.mlp_dropout_prob) if config.mlp_dropout_prob > 0.0 else nn.Identity()
+        self.Wo = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_out_bias)
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(
+            self.config,
+            self.Wi,
+            layer_dim=self.config.hidden_size,
+            layer_id=None,
+            type_of_module=ModuleType.in_module,
+        )
+        init_weights(
+            self.config,
+            self.Wo,
+            layer_dim=self.config.intermediate_size,
+            layer_id=self.layer_id,
+            type_of_module=ModuleType.out_module,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input, gate = self.Wi(hidden_states).chunk(2, dim=-1)
+        return self.Wo(self.drop(self.act(input) * gate))
+
+
+class FlexBertParallelGLU(FlexBertMLPBase):
+    """Applies the GLU at the end of each FlexBERT layer using intermediate_ff computed in parallel of the attention.
+
+    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.
+    """
+
+    def __init__(self, config: FlexBertConfig, layer_id: Optional[int] = None):
+        super().__init__(config=config, layer_id=layer_id)
+        self.act = get_act_fn(config.hidden_act)
+        self.drop = nn.Dropout(config.mlp_dropout_prob) if config.mlp_dropout_prob > 0.0 else nn.Identity()
+        self.Wo = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_out_bias)
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(
+            self.config,
+            self.Wo,
+            layer_dim=self.config.intermediate_size,
+            layer_id=self.layer_id,
+            type_of_module=ModuleType.out_module,
+        )
+
+    def forward(self, intermediate_ff: torch.Tensor) -> torch.Tensor:
+        input, gate = intermediate_ff.chunk(2, dim=-1)
+        return self.Wo(self.drop(self.act(input) * gate))
+
+
+MLP2CLS = {
+    "mlp": FlexBertMLP,
+    "glu": FlexBertGLU,
+    "parallel_glu": FlexBertParallelGLU,
+}
+
+
+def get_mlp_layer(config: FlexBertConfig, layer_id: Optional[int] = None) -> FlexBertMLPBase:
+    try:
+        mlp_layer = (
+            config.initial_mlp_layer
+            if layer_id < config.num_initial_layers and getattr(config, "initial_mlp_layer", None) is not None
+            else config.mlp_layer
+        )
+        return MLP2CLS[mlp_layer](config, layer_id=layer_id)
+    except KeyError as e:
+        if layer_id < config.num_initial_layers and getattr(config, "initial_mlp_layer", None) is not None:
+            raise ValueError(
+                f"Invalid MLP layer type: {config.initial_mlp_layer=}, must be one of {MLP2CLS.keys()}. {e}"
+            )
+        else:
+            raise ValueError(f"Invalid MLP layer type: {config.mlp_layer=}, must be one of {MLP2CLS.keys()}. {e}")

modeling_flexbert.py

@@ -0,0 +1,1920 @@
+# Copyright 2024 **AUTHORS_TODO**
+# License: Apache-2.0
+
+# RMSNorm Implementation: Copyright Meta (from their Llama RMSNorm implementation)
+# License: LLAMA 2 COMMUNITY LICENSE AGREEMENT
+
+# Copyright 2022 Jonas Geiping
+# License: MIT
+
+# Copyright 2022 MosaicML Examples authors
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright 2023 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) 2023, 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 MosaicBERT'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. MosaicBERT 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, MosaicBERT 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 MosaicBERT in Composer, and for example usage
+of the core Mosaic BERT classes.
+"""
+
+import logging
+import os
+import sys
+import warnings
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+# 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__)))
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+from transformers.modeling_outputs import (
+    MaskedLMOutput,
+    ModelOutput,
+    MultipleChoiceModelOutput,
+    SequenceClassifierOutput,
+    CausalLMOutput,
+)
+from transformers.models.bert.modeling_bert import BertPreTrainedModel
+
+from bert_padding import index_put_first_axis
+
+from src.bert_layers.activation import get_act_fn
+from src.bert_layers.attention import (
+    FlexBertPaddedAttention,
+    FlexBertPaddedParallelAttention,
+    FlexBertPaddedRopeAttention,
+    FlexBertPaddedRopeParallelAttention,
+    FlexBertUnpadAttention,
+    FlexBertUnpadParallelAttention,
+    FlexBertUnpadRopeAttention,
+    FlexBertUnpadRopeParallelAttention,
+)
+from src.bert_layers.configuration_bert import FlexBertConfig
+from src.bert_layers.embeddings import (
+    BertAlibiEmbeddings,
+    FlexBertAbsoluteEmbeddings,
+    FlexBertCompiledSansPositionEmbeddings,
+    FlexBertSansPositionEmbeddings,
+    get_embedding_layer,
+)
+from src.bert_layers.initialization import (
+    ModuleType,
+    TileLinear,
+    TileMode,
+    init_weights,
+    tile_embedding,
+    tile_linear,
+    tile_norm,
+)
+from src.bert_layers.layers import (
+    BertAlibiEncoder,
+    BertPooler,
+    BertPredictionHeadTransform,
+    FlexBertCompileUnpadPreNormLayer,
+    FlexBertPaddedEncoder,
+    FlexBertPaddedParallelPreNormLayer,
+    FlexBertPaddedPostNormLayer,
+    FlexBertPaddedPreNormLayer,
+    FlexBertUnpadEncoder,
+    FlexBertUnpadParallelPreNormLayer,
+    FlexBertUnpadPostNormLayer,
+    FlexBertUnpadPreNormLayer,
+    get_encoder_layer,
+)
+from src.bert_layers.loss import get_loss_fn
+from src.bert_layers.mlp import FlexBertGLU, FlexBertMLP, FlexBertParallelGLU
+from src.bert_layers.normalization import get_norm_layer
+from src.bert_layers.padding import pad_input, unpad_input
+
+logger = logging.getLogger(__name__)
+
+
+def _count_parameters(model: nn.Module, trainable: bool = True) -> int:
+    if trainable:
+        return sum(p.numel() for p in model.parameters() if p.requires_grad)
+    else:
+        return sum(p.numel() for p in model.parameters())
+
+
+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: bool = True,
+    ):
+        super(BertModel, self).__init__(config)
+        self.embeddings = BertAlibiEmbeddings(config)
+        self.encoder = BertAlibiEncoder(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.decoder = nn.Linear(bert_model_embedding_weights.size(1), bert_model_embedding_weights.size(0))
+        self.decoder.weight = bert_model_embedding_weights
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+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):
+    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.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.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = 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()
+            masked_token_idx = torch.nonzero(labels.flatten() > 0, as_tuple=False).flatten()
+            loss = loss_fct(prediction_scores, labels.flatten()[masked_token_idx])
+
+            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.
+    """
+
+    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):
+    """
+    Bert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
+    softmax) e.g. for RocStories/SWAG tasks.
+    """
+
+    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)
+
+        # In multiple choice tasks, all choices are submitted in a batch, and
+        # we compute a logit for each option independently. The logits are then
+        # normalized in the forward pass to get a probability distribution over
+        # the choices.
+        self.classifier = nn.Linear(config.hidden_size, 1)
+
+        # 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]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
+            num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
+            `input_ids` above)
+        """
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
+
+        input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
+        attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
+        token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+        position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+        inputs_embeds = (
+            inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
+            if inputs_embeds is not None
+            else None
+        )
+
+        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)
+        reshaped_logits = logits.view(-1, num_choices)
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(reshaped_logits, labels)
+
+        if not return_dict:
+            output = (reshaped_logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MultipleChoiceModelOutput(
+            loss=loss,
+            logits=reshaped_logits,
+            hidden_states=None,
+            attentions=None,
+        )
+
+
+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
+
+
+###################
+# FlexBert Heads
+###################
+
+
+class FlexBertPredictionHead(nn.Module):
+    def __init__(self, config: FlexBertConfig):
+        super().__init__()
+        self.config = config
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size, config.head_pred_bias)
+        self.act = get_act_fn(config.head_pred_act) if config.head_pred_act else nn.Identity()
+        self.norm = (
+            get_norm_layer(config, compiled_norm=config.compile_model) if config.head_pred_norm else nn.Identity()
+        )
+
+    def _init_weights(self, reset_params: bool = False):
+        if reset_params:
+            self.norm.reset_parameters()
+        init_weights(self.config, self.dense, layer_dim=self.config.hidden_size, type_of_module=ModuleType.in_module)
+
+    def reset_parameters(self):
+        self._init_weights(reset_params=True)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.norm(self.act(self.dense(hidden_states)))
+
+
+class FlexBertPoolingHead(nn.Module):
+    def __init__(self, config: FlexBertConfig):
+        super().__init__()
+        self.config = config
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size, config.head_class_bias)
+        self.act = get_act_fn(config.head_class_act) if config.head_class_act else nn.Identity()
+        self.norm = get_norm_layer(config) if config.head_class_norm else nn.Identity()
+        self.drop = torch.nn.Dropout(config.head_class_dropout) if config.head_class_dropout > 0 else nn.Identity()
+        self.pooling_type = config.pooling_type
+
+    def forward(self, hidden_states: torch.Tensor, pool: Optional[bool] = True) -> torch.Tensor:
+        if pool:
+            if self.pooling_type == "cls":
+                output = hidden_states[:, 0]
+            elif self.pooling_type == "mean":
+                output = hidden_states.mean(dim=1)
+            elif self.pooling_type == "max":
+                output = hidden_states.max(dim=1)[0]
+        else:
+            output = hidden_states
+
+        return self.drop(self.norm(self.act(self.dense(output))))
+
+    def _init_weights(self, reset_params: bool = False):
+        init_weights(self.config, self.dense, self.config.hidden_size, type_of_module=ModuleType.out_module)
+        if reset_params and hasattr(self.norm, "reset_parameters"):
+            self.norm.reset_parameters()
+
+    def reset_parameters(self):
+        self._init_weights(reset_params=True)
+
+
+###################
+# FlexBert Models
+###################
+
+
+@dataclass
+class MaskedLMOutput(ModelOutput):
+    """
+    Base class for masked language models outputs.
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Masked language modeling (MLM) loss.
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+    indices: Optional[torch.LongTensor] = None
+    cu_seqlens: Optional[torch.LongTensor] = None
+    max_seqlen: Optional[int] = None
+    batch_size: Optional[int] = None
+    seq_len: Optional[int] = None
+    labels: Optional[torch.LongTensor] = None
+
+
+@dataclass
+class MaskedLMOutputZLoss(ModelOutput):
+    """
+    Base class for masked language models outputs.
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Masked language modeling (MLM) loss.
+        ce_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Cross entropy loss.
+        z_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Z loss.
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+        indices (`torch.LongTensor` of shape `(batch_size,)`):
+            Indices of the tokens to be masked.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    ce_loss: Optional[torch.FloatTensor] = None
+    z_loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+    indices: Optional[torch.LongTensor] = None
+    cu_seqlens: Optional[torch.LongTensor] = None
+    max_seqlen: Optional[int] = None
+    batch_size: Optional[int] = None
+    seq_len: Optional[int] = None
+    labels: Optional[torch.LongTensor] = None
+
+
+class FlexBertPreTrainedModel(BertPreTrainedModel):
+    """
+    An abstract class to handle custom weights initialization of modules
+    """
+
+    def _init_module_weights(self, module: nn.Module):
+        """
+        Custom weight init of modules using src.bert_layers.initialization.init_weights
+        Currently only supports init of embedding modules
+        """
+        assert isinstance(module, nn.Module)
+        if isinstance(module, nn.Embedding):
+            init_weights(self.config, module, type_of_module=ModuleType.emb)
+        else:
+            raise NotImplementedError("Custom weight init for the given module is not supported")
+
+
+class FlexBertModel(FlexBertPreTrainedModel):
+    """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: FlexBertConfig):
+        super().__init__(config)
+        self.embeddings = get_embedding_layer(config)
+        self.encoder = get_encoder_layer(config)
+        if config.final_norm:
+            # if we use prenorm attention we need to add a final norm
+            self.final_norm = get_norm_layer(config)
+        else:
+            self.final_norm = None
+        self.unpad_embeddings = config.unpad_embeddings
+
+    def post_init(self):
+        self._init_weights(reset_params=False)
+        self._backward_compatibility_gradient_checkpointing()
+
+    def get_input_embeddings(self):
+        return self.embeddings.tok_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.tok_embeddings = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        indices: Optional[torch.Tensor] = None,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+        **kwargs,
+    ) -> Tuple[Union[List[torch.Tensor], torch.Tensor], Optional[torch.Tensor]]:
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        embedding_output = self.embeddings(input_ids, position_ids)
+
+        encoder_outputs = self.encoder(
+            hidden_states=embedding_output,
+            attention_mask=attention_mask,
+            indices=indices,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+
+        if self.final_norm is not None:
+            encoder_outputs = self.final_norm(encoder_outputs)
+        return encoder_outputs
+
+    def _init_weights(self, module: Optional[nn.Module] = None, reset_params: Optional[bool] = None):
+        assert (module is None) != (reset_params is None), "arg module xor reset_params must be specified"
+        if module:
+            self._init_module_weights(module)
+        else:
+            assert isinstance(reset_params, bool)
+            self.embeddings._init_weights(reset_params=reset_params)
+            self.encoder._init_weights(reset_params=reset_params)
+
+            if reset_params and self.config.final_norm:
+                self.final_norm.reset_parameters()
+
+    def reset_parameters(self):
+        self._init_weights(reset_params=True)
+
+    def get_number_parameters(self, count_embeddings: bool = True, trainable: bool = True) -> int:
+        """Returns the number of parameters in the model.
+
+        Args:
+            count_embeddings: count the parameters in the embeddings layer, excluding position embeddings.
+            trainable: only count trainable parameters.
+        """
+        params = sum([_count_parameters(layer, trainable) for layer in self.encoder.layers])
+        if count_embeddings:
+            params += _count_parameters(self.embeddings, trainable)
+            if hasattr(self.embeddings, "position_embeddings"):
+                params -= _count_parameters(self.embeddings.position_embeddings, trainable)
+        return params
+
+
+class FlexBertForMaskedLM(FlexBertPreTrainedModel):
+    def __init__(self, config: FlexBertConfig):
+        super().__init__(config)
+        self.bert = FlexBertModel(config)
+        self.head = FlexBertPredictionHead(config)
+
+        if config.tie_word_embeddings:
+            decoder_weights = self.bert.embeddings.tok_embeddings.weight
+        else:
+            decoder_weights = nn.Linear(config.hidden_size, config.vocab_size, bias=False).weight
+        self.decoder = nn.Linear(decoder_weights.size(1), decoder_weights.size(0), bias=config.decoder_bias)
+        self.decoder.weight = decoder_weights
+
+        self.loss_fn = nn.CrossEntropyLoss() if not hasattr(config, "loss_function") else get_loss_fn(config)
+        self.fa_ce = getattr(config, "loss_function", "cross_entropy") == "fa_cross_entropy"
+        self.return_z_loss = config.loss_kwargs.get("return_z_loss", False)
+        self.unpad_embeddings = config.unpad_embeddings
+        self.pad_logits = config.pad_logits
+        self.compile_model = config.compile_model
+        self.masked_prediction = config.masked_prediction
+
+        # Initialize weights and apply final processing
+        self._init_weights(reset_params=False)
+
+    def _init_weights(self, module: Optional[nn.Module] = None, reset_params: Optional[bool] = None):
+        assert (module is None) != (reset_params is None), "arg module xor reset_params must be specified"
+        if module:
+            self._init_module_weights(module)
+        else:
+            assert isinstance(reset_params, bool)
+            self.bert._init_weights(reset_params=reset_params)
+            self.head._init_weights(reset_params=reset_params)
+
+            # Output weights.
+            if not self.config.tie_word_embeddings:
+                init_weights(self.config, self.decoder, self.config.hidden_size, type_of_module=ModuleType.final_out)
+
+    @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("FlexBERT 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.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.decoder = new_embeddings
+
+    @torch.no_grad()
+    def unpad_inputs(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, position_ids: torch.Tensor, labels: torch.Tensor
+    ):
+        return unpad_input(input_ids, attention_mask, position_ids, labels)
+
+    @torch.no_grad()
+    def pad_inputs(
+        self,
+        inputs: torch.Tensor,
+        indices: torch.Tensor,
+        batch_size: int,
+        seqlen: int,
+        labels: Optional[torch.Tensor] = None,
+        ignore_index: int = -100,
+    ):
+        return pad_input(
+            inputs=inputs, indices=indices, batch=batch_size, seqlen=seqlen, labels=labels, ignore_index=ignore_index
+        )
+
+    @torch.compile(dynamic=True)
+    def compiled_head(self, output: torch.Tensor) -> torch.Tensor:
+        return self.decoder(self.head(output))
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        indices: Optional[torch.Tensor] = None,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+        batch_size: Optional[int] = None,
+        seq_len: Optional[int] = None,
+        **kwargs,
+    ) -> 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
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if self.unpad_embeddings and (indices is None and cu_seqlens is None and max_seqlen is None):
+            batch_size, seq_len = input_ids.shape[:2]
+            input_ids, indices, cu_seqlens, max_seqlen, position_ids, labels = self.unpad_inputs(
+                input_ids, attention_mask, position_ids, labels
+            )
+        
+
+        output = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            indices=indices,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+
+        if self.masked_prediction and labels is not None:
+            # flatten labels and output first
+            labels = labels.view(-1)
+            output = output.view(labels.shape[0], -1)
+
+            # then filter out the non-masked tokens
+            mask_tokens = labels != self.loss_fn.ignore_index
+            output = output[mask_tokens]
+            labels = labels[mask_tokens]
+
+        if self.compile_model:
+            logits = self.compiled_head(output)
+        else:
+            logits = self.decoder(self.head(output))
+
+        loss = None
+        if labels is not None:
+            if not self.masked_prediction:
+                labels = labels.view(-1)
+                logits = logits.view(labels.shape[0], -1)
+
+            if self.return_z_loss:
+                loss, z_loss = self.loss_fn(logits, labels)
+                if self.pad_logits:
+                    return MaskedLMOutputZLoss(
+                        loss=loss,
+                        ce_loss=loss.detach().clone() - z_loss,
+                        z_loss=z_loss,
+                        logits=self.pad_inputs(logits, indices, batch_size, seq_len)[0],
+                        hidden_states=None,
+                        attentions=None,
+                    )
+                else:
+                    return MaskedLMOutputZLoss(
+                        loss=loss,
+                        ce_loss=loss.detach().clone() - z_loss,
+                        z_loss=z_loss,
+                        logits=logits,
+                        hidden_states=None,
+                        attentions=None,
+                        indices=indices,
+                        cu_seqlens=cu_seqlens,
+                        max_seqlen=max_seqlen,
+                        batch_size=batch_size,
+                        seq_len=seq_len,
+                        labels=labels,
+                    )
+            else:
+                loss = self.loss_fn(logits, labels)
+
+        if self.pad_logits:
+            return MaskedLMOutput(
+                loss=loss,
+                logits=self.pad_inputs(logits, indices, batch_size, seq_len)[0],
+                hidden_states=None,
+                attentions=None,
+            )
+        else:
+            return MaskedLMOutput(
+                loss=loss,
+                logits=logits,
+                hidden_states=None,
+                attentions=None,
+                indices=indices,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=max_seqlen,
+                batch_size=batch_size,
+                seq_len=seq_len,
+                labels=labels,
+            )
+
+    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}
+
+    def get_number_parameters(
+        self, count_embeddings: bool = True, count_decoder: bool = False, trainable: bool = True
+    ) -> int:
+        """Returns the number of parameters in the model.
+
+        Args:
+            count_embeddings: count the parameters in the embeddings layer, excluding position embeddings.
+            count_decoder: count the parameters in the decoder layer if weights are not tied.
+            trainable: only count trainable parameters.
+        """
+        params = self.bert.get_number_parameters(count_embeddings, trainable)
+        params += _count_parameters(self.head, trainable)
+        if count_decoder and not self.config.tie_word_embeddings:
+            params += _count_parameters(self.decoder, trainable)
+        return params
+
+
+class FlexBertForSequenceClassification(FlexBertPreTrainedModel):
+    """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.
+    """
+
+    def __init__(self, config: FlexBertConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+
+        self.bert = FlexBertModel(config)
+        self.head = FlexBertPoolingHead(config)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self._init_weights(reset_params=False)
+
+    def _init_weights(self, module: Optional[nn.Module] = None, reset_params: Optional[bool] = None):
+        assert (module is None) != (reset_params is None), "arg module xor reset_params must be specified"
+        if module:
+            self._init_module_weights(module)
+        else:
+            assert isinstance(reset_params, bool)
+            self.bert._init_weights(reset_params=reset_params)
+            self.head._init_weights(reset_params=reset_params)
+            init_weights(self.config, self.classifier, self.config.hidden_size, type_of_module=ModuleType.final_out)
+
+    @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,
+        position_ids: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = 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
+
+        output = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+        )
+
+        pooled_output = self.head(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,) + output
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def get_number_parameters(self, count_embeddings: bool = True, trainable: bool = True) -> int:
+        """Returns the number of parameters in the model.
+
+        Args:
+            count_embeddings: count the parameters in the embeddings layer, excluding position embeddings.
+            trainable: only count trainable parameters.
+        """
+        params = self.bert.get_number_parameters(count_embeddings, trainable)
+        params += _count_parameters(self.head, trainable)
+        params += _count_parameters(self.classifier, trainable)
+        return params
+
+
+class FlexBertForMultipleChoice(FlexBertPreTrainedModel):
+    """
+    Bert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
+    softmax) e.g. for RocStories/SWAG tasks.
+    """
+
+    def __init__(self, config: FlexBertConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+
+        self.bert = FlexBertModel(config)
+        self.head = FlexBertPoolingHead(config)
+
+        # In multiple choice tasks, all choices are submitted in a batch, and
+        # we compute a logit for each option independently. The logits are then
+        # normalized in the forward pass to get a probability distribution over
+        # the choices.
+        self.classifier = nn.Linear(config.hidden_size, 1)
+
+        # Initialize weights and apply final processing
+        self._init_weights(reset_params=False)
+
+    def _init_weights(self, module: Optional[nn.Module] = None, reset_params: Optional[bool] = None):
+        assert (module is None) != (reset_params is None), "arg module xor reset_params must be specified"
+        if module:
+            self._init_module_weights(module)
+        else:
+            assert isinstance(reset_params, bool)
+            self.bert._init_weights(reset_params=reset_params)
+            self.head._init_weights(reset_params=reset_params)
+            init_weights(self.config, self.classifier, self.config.hidden_size, type_of_module=ModuleType.final_out)
+
+    @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,
+        position_ids: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = 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
+        num_choices = input_ids.shape[1]
+
+        input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
+        attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
+        position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+
+        output = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+        )
+
+        pooled_output = self.head(output)
+        logits = self.classifier(pooled_output)
+        reshaped_logits = logits.view(-1, num_choices)
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(reshaped_logits, labels)
+
+        if not return_dict:
+            output = (reshaped_logits,) + output
+            return ((loss,) + output) if loss is not None else output
+
+        return MultipleChoiceModelOutput(
+            loss=loss,
+            logits=reshaped_logits,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def get_number_parameters(self, count_embeddings: bool = True, trainable: bool = True) -> int:
+        """Returns the number of parameters in the model.
+
+        Args:
+            count_embeddings: count the parameters in the embeddings layer, excluding position embeddings.
+            trainable: only count trainable parameters.
+        """
+        params = self.bert.get_number_parameters(count_embeddings, trainable)
+        params += _count_parameters(self.head, trainable)
+        params += _count_parameters(self.classifier, trainable)
+        return params
+
+
+class FlexBertForCasualLM(FlexBertPreTrainedModel):
+    """Bert Model transformer with a LM head.
+
+    This head is just a standard LM head module. Used for causal language modeling tasks.
+    """
+
+    def __init__(self, config: FlexBertConfig):
+        super().__init__(config)
+        self.bert = FlexBertModel(config)
+        self.lm_head = FlexBertPredictionHead(config)
+
+        if config.tie_word_embeddings:
+            decoder_weights = self.bert.embeddings.tok_embeddings.weight
+        else:
+            decoder_weights = nn.Linear(config.hidden_size, config.vocab_size, bias=False).weight
+        self.decoder = nn.Linear(decoder_weights.size(1), decoder_weights.size(0), bias=config.decoder_bias)
+        self.decoder.weight = decoder_weights
+
+        self.loss_fn = nn.CrossEntropyLoss() if not hasattr(config, "loss_function") else get_loss_fn(config)
+        self.fa_ce = getattr(config, "loss_function", "cross_entropy") == "fa_cross_entropy"
+        self.return_z_loss = config.loss_kwargs.get("return_z_loss", False)
+        self.unpad_embeddings = config.unpad_embeddings
+        self.pad_logits = config.pad_logits
+        self.compile_model = config.compile_model
+        # self.masked_prediction = config.masked_prediction
+
+        # Initialize weights and apply final processing
+        self._init_weights(reset_params=False)
+
+    def _init_weights(self, module: Optional[nn.Module] = None, reset_params: Optional[bool] = None):
+        assert (module is None) != (reset_params is None), "arg module xor reset_params must be specified"
+        if module:
+            self._init_module_weights(module)
+        else:
+            assert isinstance(reset_params, bool)
+            self.bert._init_weights(reset_params=reset_params)
+            self.lm_head._init_weights(reset_params=reset_params)
+
+            if not self.config.tie_word_embeddings:
+                init_weights(self.config, self.decoder, self.config.hidden_size, type_of_module=ModuleType.final_out)
+
+    @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.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.decoder = new_embeddings
+
+    @torch.no_grad()
+    def unpad_inputs(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, position_ids: torch.Tensor, labels: torch.Tensor
+    ):
+        return unpad_input(input_ids, attention_mask, position_ids, labels)
+
+    @torch.no_grad()
+    def pad_inputs(
+        self,
+        inputs: torch.Tensor,
+        indices: torch.Tensor,
+        batch_size: int,
+        seqlen: int,
+        labels: Optional[torch.Tensor] = None,
+        ignore_index: int = -100,
+    ):
+        return pad_input(
+            inputs=inputs, indices=indices, batch=batch_size, seqlen=seqlen, labels=labels, ignore_index=ignore_index
+        )
+
+    @torch.compile(dynamic=True)
+    def compiled_lm_head(self, output: torch.Tensor) -> torch.Tensor:
+        return self.decoder(self.lm_head(output))
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        indices: Optional[torch.Tensor] = None,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+        batch_size: Optional[int] = None,
+        seq_len: Optional[int] = None,
+        **kwargs,
+    ) -> Union[Tuple[torch.Tensor], CausalLMOutput]:
+        # 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
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if self.unpad_embeddings and (indices is None and cu_seqlens is None and max_seqlen is None):
+            batch_size, seq_len = input_ids.shape[:2]
+            input_ids, indices, cu_seqlens, max_seqlen, position_ids, labels = self.unpad_inputs(
+                input_ids, attention_mask, position_ids, labels
+            )
+
+        hidden_states = self.bert(
+            input_ids,
+            attention_mask=None,
+            position_ids=position_ids,
+            indices=indices,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+
+        if self.compile_model:
+            logits = self.compiled_lm_head(hidden_states)
+        else:
+            logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            if indices is not None:                
+                # Unpadded case: shift within each sequence using input_ids
+                # Initialize shifted labels from input_ids
+                shift_labels = torch.full_like(input_ids, -100)
+                
+                # For each sequence, shift the input_ids to create labels
+                for i in range(len(cu_seqlens) - 1):
+                    start = cu_seqlens[i]
+                    end = cu_seqlens[i + 1]
+                    # Input: [A, B, C, D] -> Labels: [B, C, D, -100]
+                    shift_labels[start:end-1] = input_ids[start+1:end]
+                        
+                # Debug prints
+                # print(f"input_ids slice: {input_ids[:20]}")  # Show first 20 tokens
+                # print(f"shift_labels slice: {shift_labels[:20]}")  # Show first 20 token
+                        
+                # # Debug prints
+                # print(f"input_ids slice: {input_ids[:20]}")  # Show first 20 tokens
+                # print(f"shift_labels slice: {shift_labels[:20]}")  # Show first 20 tokens
+                # print(f"First sequence length: {cu_seqlens[1] - cu_seqlens[0]}")
+                    
+            else:
+                # Padded case: simple shift
+                shift_labels = input_ids[..., 1:].contiguous()
+                logits = logits[..., :-1, :].contiguous()
+
+            # For both cases, we'll use the shifted input_ids as our labels
+            labels = shift_labels
+            
+            # Flatten the tokens
+            loss = self.loss_fn(
+                logits.view(-1, logits.size(-1)),
+                shift_labels.view(-1)
+            )
+
+        if self.pad_logits:
+            return CausalLMOutput(
+                loss=loss,
+                logits=self.pad_inputs(logits, indices, batch_size, seq_len)[0],
+                hidden_states=None,
+                attentions=None,
+            )
+        else:
+            return CausalLMOutput(
+                loss=loss,
+                logits=logits,
+                hidden_states=hidden_states,
+                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}
+
+    def get_number_parameters(self, count_embeddings: bool = True, trainable: bool = True) -> int:
+        """Returns the number of parameters in the model.
+
+        Args:
+            count_embeddings: count the parameters in the embeddings layer, excluding position embeddings.
+            trainable: only count trainable parameters.
+        """
+        params = self.bert.get_number_parameters(count_embeddings, trainable)
+        params += _count_parameters(self.lm_head, trainable)
+        return params
+
+
+def init_model_from_pretrained(
+    pretrained_model: FlexBertModel,
+    new_model: FlexBertModel,
+    mode: Union[str, TileMode] = TileMode.tile_weights_from_middle,
+):
+    """
+    Initialize the new model from the pretrained model.
+
+    This method uses Gopher layer scaling and Phi-style weight tiling as selected by `mode`.
+    The new model must have the same or more layers and the same or larger dimensions than the pretrained model.
+
+    Args:
+        pretrained_model (FlexBertModel): The smaller, pre-trained model
+        new_model (FlexBertModel): The larger model to be initialized
+        mode (Union[str, TileMode]): The Phi-style weight tiling mode to use
+
+    This function assumes that the new_model has more layers and a larger hidden size
+    than the pretrained_model, but the same vocabulary size.
+    """
+
+    # Tile embeddings
+    assert isinstance(
+        new_model.embeddings, type(pretrained_model.embeddings)
+    ), f"Pretrained and new_model layers must be the same type, got {type(new_model.embeddings)} and {type(pretrained_model.embeddings)}"
+    assert isinstance(
+        new_model.embeddings,
+        (FlexBertAbsoluteEmbeddings, FlexBertSansPositionEmbeddings, FlexBertCompiledSansPositionEmbeddings),
+    ), f"Unsupported embedding layer type: {type(new_model.embeddings)}"
+
+    tile_embedding(pretrained_model.embeddings.tok_embeddings, new_model.embeddings.tok_embeddings, mode=mode)
+    if isinstance(pretrained_model.embeddings, FlexBertAbsoluteEmbeddings):
+        tile_embedding(pretrained_model.embeddings.pos_embeddings, new_model.embeddings.pos_embeddings, mode=mode)
+
+    if hasattr(pretrained_model.embeddings, "norm"):
+        tile_norm(pretrained_model.embeddings.norm, new_model.embeddings.norm, mode=mode)
+
+    # Tile encoder layers
+    assert isinstance(
+        pretrained_model.encoder, (FlexBertUnpadEncoder, FlexBertPaddedEncoder)
+    ), f"Unsupported encoder layer type: {type(pretrained_model.encoder)}"
+    assert isinstance(
+        new_model.encoder, type(pretrained_model.encoder)
+    ), f"Pretrained and new_model encoder layers must be the same type, got {type(new_model.encoder)} and {type(pretrained_model.encoder)}"
+
+    # Calculate the layer mapping
+    pretrained_layers = len(pretrained_model.encoder.layers)
+    new_layers = len(new_model.encoder.layers)
+    layer_mapping = [round(i * pretrained_layers / new_layers) for i in range(new_layers)]
+
+    # Initialize layers
+    for new_model_idx, pretrained_idx in enumerate(layer_mapping):
+        new_model_layer = new_model.encoder.layers[new_model_idx]
+        pretrained_layer = pretrained_model.encoder.layers[pretrained_idx]
+
+        # first tile the PreNorm/PostNorm layers
+        assert isinstance(
+            new_model_layer, type(pretrained_layer)
+        ), f"Pretrained and new_model prenorm/postnorm layers must be the same type, got {type(new_model_layer)} and {type(pretrained_layer)}"
+        assert isinstance(
+            new_model_layer,
+            (
+                FlexBertUnpadPreNormLayer,
+                FlexBertCompileUnpadPreNormLayer,
+                FlexBertUnpadParallelPreNormLayer,
+                FlexBertUnpadPostNormLayer,
+                FlexBertPaddedPreNormLayer,
+                FlexBertPaddedParallelPreNormLayer,
+                FlexBertPaddedPostNormLayer,
+            ),
+        ), f"Unsupported prenorm/postnorm layer type: {type(new_model_layer)}"
+
+        # First tile the normalization layers
+        if hasattr(pretrained_layer, "attn_norm"):
+            tile_norm(pretrained_layer.attn_norm, new_model_layer.attn_norm, mode=mode)
+        if hasattr(pretrained_layer, "norm"):
+            tile_norm(pretrained_layer.norm, new_model_layer.norm, mode=mode)
+        if hasattr(pretrained_layer, "mlp_norm"):
+            tile_norm(pretrained_layer.mlp_norm, new_model_layer.mlp_norm, mode=mode)
+
+        # Then tile the attention & mlp layers
+        assert isinstance(
+            new_model_layer.attn, type(pretrained_layer.attn)
+        ), f"Pretrained and new_model attention layers must be the same type, got {type(new_model_layer.attn)} and {type(pretrained_layer.attn)}"
+
+        # first try the parallel attention layers
+        if isinstance(pretrained_layer, (FlexBertUnpadParallelPreNormLayer, FlexBertPaddedParallelPreNormLayer)):
+            assert isinstance(
+                pretrained_layer.attn,
+                (
+                    FlexBertUnpadParallelAttention,
+                    FlexBertPaddedParallelAttention,
+                    FlexBertUnpadRopeParallelAttention,
+                    FlexBertPaddedRopeParallelAttention,
+                ),
+            ), f"Parallel prenorm layer must have parallel attention layer: {type(pretrained_layer.attn)}"
+            if not isinstance(pretrained_layer.mlp, (FlexBertParallelGLU)):
+                raise ValueError(f"Parallel prenorm layer must have parallel MLP layer: {type(pretrained_layer.mlp)}")
+            tile_linear(
+                pretrained_layer.Wqkvff,
+                new_model_layer.Wqkvff,
+                linear_type=TileLinear.wqkvff,
+                mode=mode,
+                pretrained_attn_size=pretrained_layer.attn_size,
+                pretrained_mlp_size=pretrained_layer.mlp_size,
+                new_attn_size=new_model_layer.attn_size,
+                new_mlp_size=new_model_layer.mlp_size,
+                wqkvff_is_glu=True,
+            )
+
+        # then try the fused attention layers
+        elif isinstance(
+            pretrained_layer.attn,
+            (
+                FlexBertUnpadAttention,
+                FlexBertPaddedAttention,
+                FlexBertUnpadRopeAttention,
+                FlexBertPaddedRopeAttention,
+            ),
+        ):
+            tile_linear(pretrained_layer.attn.Wqkv, new_model_layer.attn.Wqkv, linear_type=TileLinear.wqkv, mode=mode)
+        else:
+            raise ValueError(f"Unsupported attention layer type: {type(pretrained_layer.attn)}")
+
+        # finally, tile the attention output layer
+        tile_linear(pretrained_layer.attn.Wo, new_model_layer.attn.Wo, linear_type=TileLinear.default, mode=mode)
+
+        # tile the mlp layer if the model is not using parallel attention layers
+        if not isinstance(pretrained_layer.mlp, (FlexBertMLP, FlexBertGLU, FlexBertParallelGLU)):
+            raise ValueError(f"Unsupported MLP layer type: {type(pretrained_layer.mlp)}")
+        assert isinstance(
+            new_model_layer.mlp, type(pretrained_layer.mlp)
+        ), f"Pretrained and new_model mlp layers must be the same type, got {type(new_model_layer.mlp)} and {type(pretrained_layer.mlp)}"
+
+        # already tiled the parallel glu layer if it exists, so only need to handle mlp & glu Wi
+        if isinstance(pretrained_layer.mlp, FlexBertGLU):
+            tile_linear(pretrained_layer.mlp.Wi, new_model_layer.mlp.Wi, linear_type=TileLinear.glu, mode=mode)
+        elif isinstance(pretrained_layer.mlp, FlexBertMLP):
+            tile_linear(pretrained_layer.mlp.Wi, new_model_layer.mlp.Wi, linear_type=TileLinear.default, mode=mode)
+        # tile the output for both ParallelGLU and MLP/GLU
+        tile_linear(pretrained_layer.mlp.Wo, new_model_layer.mlp.Wo, linear_type=TileLinear.default, mode=mode)
+
+
+def init_mlm_model_from_pretrained(
+    config: FlexBertConfig,
+    pretrained_model: FlexBertForMaskedLM,
+    new_model: FlexBertForMaskedLM,
+    mode: Union[str, TileMode] = TileMode.tile_weights_from_middle,
+):
+    """
+    Initialize the new model from the pretrained model.
+
+    This method uses Gopher layer scaling and Phi-style weight tiling as selected by `mode`.
+    The new model must have the same or more layers and the same or larger dimensions than the pretrained model.
+
+    Args:
+        config (FlexBertConfig): The configuration of the new_model
+        pretrained_model (FlexBertForMaskedLM): The smaller, pre-trained model
+        new_model (FlexBertForMaskedLM): The larger model to be initialized from the pretrained model
+        mode (Union[str, TileMode]): The Phi-style weight tiling mode to use
+
+    This function assumes that the new_model has more layers and a larger hidden size
+    than the pretrained_model, but the same vocabulary size.
+    """
+    init_model_from_pretrained(pretrained_model.bert, new_model.bert, mode=mode)
+
+    # TODO: uncomment this when the repo is turned into a pip installable package
+    # if not isinstance(pretrained_model.head, FlexBertPredictionHead):
+    #     raise ValueError(f"Pretrained model must have a prediction head: {type(pretrained_model.head)}")
+    # if not isinstance(new_model.head, FlexBertPredictionHead):
+    #     raise ValueError(f"New model must have a prediction head: {type(new_model.head)}")
+
+    # tile the prediction head
+    tile_linear(pretrained_model.head.dense, new_model.head.dense, linear_type=TileLinear.default, mode=mode)
+    tile_norm(pretrained_model.head.norm, new_model.head.norm, mode=mode)
+
+    # setup weight tying
+    if config.tie_word_embeddings:
+        new_model.decoder.weight = new_model.bert.embeddings.tok_embeddings.weight
+        tile_linear(
+            pretrained_model.decoder, new_model.decoder, linear_type=TileLinear.default, mode=mode, bias_only=True
+        )
+    else:
+        tile_linear(pretrained_model.decoder, new_model.decoder, linear_type=TileLinear.default, mode=mode)

normalization.py

@@ -0,0 +1,116 @@
+# Copyright 2024 **AUTHORS_TODO**
+# License: Apache-2.0
+
+# RMSNorm Implementation: Copyright Meta (from their Llama RMSNorm implementation)
+# License: LLAMA 2 COMMUNITY LICENSE AGREEMENT
+
+
+import inspect
+import torch
+import torch.nn as nn
+from torch.nn import init
+
+from .configuration_bert import FlexBertConfig
+
+try:
+    from flash_attn.ops.triton.layer_norm import RMSNorm as TritonRMSNorm
+    from flash_attn.ops.triton.layer_norm import layer_norm_fn
+
+except ImportError:
+    TritonRMSNorm = None
+    layer_norm_fn = None
+
+
+class RMSNorm(nn.Module):
+    """Llama2 RMSNorm implementation"""
+
+    def __init__(self, dim: int, eps: float = 1e-5):
+        """
+        Initialize the RMSNorm normalization layer.
+
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+        Attributes:
+            eps (float): A small value added to the denominator for numerical stability.
+            weight (nn.Parameter): Learnable scaling parameter.
+
+        """
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        """
+        Apply the RMSNorm normalization to the input tensor.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The normalized tensor.
+
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        """
+        Forward pass through the RMSNorm layer.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor after applying RMSNorm.
+
+        """
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+    def reset_parameters(self):
+        init.ones_(self.weight)
+
+
+if layer_norm_fn is not None:
+
+    class TritonLayerNorm(nn.LayerNorm):
+        def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False):
+            return layer_norm_fn(
+                x,
+                self.weight,
+                self.bias,
+                residual=residual,
+                eps=self.eps,
+                prenorm=prenorm,
+                residual_in_fp32=residual_in_fp32,
+            )
+else:
+    TritonLayerNorm = None
+
+NORM2CLS = {
+    "layernorm": nn.LayerNorm,
+    "triton_layernorm": TritonLayerNorm if TritonLayerNorm is not None else nn.LayerNorm,
+    "rmsnorm": RMSNorm,
+    "triton_rmsnorm": TritonRMSNorm if TritonRMSNorm is not None else RMSNorm,
+}
+
+
+def get_norm_layer(config: FlexBertConfig, compiled_norm: bool = False) -> nn.Module:
+    try:
+        if compiled_norm:
+            # Use non-Triton norms when compiling
+            if config.normalization.startswith("triton_"):
+                norm = config.normalization.replace("triton_", "")
+            else:
+                norm = config.normalization
+        else:
+            norm = config.normalization
+        signature = inspect.signature(NORM2CLS[norm])
+        if hasattr(config, "norm_kwargs"):
+            norm_kwargs = {k: v for k, v in config.norm_kwargs.items() if k in signature.parameters}
+        else:
+            norm_kwargs = {}
+        return NORM2CLS[norm](config.hidden_size, **norm_kwargs)
+    except KeyError:
+        raise ValueError(f"Invalid normalization layer type: {config.normalization}, must be one of {NORM2CLS.keys()}.")

options.py

@@ -0,0 +1,32 @@
+from .normalization import NORM2CLS
+from .embeddings import EBB2CLS
+from .activation import ACT2CLS
+from .attention import ATTN2CLS
+from .mlp import MLP2CLS
+from .layers import LAYER2CLS
+
+
+def print_layer_options():
+    print("Activation options:")
+    for option in ACT2CLS:
+        print(f"  {option}")
+
+    print("\nAttention Layer options:")
+    for option in ATTN2CLS:
+        print(f"  {option}")
+
+    print("\nEmbedding Layer options:")
+    for option in EBB2CLS:
+        print(f"  {option}")
+
+    print("\nBert Layer options:")
+    for option in LAYER2CLS:
+        print(f"  {option}")
+
+    print("\nMLP Layer options:")
+    for option in MLP2CLS:
+        print(f"  {option}")
+
+    print("\nNormalization options:")
+    for option in NORM2CLS:
+        print(f"  {option}")

padding.py

@@ -0,0 +1,87 @@
+import torch
+from torch import Tensor
+from typing import Optional, Tuple
+import torch.nn.functional as F
+
+
+def unpad_input(
+    inputs: Tensor,
+    attention_mask: Tensor,
+    position_ids: Optional[Tensor] = None,
+    labels: Optional[Tensor] = None,
+) -> Tuple[Tensor, Tensor, Tensor, int, Optional[Tensor], Optional[Tensor]]:
+    """
+    Remove padding from input sequences.
+
+    Args:
+        inputs: (batch, seqlen, ...) or (batch, seqlen)
+        attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
+        position_ids: (batch, seqlen), int, position ids
+        labels: (batch, seqlen), int, labels
+
+    Returns:
+        unpadded_inputs: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask.
+        indices: (total_nnz)
+        cu_seqlens: (batch + 1), the cumulative sequence lengths
+        max_seqlen_in_batch: int
+        unpadded_position_ids: (total_nnz) or None
+        unpadded_labels: (total_nnz) or None
+    """
+    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))
+
+    if inputs.dim() == 2:
+        unpadded_inputs = inputs.flatten()[indices]
+    else:
+        batch, seqlen, *rest = inputs.shape
+        shape = batch * seqlen
+        unpadded_inputs = inputs.view(shape, *rest)[indices]
+
+    unpadded_position_ids = position_ids.flatten()[indices] if position_ids is not None else None
+    unpadded_labels = labels.flatten()[indices] if labels is not None else None
+
+    return unpadded_inputs, indices, cu_seqlens, max_seqlen_in_batch, unpadded_position_ids, unpadded_labels
+
+
+def pad_input(
+    inputs: Tensor,
+    indices: Tensor,
+    batch: int,
+    seqlen: int,
+    labels: Optional[Tensor] = None,
+    ignore_index: int = -100,
+) -> Tuple[Tensor, Optional[Tensor]]:
+    """
+    Add padding to sequences.
+
+    Args:
+        inputs: (total_nnz, ...) or (total_nnz,), where total_nnz = number of tokens selected in attention_mask.
+        indices: (total_nnz)
+        batch: int, batch size
+        seqlen: int, max sequence length
+        position_ids: (total_nnz) or None
+        labels: (total_nnz) or None
+
+    Returns:
+        padded_inputs: (batch, seqlen, ...) or (batch, seqlen)
+        padded_labels: (batch, seqlen) or None
+    """
+    if inputs.dim() == 1:
+        output = torch.zeros(batch * seqlen, dtype=inputs.dtype, device=inputs.device)
+        output[indices] = inputs
+        padded_inputs = output.view(batch, seqlen)
+    else:
+        _, *rest = inputs.shape
+        output = torch.zeros(batch * seqlen, *rest, dtype=inputs.dtype, device=inputs.device)
+        output[indices] = inputs
+        padded_inputs = output.view(batch, seqlen, *rest)
+
+    padded_labels = None
+    if labels is not None:
+        padded_labels = torch.full((batch * seqlen,), fill_value=ignore_index, dtype=labels.dtype, device=labels.device)
+        padded_labels[indices] = labels
+        padded_labels = padded_labels.view(batch, seqlen)
+
+    return padded_inputs, padded_labels

rotary.py

@@ -0,0 +1,297 @@
+# Copyright 2024 **AUTHORS_TODO**
+# License: Apache-2.0
+
+# Copyright (c) 2023, Tri Dao.
+# License: Apache-2.0
+
+import torch
+from einops import rearrange
+from flash_attn.ops.triton.rotary import apply_rotary
+
+from typing import Optional, Tuple, Union
+
+
+class ApplyRotaryEmbUnpad(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        qkv,
+        cos,
+        sin,
+        interleaved=False,
+        seqlen_offsets: Union[int, torch.Tensor] = 0,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+    ):
+        # (total_nnz, 3, nheads, headdim)
+        total_nnz, three, nheads, headdim = qkv.shape
+        assert three == 3
+        if qkv.is_contiguous():
+            # Call 1 kernel instead of 2 kernels
+            # We need qkv to be contiguous so that when we reshape to combine (3, nheads)
+            # dimensions, we get the same tensor
+            # qk = rearrange(qkv[:, :2], "b_s t h d -> b_s (t h) d")
+            qk = qkv[:, :2].view(total_nnz, -1, headdim)
+            apply_rotary(
+                qk,
+                cos,
+                sin,
+                seqlen_offsets=seqlen_offsets,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=max_seqlen,
+                interleaved=interleaved,
+                inplace=True,
+            )
+        else:
+            q, k = qkv[:, 0, :, :], qkv[:, 1, :, :]
+            apply_rotary(
+                q,
+                cos,
+                sin,
+                seqlen_offsets=seqlen_offsets,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=max_seqlen,
+                interleaved=interleaved,
+                inplace=True,
+            )
+            apply_rotary(
+                k,
+                cos,
+                sin,
+                seqlen_offsets=seqlen_offsets,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=max_seqlen,
+                interleaved=interleaved,
+                inplace=True,
+            )
+
+        if isinstance(seqlen_offsets, int):
+            ctx.save_for_backward(cos, sin, cu_seqlens)
+            ctx.seqlen_offsets = seqlen_offsets
+        else:
+            ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets)
+            ctx.seqlen_offsets = None
+        ctx.interleaved = interleaved
+        ctx.max_seqlen = max_seqlen
+        return qkv
+
+    @staticmethod
+    def backward(ctx, do):
+        seqlen_offsets = ctx.seqlen_offsets
+        if seqlen_offsets is None:
+            cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors
+        else:
+            cos, sin, cu_seqlens = ctx.saved_tensors
+        if do.is_contiguous():
+            total_nnz, three, nheads, headdim = do.shape
+            # Call 1 kernel instead of 2 kernels
+            # We need dqkv to be contiguous so that when we reshape to combine (3, nheads)
+            # dimensions, we get the same tensor
+            dqk = do[:, :2].view(total_nnz, -1, headdim)
+            apply_rotary(
+                dqk,
+                cos,
+                sin,
+                seqlen_offsets=seqlen_offsets,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=ctx.max_seqlen,
+                interleaved=ctx.interleaved,
+                inplace=True,
+                conjugate=True,
+            )
+        else:
+            dq, dk = do[:, 0, :, :], do[:, 1, :, :]
+            apply_rotary(
+                dq,
+                cos,
+                sin,
+                seqlen_offsets=seqlen_offsets,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=ctx.max_seqlen,
+                interleaved=ctx.interleaved,
+                inplace=True,
+                conjugate=True,
+            )
+            apply_rotary(
+                dk,
+                cos,
+                sin,
+                seqlen_offsets=seqlen_offsets,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=ctx.max_seqlen,
+                interleaved=ctx.interleaved,
+                inplace=True,
+                conjugate=True,
+            )
+
+        return do, None, None, None, None, None, None
+
+
+def apply_rotary_emb_unpad(
+    qkv,
+    cos,
+    sin,
+    interleaved=False,
+    seqlen_offsets: Union[int, torch.Tensor] = 0,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[int] = None,
+):
+    """
+    Arguments:
+        qkv: (total_nnz, 3, nheads, headdim) - input tensor for packed QKV.
+        cos, sin: (seqlen_rotary, rotary_dim / 2)
+        interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
+            of 1st half and 2nd half (GPT-NeoX style).
+        inplace: if True, apply rotary embedding in-place.
+        seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount.
+            Most commonly used in inference when we have KV cache.
+        cu_seqlens: (batch + 1,) or None
+        max_seqlen: int
+    Return:
+        out: (total_nnz, dim)
+    rotary_dim must be <= headdim
+    Apply rotary embedding to the first rotary_dim of x.
+    """
+    return ApplyRotaryEmbUnpad.apply(qkv, cos, sin, interleaved, seqlen_offsets, cu_seqlens, max_seqlen)
+
+
+class UnpaddedRotaryEmbedding(torch.nn.Module):
+    """
+    The rotary position embeddings applied directly to unpadded sequences.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        base: float = 10000.0,
+        interleaved: bool = False,
+        max_seqlen: Optional[int] = None,
+        scale_base: Optional[bool] = None,
+        pos_idx_in_fp32: bool = True,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        """
+        interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
+            of 1st half and 2nd half (GPT-NeoX style).
+        pos_idx_in_fp32: if True, the position indices [0.0, ..., seqlen - 1] are in fp32,
+            otherwise they might be in lower precision.
+            This option was added because previously (before 2023-07-02), when we construct
+            the position indices, we use the dtype of self.inv_freq. In most cases this would
+            be fp32, but if the model is trained in pure bf16 (not mixed precision), then
+            self.inv_freq would be bf16, and the position indices are also in bf16.
+            Because of the limited precision of bf16 (e.g. 1995.0 is rounded to 2000.0), the
+            embeddings for some positions will coincide.
+            To maintain compatibility with models previously trained in pure bf16,
+            we add this option.
+        max_seqlen: if max_seqlen, device, and dtype are provided, we precompute the cos_sin_cache
+            up to max_seqlen. If the max_seqlen, device, or dtype during training/inference differ,
+            the cos_sin_cache wll be recomputed during the forward pass.
+        """
+        super().__init__()
+        self.dim = dim
+        self.base = float(base)
+        self.pos_idx_in_fp32 = pos_idx_in_fp32
+        # Generate and save the inverse frequency buffer (non trainable)
+        inv_freq = self._compute_inv_freq(device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.interleaved = interleaved
+        self.scale_base = scale_base
+        scale = (
+            (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
+            if scale_base is not None
+            else None
+        )
+        self.register_buffer("scale", scale, persistent=False)
+
+        self._seq_len_cached = 0
+        self._cos_cached = None
+        self._sin_cached = None
+        self._cos_k_cached = None
+        self._sin_k_cached = None
+
+        if max_seqlen is not None and device is not None and dtype is not None:
+            self._update_cos_sin_cache(max_seqlen, device=device, dtype=dtype)
+
+    def _compute_inv_freq(self, device=None):
+        return 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
+
+    def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
+        # Reset the tables if the sequence length has changed,
+        # if we're on a new device (possibly due to tracing for instance),
+        # or if we're switching from inference mode to training
+        if (
+            seqlen > self._seq_len_cached
+            or self._cos_cached is None
+            or self._cos_cached.device != device
+            or self._cos_cached.dtype != dtype
+            or (self.training and self._cos_cached.is_inference())
+        ):
+            self._seq_len_cached = seqlen
+            # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16
+            # And the output of arange can be quite large, so bf16 would lose a lot of precision.
+            # However, for compatibility reason, we add an option to use the dtype of self.inv_freq.
+            if self.pos_idx_in_fp32:
+                t = torch.arange(seqlen, device=device, dtype=torch.float32)
+                # We want fp32 here as well since inv_freq will be multiplied with t, and the output
+                # will be large. Having it in bf16 will lose a lot of precision and cause the
+                # cos & sin output to change significantly.
+                # We want to recompute self.inv_freq if it was not loaded in fp32
+                if self.inv_freq.dtype != torch.float32:
+                    inv_freq = self._compute_inv_freq(device=device)
+                else:
+                    inv_freq = self.inv_freq
+            else:
+                t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
+                inv_freq = self.inv_freq
+            # Don't do einsum, it converts fp32 to fp16 under AMP
+            # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            freqs = torch.outer(t, inv_freq)
+            if self.scale is None:
+                self._cos_cached = torch.cos(freqs).to(dtype)
+                self._sin_cached = torch.sin(freqs).to(dtype)
+            else:
+                power = (
+                    torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
+                ) / self.scale_base
+                scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")
+                # We want the multiplication by scale to happen in fp32
+                self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
+                self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
+                self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
+                self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
+
+    def forward(
+        self,
+        qkv: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: Optional[int] = None,
+        seqlen_offset: Union[int, torch.Tensor] = 0,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        qkv: (total_nnz, 3, nheads, headdim)
+        cu_seqlens: (batch + 1,) cumulative sequence lengths
+        max_seqlen: int max seq length in the batch
+        seqlen_offset: (batch_size,) or int. Each sequence in x is shifted by this amount.
+            Most commonly used in inference when we have KV cache.
+            If it's a tensor of shape (batch_size,), then to update the cos / sin cache, one
+            should pass in max_seqlen, which will update the cos / sin cache up to that length.
+        Apply rotary embedding *inplace* to qkv.
+        """
+        if max_seqlen is not None:
+            self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)
+
+        qkv = apply_rotary_emb_unpad(
+            qkv,
+            self._cos_cached,
+            self._sin_cached,
+            interleaved=self.interleaved,
+            seqlen_offsets=seqlen_offset,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+
+        return qkv
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, base={self.base}, scale_base={self.scale_base}"

utils.py

@@ -0,0 +1,38 @@
+# Copyright 2020 Optuna, Hugging Face
+# License: Apache-2.0
+
+# Copyright 2023 OLMo Authors
+# License: Apache-2.0
+
+import functools
+import logging
+from enum import Enum
+
+
+@functools.lru_cache(None)
+def warning_once(self, *args, **kwargs):
+    """
+    This method is identical to `logger.warning()`, but will emit the warning with the same message only once
+
+    Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the cache.
+    The assumption here is that all warning messages are unique across the code. If they aren't then need to switch to
+    another type of cache that includes the caller frame information in the hashing function.
+    """
+    self.warning(*args, **kwargs)
+
+
+logging.Logger.warning_once = warning_once
+logging.Logger.warn_once = warning_once
+
+
+class StrEnum(str, Enum):
+    """
+    This is equivalent to Python's :class:`enum.StrEnum` since version 3.11.
+    We include this here for compatibility with older version of Python.
+    """
+
+    def __str__(self) -> str:
+        return self.value
+
+    def __repr__(self) -> str:
+        return f"'{str(self)}'"