modeling_flexbert.py

# 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,
    CausalLMOutput,
    MultipleChoiceModelOutput,
    SequenceClassifierOutput,
)
from transformers.models.bert.modeling_bert import BertPreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions
from .bert_padding import index_put_first_axis

from .activation import get_act_fn
from .attention import (
    FlexBertPaddedAttention,
    FlexBertPaddedParallelAttention,
    FlexBertPaddedRopeAttention,
    FlexBertPaddedRopeParallelAttention,
    FlexBertUnpadAttention,
    FlexBertUnpadParallelAttention,
    FlexBertUnpadRopeAttention,
    FlexBertUnpadRopeParallelAttention,
)
from .configuration_bert import FlexBertConfig
from .embeddings import (
    BertAlibiEmbeddings,
    FlexBertAbsoluteEmbeddings,
    FlexBertCompiledSansPositionEmbeddings,
    FlexBertSansPositionEmbeddings,
    get_embedding_layer,
)
from .initialization import (
    ModuleType,
    TileLinear,
    TileMode,
    init_weights,
    tile_embedding,
    tile_linear,
    tile_norm,
)
from .layers import (
    BertAlibiEncoder,
    BertPooler,
    BertPredictionHeadTransform,
    FlexBertCompileUnpadPreNormLayer,
    FlexBertPaddedEncoder,
    FlexBertPaddedParallelPreNormLayer,
    FlexBertPaddedPostNormLayer,
    FlexBertPaddedPreNormLayer,
    FlexBertUnpadEncoder,
    FlexBertUnpadParallelPreNormLayer,
    FlexBertUnpadPostNormLayer,
    FlexBertUnpadPreNormLayer,
    get_encoder_layer,
)
from .mlp import FlexBertGLU, FlexBertMLP, FlexBertParallelGLU
from .normalization import get_norm_layer
from .padding import pad_input, unpad_input
from .loss import get_loss_fn

# TODO: This is not used here, but this is so these files are copied when saving the model in ST/PyLate
from .utils import StrEnum
from .rotary import UnpaddedRotaryEmbedding


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 .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):
    config_class = FlexBertConfig 
    
    """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):
        # Handle the XOR condition
        assert (module is None) != (reset_params is None), "arg module xor reset_params must be specified"
        
        if module is not None:
            # Add basic initialization for common module types
            if isinstance(module, (nn.Linear, nn.Embedding)):
                module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
                if isinstance(module, nn.Linear) and module.bias is not None:
                    module.bias.data.zero_()
            elif isinstance(module, nn.LayerNorm):
                module.bias.data.zero_()
                module.weight.data.fill_(1.0)
        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:
            print(f"Padding logits: {logits.shape}")
            new_logits = self.pad_inputs(logits, indices, batch_size, seq_len)[0]
            if len(new_logits.shape) == 2:
                new_logits = new_logits.unsqueeze(0)
            return CausalLMOutput(
                loss=loss,
                logits=new_logits,
                hidden_states=None,
                attentions=None,
            )
        else:
            print(f"Non-padding logits: {logits.shape}")
            if len(logits.shape) == 2:
                logits = logits.unsqueeze(0)
            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

FlexBertForCasualLM.register_for_auto_class("AutoModelForCausalLM")

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)