Update model code

configuration.py

@@ -1,45 +1,242 @@
-# Adapted from https://github.com/mosaicml/llm-foundry
-# Classes changed: MPTConfig
-# SPDX-License-Identifier: Apache-2.0
-
-
-"""A HuggingFace-style model configuration."""
-
-from typing import Dict, List, Optional, Union
-from transformers import PretrainedConfig
-
-attn_config_defaults: Dict = {
-    'attn_type': 'multihead_attention',
-    'attn_pdrop': 0.0,
-    'attn_impl': 'torch',
-    'qk_ln': False,
-    'clip_qkv': None,
-    'softmax_scale': None,
-    'prefix_lm': False,
-    'attn_uses_sequence_id': False,
-    'alibi': True,
-    'alibi_bias_max': 8,
-    "topk": 10,
-    'mask_by_sim':True,
-    'sim_threshold':0.25,
-    'use_active_externalism':True,
-    'memory_type':'manual'
-}
-
-init_config_defaults: Dict = {
-    'name': 'kaiming_normal_',
-    'fan_mode': 'fan_in',
-    'init_nonlinearity': 'relu',
-    'init_div_is_residual': True,
-    'emb_init_std': None,
-    'emb_init_uniform_lim': None,
-    'init_std': None,
-    'init_gain': 0.0,
-}
-
-
-class ExtendedMPTConfig(PretrainedConfig):
-    model_type = 'extended-mpt'
+# Copyright 2023 HuggingFace Inc. team and MosaicML NLP team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This code has been adapted from Mosaic ML and Huggingface and inherits the above lisence.
+# The original code can be found here:
+# https://github.com/huggingface/transformers/blob/main/src/transformers/models/mpt/configuration_mpt.py
+
+"""Extended Mind Mpt configuration"""
+from typing import Optional, Union
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class ExtendedMptAttentionConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`ExtendedMptAttention`] class. It is used to instantiate
+    attention layers according to the specified arguments, defining the layers architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the MPT
+    [mosaicml/mpt-7b](https://huggingface.co/mosaicml/mpt-7b) architecture. Most of the arguments are kept for backward
+    compatibility with previous MPT models that are hosted on the Hub (previously with `trust_remote_code=True`).
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        attn_type (`str`, *optional*, defaults to `"multihead_attention"`):
+            type of attention to use. Options: `"multihead_attention"`, `"multiquery_attention"`.
+        attn_pdrop (`float`, *optional*, defaults to 0.0):
+            The dropout probability for the attention layers.
+        attn_impl (`str`, *optional*, defaults to `"torch"`):
+            The attention implementation to use. One of `"torch"`, `"flash"`, or `"triton"`.
+        clip_qkv (`float`, *optional*):
+            If not `None`, clip the queries, keys, and values in the attention layer to this value.
+        softmax_scale (`float`, *optional*, defaults to `None`):
+            If not `None`, scale the softmax in the attention layer by this value. If `None`, will default to
+            `1/sqrt(hidden_size)`.
+        prefix_lm (`bool`, *optional*, defaults to `False`)):
+            Whether the model should operate as a Prefix LM. This requires passing an extra `prefix_mask` argument
+            which indicates which tokens belong to the prefix. Tokens in the prefix can attend to one another
+            bi-directionally. Tokens outside the prefix use causal attention.
+        qk_ln (`bool`, *optional*, defaults to `False`):
+            Whether to apply layer normalization to the queries and keys in the attention layer.
+        attn_uses_sequence_id (`bool`, *optional*, defaults to `False`)):
+            Whether to restrict attention to tokens that have the same token_type_ids. When the model is in `train`
+            mode, this requires passing an extra *token_type_ids* argument which indicates which sub-sequence each
+            token belongs to. Defaults to `False` meaning any provided *token_type_ids* will be ignored.
+        alibi (`bool`, *optional*, defaults to `True`):
+            Whether or not to use the alibi bias instead of positional embedding.
+        alibi_bias_max (`int`, *optional*, defaults to 8):
+            The maximum value of the alibi bias.
+
+        #### Memory Configuration ####
+        topk (`int`, *optional*, defaults to `10`):
+            Number of external memories for each query token to retrieve and attend to.
+        memory_type (`string`, *optional*, defaults to `manual`):
+            Whether to store external memories manually or in a vector database.
+        memory_device (`string`, *optional*, defaults to `cpu`):
+            Specify device to store memory.
+        mask_by_sim (`bool`, *optional*, defaults to `True`):
+            Whether or not to mask retrieved memories by similarity.
+        sim_threshold (`float`, *optional*, defaults to `0.25`):
+            Threshold for masking retrieved memories.
+        tokenizer_all_special_ids (`list`, *optional*, defaults to `[0, 50278]`):
+            Ids for special tokens to remove from memories.
+        remove_special_tokens (`bool`, *optional*, defaults to `True`):
+            Remove memories that correspond to tokenizer special ids.
+        #### Memory Configuration ####
+
+    """
+
+    def __init__(
+        self,
+        attn_type="multihead_attention",
+        attn_pdrop=0,
+        attn_impl="torch",
+        clip_qkv=None,
+        softmax_scale=None,
+        prefix_lm=False,
+        qk_ln=False,
+        attn_uses_sequence_id=False,
+        alibi=True,
+        alibi_bias_max=8,
+        topk=10,
+        memory_type="manual",
+        memory_device="cpu",
+        mask_by_sim=True,
+        sim_threshold=0.25,
+        tokenizer_all_special_ids=[0, 50278],
+        remove_special_ids=False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.attn_type = attn_type
+        self.attn_pdrop = attn_pdrop
+        self.attn_impl = attn_impl
+        self.clip_qkv = clip_qkv
+        self.softmax_scale = softmax_scale
+        self.prefix_lm = prefix_lm
+        self.attn_uses_sequence_id = attn_uses_sequence_id
+        self.alibi = alibi
+        self.qk_ln = qk_ln
+        self.alibi_bias_max = alibi_bias_max
+        self.topk = topk
+        self.memory_type = memory_type
+        self.memory_device = memory_device
+        self.mask_by_sim = mask_by_sim
+        self.sim_threshold = sim_threshold
+        self.tokenizer_all_special_ids = tokenizer_all_special_ids
+        self.remove_special_ids = remove_special_ids
+
+        if attn_type not in ["multihead_attention", "multiquery_attention"]:
+            raise ValueError(
+                f"`attn_type` has to be either `multihead_attention` or `multiquery_attention`. Received: {attn_type}"
+            )
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_name_or_path, **kwargs
+    ) -> "PretrainedConfig":
+        cls._set_token_in_kwargs(kwargs)
+
+        config_dict, kwargs = cls.get_config_dict(
+            pretrained_model_name_or_path, **kwargs
+        )
+
+        if config_dict.get("model_type") == "mpt":
+            config_dict = config_dict["attn_config"]
+
+        if (
+            "model_type" in config_dict
+            and hasattr(cls, "model_type")
+            and config_dict["model_type"] != cls.model_type
+        ):
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class ExtendedMptConfig(PretrainedConfig):
+    """
+    This is the configuration class to store the configuration of a [`MptModel`]. It is used to instantiate a Mpt model
+    according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to the Mpt-7b architecture
+    [mosaicml/mpt-7b](https://huggingface.co/mosaicml/mpt-7b).
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        d_model (`int`, *optional*, defaults to 2048):
+            Dimensionality of the embeddings and hidden states.
+        n_heads (`int`, *optional*, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        n_layers (`int`, *optional*, defaults to 24):
+            Number of hidden layers in the Transformer encoder.
+        expansion_ratio (`int`, *optional*, defaults to 4):
+            The ratio of the up/down scale in the MLP.
+        max_seq_len (`int`, *optional*, defaults to 2048):
+            The maximum sequence length of the model.
+        vocab_size (`int`, *optional*, defaults to 50368):
+            Vocabulary size of the Mpt model. Defines the maximum number of different tokens that can be represented by
+            the `inputs_ids` passed when calling [`MptModel`]. Check [this
+            discussion](https://huggingface.co/bigscience/mpt/discussions/120#633d28389addb8530b406c2a) on how the
+            `vocab_size` has been defined.
+        resid_pdrop (`float`, *optional*, defaults to 0.1):
+            The dropout probability applied to the attention output before combining with residual.
+        layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
+            The epsilon to use in the layer normalization layers.
+        emb_pdrop (`float`, *optional*, defaults to 0.1):
+            The dropout probability for the embedding layer.
+        learned_pos_emb (`bool`, *optional*, defaults to `False`):
+            Whether to use learned positional embeddings.
+        attn_config (`dict`, *optional*):
+            A dictionary used to configure the model's attention module.
+        init_device (`str`, *optional*):
+            The device to use for parameter initialization. Defined for backward compatibility
+        logit_scale (`float`, *optional*):
+            If not None, scale the logits by this value.
+        no_bias (`bool`, *optional*, defaults to `True`):
+            Whether to use bias in all linear layers.
+        verbose (`int`, *optional*, defaults to 0):
+            The verbosity level to use for logging. Used in the previous versions of MPT models for logging. This
+            argument is deprecated.
+        embedding_fraction (`float`, *optional*, defaults to 1.0):
+            The fraction to scale the gradients of the embedding layer by.
+        norm_type (`str`, *optional*, defaults to `"low_precision_layernorm"`):
+            Type of layer norm to use. All MPT models uses the same layer norm implementation. Defined for backward
+            compatibility.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models).
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+
+        #### Memory Configuration ####
+        use_external_mind (`bool`, *optional*, defaults to `True`):
+            Whether to attend to external memories.
+        use_external_mind_by_layer (`List[bool]`, *optional*, defaults to List[`True`, ..., `True`]):
+            Whether to attend to external memories, on each decoder layer.
+        #### Memory Configuration ####
+
+    Example:
+
+    ```python
+    >>> from transformers import MptConfig, MptModel
+
+    >>> # Initializing a Mpt configuration
+    >>> configuration = MptConfig()
+
+    >>> # Initializing a model (with random weights) from the configuration
+    >>> model = MptModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "extended-mpt"
+    attribute_map = {
+        "num_attention_heads": "n_heads",
+        "hidden_size": "d_model",
+        "num_hidden_layers": "n_layers",
+    }
 
     def __init__(
         self,
@@ -47,166 +244,51 @@ class ExtendedMPTConfig(PretrainedConfig):
         n_heads: int = 32,
         n_layers: int = 32,
         expansion_ratio: int = 4,
-        max_seq_len: int = 2048,
+        max_seq_len_inference: int = 2048,
+        max_seq_len_train: int = 2048,
         vocab_size: int = 50432,
         resid_pdrop: float = 0.0,
+        layer_norm_epsilon: float = 1e-5,
         emb_pdrop: float = 0.0,
         learned_pos_emb: bool = True,
-        attn_config: Dict = attn_config_defaults,
-        init_device: str = 'cpu',
+        attn_config: ExtendedMptAttentionConfig = None,
+        init_device: str = "cpu",
         logit_scale: Optional[Union[float, str]] = None,
         no_bias: bool = True,
         verbose: int = 0,
         embedding_fraction: float = 1.0,
-        norm_type: str = 'low_precision_layernorm',
+        norm_type: str = "low_precision_layernorm",
         use_cache: bool = False,
-        init_config: Dict = init_config_defaults,
-        use_active_externalism_by_layer: List[bool] = [True for _ in range(32)],
-        memory_device:str = 'cpu',
+        initializer_range=0.02,
+        use_external_mind: bool = True,
+        use_external_mind_by_layer: list[bool] = [True for _ in range(32)],
         **kwargs,
     ):
-        """The MPT configuration class.
-
-        Args:
-            d_model (int): The size of the embedding dimension of the model.
-            n_heads (int): The number of attention heads.
-            n_layers (int): The number of layers in the model.
-            expansion_ratio (int): The ratio of the up/down scale in the MLP.
-            max_seq_len (int): The maximum sequence length of the model.
-            vocab_size (int): The size of the vocabulary.
-            resid_pdrop (float): The dropout probability applied to the attention output before combining with residual.
-            emb_pdrop (float): The dropout probability for the embedding layer.
-            learned_pos_emb (bool): Whether to use learned positional embeddings
-            attn_config (Dict):  A dictionary used to configure the model's attention module:
-                attn_type (str): type of attention to use. Options: multihead_attention, multiquery_attention
-                attn_pdrop (float): The dropout probability for the attention layers.
-                attn_impl (str): The attention implementation to use. One of 'torch', 'flash', or 'triton'.
-                qk_ln (bool): Whether to apply layer normalization to the queries and keys in the attention layer.
-                clip_qkv (Optional[float]): If not None, clip the queries, keys, and values in the attention layer to
-                    this value.
-                softmax_scale (Optional[float]): If not None, scale the softmax in the attention layer by this value. If None,
-                    use the default scale of ``1/sqrt(d_keys)``.
-                prefix_lm (Optional[bool]): Whether the model should operate as a Prefix LM. This requires passing an
-                    extra `prefix_mask` argument which indicates which tokens belong to the prefix. Tokens in the prefix
-                    can attend to one another bi-directionally. Tokens outside the prefix use causal attention.
-                attn_uses_sequence_id (Optional[bool]): Whether to restrict attention to tokens that have the same sequence_id.
-                    When the model is in `train` mode, this requires passing an extra `sequence_id` argument which indicates
-                    which sub-sequence each token belongs to.
-                    Defaults to ``False`` meaning any provided `sequence_id` will be ignored.
-                alibi (bool): Whether to use the alibi bias instead of position embeddings.
-                alibi_bias_max (int): The maximum value of the alibi bias.
-            init_device (str): The device to use for parameter initialization.
-            logit_scale (Optional[Union[float, str]]): If not None, scale the logits by this value.
-            no_bias (bool): Whether to use bias in all layers.
-            verbose (int): The verbosity level. 0 is silent.
-            embedding_fraction (float): The fraction to scale the gradients of the embedding layer by.
-            norm_type (str): choose type of norm to use
-            multiquery_attention (bool): Whether to use multiquery attention implementation.
-            use_cache (bool): Whether or not the model should return the last key/values attentions
-            init_config (Dict): A dictionary used to configure the model initialization:
-                init_config.name: The parameter initialization scheme to use. Options: 'default_', 'baseline_',
-                    'kaiming_uniform_', 'kaiming_normal_', 'neox_init_', 'small_init_', 'xavier_uniform_', or
-                    'xavier_normal_'. These mimic the parameter initialization methods in PyTorch.
-                init_div_is_residual (Union[int, float, str, bool]): Value to divide initial weights by if ``module._is_residual`` is True.
-                emb_init_std (Optional[float]): The standard deviation of the normal distribution used to initialize the embedding layer.
-                emb_init_uniform_lim (Optional[Union[Tuple[float, float], float]]): The lower and upper limits of the uniform distribution
-                    used to initialize the embedding layer. Mutually exclusive with ``emb_init_std``.
-                init_std (float): The standard deviation of the normal distribution used to initialize the model,
-                    if using the baseline_ parameter initialization scheme.
-                init_gain (float): The gain to use for parameter initialization with kaiming or xavier initialization schemes.
-                fan_mode (str): The fan mode to use for parameter initialization with kaiming initialization schemes.
-                init_nonlinearity (str): The nonlinearity to use for parameter initialization with kaiming initialization schemes.
-                ---
-                See llmfoundry.models.utils.param_init_fns.py for info on other param init config options
-        """
+        if attn_config is None:
+            self.attn_config = ExtendedMptAttentionConfig()
+        elif not isinstance(attn_config, ExtendedMptAttentionConfig):
+            self.attn_config = ExtendedMptAttentionConfig(**attn_config)
+        else:
+            self.attn_config = attn_config
         self.d_model = d_model
         self.n_heads = n_heads
         self.n_layers = n_layers
         self.expansion_ratio = expansion_ratio
-        self.max_seq_len = max_seq_len
+        self.max_seq_len = max_seq_len_inference
+        self.max_seq_len_train = max_seq_len_train
         self.vocab_size = vocab_size
         self.resid_pdrop = resid_pdrop
         self.emb_pdrop = emb_pdrop
         self.learned_pos_emb = learned_pos_emb
-        self.attn_config = attn_config
         self.init_device = init_device
         self.logit_scale = logit_scale
         self.no_bias = no_bias
         self.verbose = verbose
         self.embedding_fraction = embedding_fraction
         self.norm_type = norm_type
+        self.layer_norm_epsilon = layer_norm_epsilon
         self.use_cache = use_cache
-        self.init_config = init_config
-        self.use_active_externalism_by_layer = use_active_externalism_by_layer
-        self.memory_device = memory_device
-        if 'name' in kwargs:
-            del kwargs['name']
-        if 'loss_fn' in kwargs:
-            del kwargs['loss_fn']
+        self.initializer_range = initializer_range
+        self.use_external_mind = use_external_mind
+        self.use_external_mind_by_layer = use_external_mind_by_layer
         super().__init__(**kwargs)
-
-        self._validate_config()
-
-    def _set_config_defaults(self, config, config_defaults):
-        # set config defaults
-        for k, v in config_defaults.items():
-            if k not in config:
-                config[k] = v
-        return config
-
-    def _validate_config(self):
-        # set config defaults
-        self.attn_config = self._set_config_defaults(
-            self.attn_config,
-            attn_config_defaults,
-        )
-        self.init_config = self._set_config_defaults(
-            self.init_config,
-            init_config_defaults,
-        )
-
-        if self.attn_config['memory_type']=='faiss' and self.attn_config['mask_by_sim'] is True:
-            raise ValueError(
-                'mask_by_sim is not supported for faiss memory type.'
-            ) 
-
-        if self.d_model % self.n_heads != 0:
-            raise ValueError('d_model must be divisible by n_heads')
-        if any(
-                prob < 0 or prob > 1 for prob in
-            [self.attn_config['attn_pdrop'], self.resid_pdrop, self.emb_pdrop]):
-            raise ValueError(
-                "self.attn_config['attn_pdrop'], resid_pdrop, emb_pdrop are probabilities and must be between 0 and 1"
-            )
-        if self.attn_config['attn_impl'] not in ['torch', 'flash', 'triton']:
-            raise ValueError(
-                f"Unknown attn_impl={self.attn_config['attn_impl']}")
-        if self.attn_config['prefix_lm'] and self.attn_config[
-                'attn_impl'] not in ['torch', 'triton']:
-            raise NotImplementedError(
-                'prefix_lm only implemented with torch and triton attention.')
-        if self.attn_config['alibi'] and self.attn_config['attn_impl'] not in [
-                'torch', 'triton'
-        ]:
-            raise NotImplementedError(
-                'alibi only implemented with torch and triton attention.')
-        if self.attn_config['attn_uses_sequence_id'] and self.attn_config[
-                'attn_impl'] not in ['torch', 'triton']:
-            raise NotImplementedError(
-                'attn_uses_sequence_id only implemented with torch and triton attention.'
-            )
-        if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
-            raise ValueError(
-                'model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!'
-            )
-        if isinstance(self.logit_scale,
-                      str) and self.logit_scale != 'inv_sqrt_d_model':
-            raise ValueError(
-                f"{self.logit_scale=} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'."
-            )
-        if self.init_config.get('name', None) is None:
-            raise ValueError(f"{self.init_config=} 'name' needs to be set.")
-        if not self.learned_pos_emb and not self.attn_config['alibi']:
-            raise ValueError(
-                f'Positional information must be provided to the model using either learned_pos_emb or alibi.'
-            )

modeling.py

@@ -0,0 +1,1297 @@
+# Copyright 2023 HuggingFace Inc. team and MosaicML NLP team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This code has been adapted from Mosaic ML and Huggingface and inherits the above lisence.
+# The original code can be found here:
+# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
+# We annotate the edited code below with 'EM' comments to indicate where we have made changes.
+"""PyTorch MPT model."""
+
+import math
+from typing import Optional, Tuple, Union
+
+import faiss
+import numpy as np
+import torch
+import torch.utils.checkpoint
+from einops import rearrange
+from torch import nn
+from torch.linalg import vector_norm
+from torch.nn import CrossEntropyLoss, LayerNorm
+from torch.nn import functional as F
+from transformers.file_utils import (
+    add_code_sample_docstrings,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from .configuration import ExtendedMptConfig
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "mosaicml/mpt-7b"
+_CONFIG_FOR_DOC = "MptConfig"
+
+
+# Copied from transformers.models.bloom.modeling_bloom._make_causal_mask
+def _make_causal_mask(
+    input_ids_shape: torch.Size, device: torch.device, past_key_values_length: int
+) -> torch.BoolTensor:
+    """
+    Make causal mask used for self-attention.
+    """
+    batch_size, target_length = input_ids_shape
+    mask = torch.empty(
+        (target_length, target_length + past_key_values_length),
+        dtype=torch.bool,
+        device=device,
+    )
+    # ONNX doesn't support `torch.Tensor.triu` properly, thus we use this workaround
+    seq_ids = torch.arange(target_length, device=device)
+    mask[:, past_key_values_length:] = seq_ids[:, None] < seq_ids[None, :]
+
+    if past_key_values_length > 0:
+        mask[:, :past_key_values_length] = False
+
+    expanded_mask = mask[None, None, :, :].expand(
+        batch_size, 1, target_length, target_length + past_key_values_length
+    )
+    return expanded_mask
+
+
+# Copied from transformers.models.bloom.modeling_bloom._expand_mask
+def _expand_mask(mask: torch.Tensor, tgt_length: int) -> torch.BoolTensor:
+    """
+    Expands attention_mask from `[batch_size, src_length]` to `[batch_size, 1, tgt_length, src_length]`.
+    """
+    batch_size, src_length = mask.shape
+    tgt_length = tgt_length if tgt_length is not None else src_length
+
+    expanded_mask = ~(mask[:, None, None, :].to(torch.bool))
+    return expanded_mask.expand(batch_size, 1, tgt_length, src_length)
+
+
+def build_mpt_alibi_tensor(
+    num_heads,
+    sequence_length,
+    sequence_length_with_past,
+    alibi_bias_max=8,
+    device=None,
+    for_ae=False,
+    topk=None,
+):
+    r"""
+    Link to paper: https://arxiv.org/abs/2108.12409 - Alibi tensor is not causal as the original paper mentions, it
+    relies on a translation invariance of softmax for quick implementation. This implementation has been copied from
+    the alibi implementation of MPT source code that led to slightly different results than the Bloom alibi:
+    https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py#L292
+    """
+    if not for_ae:
+        alibi = torch.arange(
+            1 - sequence_length, 1, dtype=torch.int32, device=device
+        ).view(1, 1, 1, sequence_length)
+    else:  # EM: All memory tokens get same bias
+        alibi = (
+            torch.tensor(-sequence_length_with_past, dtype=torch.int32, device=device)
+            .repeat(sequence_length * topk)
+            .view(1, 1, 1, sequence_length * topk)
+        )
+    num_heads_power_of_2 = 2 ** math.ceil(math.log2(num_heads))
+
+    base = torch.arange(1, num_heads_power_of_2 + 1, dtype=torch.float32, device=device)
+    base = base * (alibi_bias_max / num_heads_power_of_2)
+
+    slopes = 1.0 / torch.pow(2, base)
+    slopes = slopes.view(1, num_heads, 1, 1)
+
+    if num_heads_power_of_2 != num_heads:
+        slopes = torch.concat([slopes[1::2], slopes[::2]])[:num_heads]
+
+    alibi = alibi * slopes
+    return alibi.squeeze(0)
+
+
+class ExtendedMptAttention(nn.Module):
+    """Multi-head self attention.
+    Using torch or triton attention implemetation enables user to also use additive bias.
+    """
+
+    def __init__(self, config: ExtendedMptConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.n_heads = config.n_heads
+        self.n_layers = config.n_layers
+        self.head_dim = self.hidden_size // self.n_heads
+        self.softmax_scale = config.attn_config.softmax_scale
+        if self.softmax_scale is None:
+            self.softmax_scale = 1 / math.sqrt(self.hidden_size / self.n_heads)
+
+        self.attn_dropout_p = config.attn_config.attn_pdrop
+        self.Wqkv = nn.Linear(self.hidden_size, 3 * self.hidden_size, bias=False)
+        self.out_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_bias: torch.Tensor,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        long_range_past_key_value=None,
+        topk=None,
+        faiss_indexes=None,
+        mask_by_sim=None,
+        sim_threshold=None,
+        position_bias_ae=None,
+        current_layer=None,
+        output_retrieved_memory_idx=False,
+    ):
+        batch_size, seq_length = hidden_states.shape[:2]
+
+        mixed_qkv = self.Wqkv(hidden_states)
+        query_states, key_states, value_states = mixed_qkv.chunk(3, dim=2)
+        query_states = query_states.reshape(
+            batch_size, seq_length, self.n_heads, self.head_dim
+        ).transpose(1, 2)
+        key_states = key_states.reshape(
+            batch_size, seq_length, self.n_heads, self.head_dim
+        ).transpose(1, 2)
+        value_states = value_states.reshape(
+            batch_size, seq_length, self.n_heads, self.head_dim
+        ).transpose(1, 2)
+
+        if past_key_value is not None:
+            if len(past_key_value) != 0:
+                key_states = torch.cat([past_key_value[0], key_states], dim=2)
+                value_states = torch.cat([past_key_value[1], value_states], dim=2)
+        past_key_value = (key_states, value_states)
+        bsz, nh, s_q, d = query_states.shape
+
+        attention_scores = (
+            torch.matmul(query_states, key_states.transpose(-1, -2))
+            * self.softmax_scale
+        )
+        key_length = key_states.shape[-2]
+        query_length = (
+            seq_length
+            if past_key_value is None
+            else seq_length + past_key_value[0].shape[2]
+        )
+        if position_bias is not None:
+            if len(position_bias.shape) != 3:
+                raise ValueError(
+                    f"Expecting position_bias shape to be 3 dimensions, got {len(position_bias.shape)}"
+                )
+
+            position_bias_query_index = max(0, position_bias.size(1) - query_length)
+            position_bias_key_index = max(0, position_bias.size(2) - key_length)
+
+            position_bias = position_bias[
+                :, position_bias_query_index:, position_bias_key_index:
+            ]
+
+            attention_scores = attention_scores + position_bias
+
+        # EM: Retrieve memories from cache or faiss indexes
+        if long_range_past_key_value is not None or faiss_indexes is not None:
+            if long_range_past_key_value is not None:  # Manual store
+                k_cache, v_cache = long_range_past_key_value
+                s_cache = k_cache.size(-2)
+
+                k_cache = k_cache.to(key_states.device)
+                v_cache = v_cache.to(key_states.device)
+
+                # Normalize query and key vectors
+                q_n = query_states / vector_norm(
+                    query_states, ord=2, dim=-1, keepdim=True
+                ) 
+                k_n = k_cache / vector_norm(k_cache, ord=2, dim=-1, keepdim=True)
+                sim = q_n.matmul(k_n.transpose(-1, -2))
+                if s_cache < topk:   # number of tokens in cache < topk
+                    topk = s_cache
+                val, idx = torch.topk(sim, k=topk, dim=-1)  # Retrieve topk memories
+
+                reshaped_idx = idx.reshape(bsz, nh, s_q * topk)
+                selected_k = k_cache.gather(
+                    dim=-2, index=reshaped_idx.unsqueeze(-1).expand(-1, -1, -1, d)
+                )
+                selected_v = v_cache.gather(
+                    dim=-2, index=reshaped_idx.unsqueeze(-1).expand(-1, -1, -1, d)
+                )
+
+            elif faiss_indexes is not None:  # FAISS indexes
+                kn_index, kv_index = faiss_indexes
+                q_n = query_states / vector_norm(
+                    query_states, ord=2, dim=-1, keepdim=True
+                )
+                # One-hot encoding for layer, head to only retrieve memories from the same layer, head
+                one_hot_encodings = (
+                    F.one_hot(
+                        torch.arange(0, nh * self.n_layers, device=query_states.device)
+                    )
+                    * 10
+                )
+                q_n = torch.concat(
+                    [
+                        rearrange(q_n, "b h s d -> b (h s) d", h=nh),
+                        one_hot_encodings[nh * current_layer : nh * (current_layer + 1)]
+                        .unsqueeze(0)
+                        .repeat_interleave(repeats=query_states.size(-2), dim=-2),
+                    ],
+                    dim=-1,
+                ).squeeze()
+
+                if kn_index.ntotal / (nh * self.n_layers) < topk:
+                    topk = int(kn_index.ntotal / (nh * self.n_layers))
+
+                val, idx = kn_index.search(q_n.to("cpu").detach().numpy(), k=topk)
+                val = torch.tensor(val - 100).reshape(bsz, nh, s_q, topk)  #Similarity includes scale factor from one-hot encoding
+                reshaped_idx = torch.tensor(
+                    idx % (kn_index.ntotal / (nh * self.n_layers))
+                ).reshape(bsz, nh, s_q * topk)
+
+                # Retrieve tensors
+                selected_k = rearrange(
+                    torch.tensor(kv_index.reconstruct_batch(idx.flatten()))[:, :d],
+                    "(h s) d -> 1 h s d",
+                    h=nh,
+                ).to(query_states.device)
+                selected_v = rearrange(
+                    torch.tensor(kv_index.reconstruct_batch(idx.flatten()))[:, d:],
+                    "(h s) d -> 1 h s d",
+                    h=nh,
+                ).to(query_states.device)
+
+            selected_key_length = selected_k.size(-2)
+            key_length += selected_key_length
+            attention_scores_cache = (
+                query_states.matmul(selected_k.transpose(-1, -2)) * self.softmax_scale
+            )
+            # EM: Mask by similarity
+            if mask_by_sim:
+                sim_mask = (
+                    rearrange(~(val > sim_threshold).bool(), "b h s i -> b h (s i)")
+                    .unsqueeze(-2)
+                    .expand(-1, -1, s_q, -1)
+                ).to(query_states.device)
+
+                attention_scores_cache = attention_scores_cache.masked_fill(
+                    sim_mask, torch.finfo(query_states.dtype).min
+                )
+
+            # EM: Add position bias to cache
+            if position_bias_ae is not None:
+                if len(position_bias_ae.shape) != 3:
+                    raise ValueError(
+                        f"Expecting position_bias shape to be 3 dimensions, got {len(position_bias_ae.shape)}"
+                    )
+
+                position_bias_query_index = max(
+                    0, position_bias_ae.size(1) - query_length
+                )
+                position_bias_key_index = max(
+                    0, position_bias_ae.size(2) - selected_key_length
+                )
+
+                position_bias_ae = position_bias_ae[
+                    :, position_bias_query_index:, position_bias_key_index:
+                ]
+
+                attention_scores_cache = attention_scores_cache + position_bias_ae
+
+            # EM: Concatenate cache and current attention weights, values
+            attention_scores = torch.cat(
+                [attention_scores_cache, attention_scores], dim=-1
+            )  # Concat attention scores, values
+            value_states = torch.cat([selected_v, value_states], dim=-2)
+
+        # EM: Create mask for external memories, queries only attend to their own memories
+        def _create_external_memories_mask(k, s_q, device):
+            mask = torch.zeros(s_q, s_q * k, device=device, dtype=torch.bool)
+            for i in range(s_q):
+                mask[i, i * k : (i + 1) * k] = 1
+            return ~mask
+
+        if attention_mask is not None:
+            # EM: Concatenate attention mask with external memories mask
+            if long_range_past_key_value is not None or faiss_indexes is not None:
+                mask = _create_external_memories_mask(
+                    k=topk, s_q=s_q, device=attention_scores.device
+                )
+                attention_mask = attention_mask.squeeze(dim=0).squeeze(dim=0)
+                attention_mask = torch.cat([mask, attention_mask], dim=1)
+            attention_scores = attention_scores.masked_fill(
+                attention_mask, torch.finfo(query_states.dtype).min
+            )
+
+        # (batch_size, n_heads, seq_length, key_length)
+        attn_weights = nn.functional.softmax(attention_scores.float(), dim=-1).to(
+            value_states.dtype
+        )
+        attn_weights = nn.functional.dropout(
+            attn_weights, p=self.attn_dropout_p, training=self.training
+        )
+
+        context_states = torch.matmul(attn_weights, value_states)
+        context_states = (
+            context_states.permute(0, 2, 1, 3)
+            .contiguous()
+            .view(batch_size, seq_length, -1)
+        )
+        attn_output = self.out_proj(context_states)
+
+        if not output_retrieved_memory_idx:
+            reshaped_idx = None
+
+        return attn_output, attn_weights, past_key_value, reshaped_idx
+
+
+class MptMLP(nn.Module):
+    def __init__(self, config: ExtendedMptConfig):
+        super().__init__()
+        hidden_size = config.hidden_size
+
+        self.up_proj = nn.Linear(hidden_size, 4 * hidden_size, bias=False)
+        self.act = nn.GELU(approximate="none")
+        self.down_proj = nn.Linear(4 * hidden_size, hidden_size, bias=False)
+        self.hidden_dropout = config.attn_config.attn_pdrop
+
+    def forward(
+        self, hidden_states: torch.Tensor, residual: torch.Tensor
+    ) -> torch.Tensor:
+        hidden_states = self.act(self.up_proj(hidden_states))
+
+        intermediate_output = self.down_proj(hidden_states)
+
+        output = F.dropout(
+            intermediate_output, p=self.hidden_dropout, training=self.training
+        )
+        output = output + residual
+
+        return output
+
+
+class MptBlock(nn.Module):
+    """MPTBlock"""
+
+    def __init__(self, config: ExtendedMptConfig):
+        super().__init__()
+        hidden_size = config.hidden_size
+
+        self.norm_1 = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        # backward compatibility with weights on the Hub
+        self.norm_1.bias = None
+
+        self.num_heads = config.n_heads
+        self.attn = ExtendedMptAttention(config)
+
+        self.norm_2 = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        # backward compatibility with weights on the Hub
+        self.norm_2.bias = None
+
+        self.ffn = MptMLP(config)
+
+        self.dropout_rate = config.attn_config.attn_pdrop
+        self.resid_attn_dropout = nn.Dropout(self.dropout_rate)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_bias: torch.Tensor,
+        attention_mask: torch.Tensor,
+        layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+        output_attentions: bool = False,
+        output_retrieved_memory_idx: bool = False,
+        topk: int = None,
+        long_range_past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        faiss_indexes: Tuple = None,
+        position_bias_ae=None,
+        current_layer: int = None,
+        mask_by_sim: bool = False,
+        sim_threshold: float = None,
+    ):
+        # hidden_states: [batch_size, seq_length, hidden_size]
+        # Layer norm at the beginning of the transformer layer.
+        layernorm_output = self.norm_1(hidden_states)
+
+        residual = hidden_states
+
+        # Self attention.
+        attn_outputs, attn_weights, past_key_value, reshaped_idx = self.attn(
+            layernorm_output,
+            position_bias=position_bias,
+            attention_mask=attention_mask,
+            past_key_value=layer_past,
+            long_range_past_key_value=long_range_past_key_value,
+            topk=topk,
+            faiss_indexes=faiss_indexes,
+            position_bias_ae=position_bias_ae,
+            current_layer=current_layer,
+            mask_by_sim=mask_by_sim,
+            sim_threshold=sim_threshold,
+            output_retrieved_memory_idx=output_retrieved_memory_idx,
+        )
+
+        hidden_states = self.resid_attn_dropout(attn_outputs) + residual
+
+        layernorm_output = self.norm_2(hidden_states)
+
+        # Get residual
+        residual = hidden_states
+
+        # MLP.
+        output = self.ffn(layernorm_output, residual)
+        outputs = (output,)
+
+        if use_cache:
+            outputs += (past_key_value,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+        if output_retrieved_memory_idx:
+            outputs += (reshaped_idx,)
+
+        return outputs  # hidden_states, present, attentions
+
+
+class MptPreTrainedModel(PreTrainedModel):
+    """MPT Pretrained Model"""
+
+    config_class = ExtendedMptConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MptBlock"]
+    _keys_to_ignore_on_load_missing = [r"lm_head.*."]
+
+    def __init__(self, *inputs, **kwargs):
+        super().__init__(*inputs, **kwargs)
+
+    def _init_weights(self, module: nn.Module):
+        """Initialize the weights."""
+        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=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, LayerNorm):
+            if module.bias is not None:
+                module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False):
+        if isinstance(module, ExtendedMptConfig):
+            module.gradient_checkpointing = value
+
+    @staticmethod
+    def _convert_to_mpt_cache(
+        past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]]
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Converts the cache to the format expected by Mpt, i.e. to tuple(tuple([batch_size * num_heads, ...]))
+        """
+        batch_size, num_heads, head_dim, seq_length = past_key_value[0][0].shape
+        batch_size_times_num_heads = batch_size * num_heads
+        # key:  [batch_size, num_heads, head_dim, seq_length] -> [batch_size * num_heads, head_dim, seq_length]
+        # value: [batch_size, num_heads, seq_length, head_dim] -> [batch_size * num_heads, seq_length, head_dim]
+        return tuple(
+            (
+                layer_past[0].reshape(batch_size_times_num_heads, head_dim, seq_length),
+                layer_past[1].reshape(batch_size_times_num_heads, seq_length, head_dim),
+            )
+            for layer_past in past_key_value
+        )
+
+
+MPT_START_DOCSTRING = r"""
+
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`ExtendedMptConfig`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+MPT_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+            `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values[0][0].shape[2]`
+            (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
+
+            If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+            `input_ids`.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        past_key_values (`Tuple[Tuple[torch.Tensor]]` of length `config.n_layers`):
+            Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see
+            `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have
+            their past given to this model should not be passed as `input_ids` as they have already been computed.
+
+            Each element of `past_key_values` is a tuple (past_key, past_value):
+            - past_key: [batch_size * num_heads, head_dim, kv_length]
+            - past_value: [batch_size * num_heads, kv_length, head_dim]
+        attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+
+            If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see
+            `past_key_values`).
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
+        use_external_mind (`bool`, *optional*, defaults to `True`):
+            Whether to attend to external memories.
+        long_range_past_key_values (`List[Tuple[torch.FloatTensor]]`, *optional*, defaults to None):
+            Manual store for memories.
+        faiss_indexes (`Tuple[faiss.swigfaiss_avx2.IndexFlatIP]`, *optional*, defaults to None):
+            Vector store for memories.
+        topk (`int`, *optional*, defaults to `10`):
+            Number of external memories for each query token to retrieve and attend to.
+"""
+
+
+@add_start_docstrings(
+    "The bare Mpt Model transformer outputting raw hidden-states without any specific head on top.",
+    MPT_START_DOCSTRING,
+)
+class ExtendedMptModel(MptPreTrainedModel):
+    """Extended MPT Model"""
+
+    def __init__(self, config: ExtendedMptConfig):
+        super().__init__(config)
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.n_heads
+
+        # Embedding + LN Embedding
+        self.wte = nn.Embedding(config.vocab_size, self.hidden_size)
+
+        # Transformer blocks
+        self.blocks = nn.ModuleList([MptBlock(config) for _ in range(config.n_layers)])
+
+        # Final Layer Norm
+        self.norm_f = LayerNorm(self.hidden_size, eps=config.layer_norm_epsilon)
+        # backward compatibility with weights on the Hub
+        self.norm_f.bias = None
+
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+        self.mask_by_sim = config.attn_config.mask_by_sim
+        self.sim_threshold = config.attn_config.sim_threshold
+        self.topk = config.attn_config.topk
+        self.use_external_mind = config.use_external_mind
+        self.use_external_mind_by_layer = config.use_external_mind_by_layer
+
+    def get_input_embeddings(self):
+        return self.wte
+
+    def build_mpt_alibi_tensor(
+        self,
+        num_heads,
+        sequence_length,
+        sequence_length_with_past,
+        alibi_bias_max=8,
+        device=None,
+        for_ae=None,
+        topk=None,
+    ):
+        return build_mpt_alibi_tensor(
+            num_heads,
+            sequence_length,
+            sequence_length_with_past,
+            alibi_bias_max,
+            device,
+            for_ae=for_ae,
+            topk=topk,
+        )
+
+    def _prepare_attn_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_shape: Tuple[int, int],
+        past_key_values_length: int,
+    ) -> torch.BoolTensor:
+        # create causal mask
+        # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
+        if input_shape[1] + past_key_values_length != attention_mask.shape[1]:
+            raise ValueError(
+                "Attention mask shape should be (batch_size, seq_length + past_key_values_length)"
+                f" but is {attention_mask.shape} with input_ids shape {input_shape} and past length"
+                f" {past_key_values_length}."
+            )
+        combined_attention_mask = None
+        device = attention_mask.device
+        _, src_length = input_shape
+
+        if src_length > 1:
+            combined_attention_mask = _make_causal_mask(
+                input_shape,
+                device=device,
+                past_key_values_length=past_key_values_length,
+            )
+
+        # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
+        expanded_attn_mask = _expand_mask(attention_mask, tgt_length=src_length)
+        combined_attention_mask = (
+            expanded_attn_mask
+            if combined_attention_mask is None
+            else expanded_attn_mask | combined_attention_mask
+        )
+
+        return combined_attention_mask
+
+    def set_input_embeddings(self, new_embeddings: torch.Tensor):
+        self.wte = new_embeddings
+
+    @add_start_docstrings_to_model_forward(MPT_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=BaseModelOutputWithPastAndCrossAttentions,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_retrieved_memory_idx: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        use_external_mind: Optional[bool] = None,
+        long_range_past_key_values: Optional[list[Tuple[torch.FloatTensor]]] = None,
+        faiss_indexes: Tuple = None,
+        topk: int = None,
+    ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_retrieved_memory_idx = (
+            output_retrieved_memory_idx
+            if output_retrieved_memory_idx is not None
+            else False
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        use_external_mind = (
+            use_external_mind
+            if use_external_mind is not None
+            else self.use_external_mind
+        )
+        topk = topk if topk is not None else self.topk
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time"
+            )
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if past_key_values is None:
+            past_key_values = tuple([None] * len(self.blocks))
+
+        if inputs_embeds is None:
+            inputs_embeds = self.wte(input_ids)
+
+        hidden_states = inputs_embeds
+
+        presents = () if use_cache else None
+        all_self_attentions = () if output_attentions else None
+        all_hidden_states = () if output_hidden_states else None
+        all_idx = () if output_retrieved_memory_idx else None
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        # Compute alibi tensor: check build_alibi_tensor documentation
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+        if past_key_values[0] is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                (batch_size, seq_length_with_past), device=hidden_states.device
+            )
+        else:
+            attention_mask = attention_mask.to(hidden_states.device)
+
+        alibi = self.build_mpt_alibi_tensor(
+            self.num_heads,
+            self.config.max_seq_len,
+            seq_length_with_past,
+            device=hidden_states.device,
+        )
+        # EM: Alibi tensor for retrieved kvs
+        alibi_ae = self.build_mpt_alibi_tensor(
+            self.num_heads,
+            seq_length,
+            seq_length_with_past,
+            device=hidden_states.device,
+            for_ae=True,
+            topk=topk,
+        )
+
+        causal_mask = self._prepare_attn_mask(
+            attention_mask,
+            input_shape=(batch_size, seq_length),
+            past_key_values_length=past_key_values_length,
+        )
+
+        for i, (block, layer_past) in enumerate(zip(self.blocks, past_key_values)):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            long_range_past_key_value = (
+                long_range_past_key_values[i]
+                if (
+                    long_range_past_key_values is not None
+                    and self.use_external_mind_by_layer[i]
+                    and use_external_mind is True
+                )
+                else None
+            )
+            if long_range_past_key_value is not None and faiss_indexes is not None:
+                raise NotImplementedError(
+                    """Using faiss and passing key value pairs
+                    manually are mutually exclusive right now."""
+                )
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(
+                            *inputs,
+                            use_cache=use_cache,
+                            output_attentions=output_attentions,
+                        )
+
+                    return custom_forward
+
+                outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    alibi,
+                    causal_mask,
+                    layer_past,
+                )
+            else:
+                outputs = block(
+                    hidden_states,
+                    layer_past=layer_past,
+                    attention_mask=causal_mask,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                    output_retrieved_memory_idx=output_retrieved_memory_idx,
+                    position_bias=alibi,
+                    position_bias_ae=alibi_ae,
+                    topk=topk,
+                    long_range_past_key_value=long_range_past_key_value,
+                    faiss_indexes=faiss_indexes,
+                    mask_by_sim=self.mask_by_sim,
+                    sim_threshold=self.sim_threshold,
+                    current_layer=i,
+                )
+
+            hidden_states = outputs[0]
+            if use_cache is True:
+                presents = presents + (outputs[1],)
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (
+                    outputs[2 if use_cache else 1],
+                )
+            if output_retrieved_memory_idx:
+                idx = (
+                    3
+                    if (use_cache & output_attentions)
+                    else 2
+                    if (use_cache or output_attentions)
+                    else 1
+                )
+                all_idx = all_idx + (outputs[idx],)
+
+        # Add last hidden state
+        hidden_states = self.norm_f(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    presents,
+                    all_hidden_states,
+                    all_self_attentions,
+                    all_idx,
+                ]
+                if v is not None
+            )
+
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=presents,
+            hidden_states=all_hidden_states,
+            attentions=(all_self_attentions, all_idx),  # EM: Return idx of retrieved memories
+        )
+
+
+@add_start_docstrings(
+    """
+    The MPT Model transformer with a language modeling head on top (linear layer with weights tied to the input
+    embeddings).
+    """,
+    MPT_START_DOCSTRING,
+)
+class ExtendedMptForCausalLM(MptPreTrainedModel):
+    """Extended MPT for Causal LM."""
+
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: ExtendedMptConfig, external_memories=None):
+        super().__init__(config)
+        self.transformer: ExtendedMptModel = ExtendedMptModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.use_external_mind = config.use_external_mind
+        self.memory_type = config.attn_config.memory_type
+        self.memory_ids = None
+        self.memories = None
+        self.memory_device = config.attn_config.memory_device
+        self.remove_special_ids = config.attn_config.remove_special_ids
+        self.tokenizer_all_special_ids = config.attn_config.tokenizer_all_special_ids
+
+        # EM: Memory token ids
+        if external_memories is not None:
+            self.memory_ids = external_memories
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: torch.Tensor):
+        self.lm_head = new_embeddings
+
+    # EM: Clear memory cache
+    def clear_memory(self):
+        """Clear memory cache."""
+        self.memory_ids = None
+        self.memories = None
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        **kwargs,
+    ) -> dict:
+        # only last token for input_ids if past is not None
+        if past_key_values:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,  # NITS should it be layer_past?
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+                "use_external_mind": kwargs.get("use_external_mind"), # EM: Add config here
+                "topk": kwargs.get("topk"),
+            }
+        )
+        return model_inputs
+
+    @add_start_docstrings_to_model_forward(MPT_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=CausalLMOutputWithCrossAttentions,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_retrieved_memory_idx: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        use_external_mind: Optional[bool] = None,
+        topk: int = None,
+    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+        """
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        # EM: Generate key value cache once on first call
+        if (
+            self.memory_ids is not None and self.memories is None
+        ): 
+            self.memories = self.generate_cache(
+                self.memory_ids, cache_type=self.memory_type
+            )
+            # EM: Remove special tokens from memory cache
+            if self.remove_special_ids:
+                idx_to_remove = [
+                    token_idx
+                    for token_idx, token in enumerate(self.memory_ids[0])
+                    if token in self.tokenizer_all_special_ids
+                ]
+                if self.memory_type == "manual":
+                    mask = torch.ones(self.memories[0][0].size(), dtype=torch.bool)
+                    mask[:, :, idx_to_remove, :] = False
+
+                    new_size = (
+                        self.memories[0][0].size(0),
+                        self.memories[0][0].size(1),
+                        -1,
+                        self.memories[0][0].size(3),
+                    )
+                    self.memories = [
+                        (ks[mask].view(new_size), vs[mask].view(new_size))
+                        for ks, vs in self.memories
+                    ]
+                else:
+                    kn_index, kv_index = self.memories
+                    all_idx_to_remove = [
+                        [
+                            i
+                            for i in range(0, kn_index.ntotal)
+                            if (
+                                i
+                                % (
+                                    kn_index.ntotal
+                                    / (
+                                        self.config.num_attention_heads
+                                        * self.config.num_hidden_layers
+                                    )
+                                )
+                            )
+                            == j
+                        ]
+                        for j in idx_to_remove
+                    ]
+                    kn_index.remove_ids(
+                        np.array(all_idx_to_remove).flatten().astype("int64")
+                    )
+                    kv_index.remove_ids(
+                        np.array(all_idx_to_remove).flatten().astype("int64")
+                    )
+
+        use_external_mind = (
+            use_external_mind
+            if use_external_mind is not None
+            else self.use_external_mind
+        )
+        topk = topk if topk is not None else None
+
+        long_range_past_key_values = None
+        faiss_indexes = None
+        if hasattr(self, "memories") and isinstance(self.memories, list):
+            long_range_past_key_values = self.memories
+        elif hasattr(self, "memories"):
+            faiss_indexes = self.memories
+
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_retrieved_memory_idx=output_retrieved_memory_idx,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            long_range_past_key_values=long_range_past_key_values,
+            faiss_indexes=faiss_indexes,
+            use_external_mind=use_external_mind,
+            topk=topk,
+        )
+        hidden_states = transformer_outputs[0]
+
+        lm_logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(lm_logits.device)
+            # Shift so that tokens < n predict n
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            batch_size, seq_length, vocab_size = shift_logits.shape
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(batch_size * seq_length, vocab_size),
+                shift_labels.view(batch_size * seq_length),
+            )
+
+        if not return_dict:
+            output = (lm_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=lm_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )
+
+    def _reorder_cache(
+        self,
+        past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...],
+        beam_idx: torch.LongTensor,
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
+        """
+        This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
+        [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
+        beam_idx at every generation step.
+
+        Output shares the same memory storage as `past`.
+        """
+        # Get a copy of `beam_idx` on all the devices where we need those indices.
+        device_to_beam_idx = {
+            past_state.device: beam_idx.to(past_state.device)
+            for layer_past in past
+            for past_state in layer_past
+        }
+        reordered_past = tuple(
+            (
+                layer_past[0].index_select(0, device_to_beam_idx[layer_past[0].device]),
+                layer_past[1].index_select(0, device_to_beam_idx[layer_past[0].device]),
+            )
+            for layer_past in past
+        )
+        return reordered_past
+
+    # EM: Add method to generate key-value cache
+    def generate_cache(
+        self,
+        input_ids: torch.LongTensor,
+        stride: int = 512,
+        max_len: int = 3072,
+        cache_type: str = "manual",
+    ):
+        """Generate cache for long range attention."""
+        if cache_type not in ["manual", "faiss"]:
+            raise NotImplementedError(f"Cache type {cache_type} not implemented.")
+
+        prev_end_loc = 0
+        long_range_past_key_values = None
+        faiss_indexes = None
+        for b_idx in range(
+            0, input_ids.size(-1), stride
+        ):  # generate kv-pairs using stride
+            end_loc = min(b_idx + max_len, input_ids.size(-1))
+            trg_len = end_loc - prev_end_loc
+            subseq = input_ids[:, b_idx:end_loc].to(self.device)
+            with torch.no_grad():
+                outputs = self.transformer(
+                    subseq, use_cache=True, use_external_mind=False
+                )
+            to_cache = [
+                (kv[0][:, :, -trg_len:], kv[1][:, :, -trg_len:])
+                for kv in outputs.past_key_values
+            ]
+            long_range_past_key_values, faiss_indexes = self.cache(
+                to_cache,
+                cache_type,
+                long_range_past_key_values=long_range_past_key_values,
+                faiss_indexes=faiss_indexes,
+            )
+
+            prev_end_loc = end_loc
+            if end_loc == input_ids.size(-1):
+                break
+        if long_range_past_key_values is not None:
+            return long_range_past_key_values
+        else:
+            return faiss_indexes
+        
+    # EM: Add method to cache key value pairs
+    def cache(
+        self,
+        to_cache: list,
+        cache_type: str = "manual",
+        long_range_past_key_values: list = None,
+        faiss_indexes: faiss.IndexFlatIP = None,
+        max_length_cache=100000,
+        verbose=False,
+    ):
+        """Cache long range attention."""
+        if (long_range_past_key_values is not None) & (faiss_indexes is not None):
+            raise NotImplementedError(
+                "Using faiss and passing key value pairs manually are mutually exclusive right now."
+            )
+        
+        # To avoid spinning up a new index for each layer, we add one-hot encodings to the keys so that queries match with the appropriate layer, head
+        if cache_type == "faiss":  # add one-hot encoding to match layer, head indices
+            one_hot_encodings = (
+                F.one_hot(torch.arange(0, self.config.n_heads * self.config.n_layers))
+                * 10
+            )
+            # New indices, one to store normalized keys with one-hot encodings, another to retrieve kv pairs without normalization
+            if faiss_indexes is None:
+                faiss_indexes = (
+                    faiss.IndexFlatIP(
+                        to_cache[0][0].size(-1) + one_hot_encodings.size(-1)
+                    ),
+                    faiss.IndexFlatIP(to_cache[0][0].size(-1) * 2),
+                )
+            kn_index, kv_index = faiss_indexes
+            for l_idx, (k, v) in enumerate(to_cache):
+                k_n = (k / vector_norm(k, ord=2, dim=-1, keepdim=True)).to("cpu") #Normalize keys for cosine sim
+                
+                # Indices are 2 dimensional, so flatten 
+                # Add normalized keys with one-hot encodings
+                k_n = torch.concat(
+                    [
+                        rearrange(k_n, "b h s d -> b (h s) d", h=self.config.n_heads),
+                        one_hot_encodings[
+                            self.config.n_heads
+                            * l_idx : self.config.n_heads
+                            * (l_idx + 1)
+                        ]
+                        .unsqueeze(0)
+                        .repeat_interleave(repeats=k.size(-2), dim=-2),
+                    ],
+                    dim=-1,
+                )
+                kn_index.add(k_n.squeeze().numpy())
+
+                # Add unnormalized keys and values
+                k = rearrange(k, "b h s d -> b (h s) d", h=self.config.n_heads)
+                v = rearrange(v, "b h s d -> b (h s) d", h=self.config.n_heads)
+                kv_index.add(
+                    torch.concat([k.squeeze(), v.squeeze()], dim=1).to("cpu").numpy()
+                )
+        else:
+            # Simply use list to store key value pairs
+            if long_range_past_key_values is None:
+                long_range_past_key_values = [
+                    (k.to(self.memory_device), v.to(self.memory_device))
+                    for k, v in to_cache
+                ]
+            else:
+                long_range_past_key_values = [
+                    (
+                        torch.concat(
+                            [kv[0], to_cache[ind][0].to(self.memory_device)], dim=2
+                        ),
+                        torch.concat(
+                            [kv[1], to_cache[ind][1].to(self.memory_device)], dim=2
+                        ),
+                    )
+                    for ind, kv in enumerate(long_range_past_key_values)
+                ]
+        if (
+            long_range_past_key_values is not None
+        ):  # set a limit on manual memory length
+            if long_range_past_key_values[0][0].size(-2) > max_length_cache:
+                long_range_past_key_values = [
+                    (
+                        kv[0][:, :, -max_length_cache:],
+                        kv[1][:, :, -max_length_cache:],
+                    )
+                    for kv in long_range_past_key_values
+                ]
+        if verbose:
+            if cache_type == "faiss":
+                print(f"{kn_index.ntotal} keys in faiss index")
+            else:
+                print(f"{long_range_past_key_values[0][0].size(-2)} cached kvs")
+
+        return (
+            long_range_past_key_values,
+            (kn_index, kv_index) if cache_type == "faiss" else None,
+        )