Upload folder using huggingface_hub

README.md

@@ -0,0 +1,94 @@
+---
+datasets:
+- tiiuae/falcon-refinedweb
+- togethercomputer/RedPajama-Data-1T
+- CarperAI/pilev2-dev
+- bigcode/starcoderdata
+- JeanKaddour/minipile
+language:
+- en
+tags:
+- causal-lm
+license: cc-by-sa-4.0
+---
+# `StableLM-Base-Alpha-7B-v2`
+
+## Model Description
+
+`StableLM-Base-Alpha-7B-v2` is a 7 billion parameter decoder-only language model pre-trained on diverse English datasets. This model is the successor to the first [`StableLM-Base-Alpha-7B`](https://huggingface.co/stabilityai/stablelm-base-alpha-7b) model, addressing previous shortcomings through the use of improved data sources and mixture ratios.
+
+## Usage
+
+Get started generating text with `StableLM-Base-Alpha-7B-v2` by using the following code snippet:
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-base-alpha-7b-v2")
+model = AutoModelForCausalLM.from_pretrained(
+  "stabilityai/stablelm-base-alpha-7b-v2",
+  trust_remote_code=True,
+  torch_dtype="auto",
+)
+model.cuda()
+inputs = tokenizer("The weather is always wonderful", return_tensors="pt").to("cuda")
+tokens = model.generate(
+  **inputs,
+  max_new_tokens=64,
+  temperature=0.75,
+  top_p=0.95,
+  do_sample=True,
+)
+print(tokenizer.decode(tokens[0], skip_special_tokens=True))
+```
+
+## Model Details
+
+* **Developed by**: [Stability AI](https://stability.ai/)
+* **Model type**: `StableLM-Base-Alpha-v2` models are auto-regressive language models based on the transformer decoder architecture.
+* **Language(s)**: English
+* **Library**: [GPT-NeoX](https://github.com/EleutherAI/gpt-neox)
+* **License**: Model checkpoints are licensed under the Creative Commons license ([CC BY-SA-4.0](https://creativecommons.org/licenses/by-sa/4.0/)). Under this license, you must give [credit](https://creativecommons.org/licenses/by/4.0/#) to Stability AI, provide a link to the license, and [indicate if changes were made](https://creativecommons.org/licenses/by/4.0/#). You may do so in any reasonable manner, but not in any way that suggests the Stability AI endorses you or your use.
+* **Contact**: For questions and comments about the model, please email `[email protected]`
+
+### Model Architecture
+
+| Parameters     | Hidden Size | Layers | Heads | Sequence Length |
+|----------------|-------------|--------|-------|-----------------|
+| 6,890,209,280  | 4096        | 32     | 32    | 4096            |
+
+The model is a decoder-only transformer similar to the `StableLM-Base-Alpha` (v1) with the following configurations:
+
+* **Activation**: SwiGLU ([Shazeer, 2020](https://arxiv.org/abs/2002.05202))
+* **Decoder Layer**: Parallel Attention and MLP residuals with a single input LayerNorm ([Wang & Komatsuzaki, 2021](https://github.com/kingoflolz/mesh-transformer-jax/tree/master))
+* **Position Embeddings**: Rotary Position Embeddings ([Su et al., 2021](https://arxiv.org/abs/2104.09864))
+* **Bias**: LayerNorm bias terms only
+
+## Training
+
+`StableLM-Base-Alpha-7B-v2` is pre-trained using a multi-stage context length extension schedule following similar work ([Nijkamp et al. 2023](https://blog.salesforceairesearch.com/xgen/)); first pre-training at a context length of 2048 for 1 trillion tokens, then fine-tuning at a context length of 4096 for another 100B tokens.
+
+### Training Dataset
+
+The first pre-training stage relies on 1 trillion tokens sourced from a mix of the public Falcon RefinedWeb extract ([Penedo et al., 2023](https://huggingface.co/datasets/tiiuae/falcon-refinedweb)), RedPajama-Data ([Together Computer 2023](https://github.com/togethercomputer/RedPajama-Data), The Pile ([Gao et al., 2020](https://arxiv.org/abs/2101.00027)), and internal datasets with web text sampled at a rate of 71%.
+
+In the second stage, we include the StarCoder ([Li et al., 2023](https://arxiv.org/abs/2305.06161)) dataset and down sample web text to 55% while increasing sampling proportions of naturally long text examples in the aforementioned sources.
+
+### Training Procedure
+
+The model is pre-trained on the dataset mixes mentioned above in mixed-precision (FP16), optimized with AdamW, and trained using the NeoX tokenizer with a vocabulary size of 50,257. We outline the complete hyperparameters choices in the project's [GitHub repository - config](https://github.com/Stability-AI/StableLM/blob/main/configs/stablelm-base-alpha-7b-v2.yaml).
+
+### Training Infrastructure
+
+* **Hardware**: `StableLM-Base-Alpha-7B-v2` was trained on the Stability AI cluster - occupying 384 NVIDIA A100 40GB GPUs across AWS P4d instances. Training took approximately 16.33 days to complete across both stages.
+
+* **Software**: We use a fork of gpt-neox ([EleutherAI, 2021](https://github.com/EleutherAI/gpt-neox)) and train under 2D parallelism (Data and Tensor Parallel) with ZeRO-1 ([Rajbhandari et al., 2019](https://arxiv.org/abs/1910.02054v3)) and rely on flash-attention as well as rotary embedding kernels from FlashAttention-2 ([Dao et al., 2023](https://tridao.me/publications/flash2/flash2.pdf))
+
+## Use and Limitations
+
+### Intended Use
+
+These models are intended to be used by all individuals as foundational models for application-specific fine-tuning without strict limitations on commercial use.
+
+### Limitations and bias
+
+The pre-training dataset may have contained offensive or inappropriate content even after applying data cleansing filters which can be reflected in the model-generated text. We recommend that users exercise caution when using these models in production systems. Do not use the models for any applications that may cause harm or distress to individuals or groups.

config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "StableLMAlphaForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_stablelm_alpha.StableLMAlphaConfig",
+    "AutoModelForCausalLM": "modeling_stablelm_alpha.StableLMAlphaForCausalLM"
+  },
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "max_position_embeddings": 4096,
+  "model_type": "stablelm_alpha",
+  "norm_eps": 1e-05,
+  "num_heads": 32,
+  "num_hidden_layers": 32,
+  "rotary_emb_base": 10000,
+  "rotary_pct": 0.25,
+  "rotary_scaling_factor": 1.0,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.30.2",
+  "use_cache": true,
+  "vocab_size": 50432
+}

configuration_stablelm_alpha.py

@@ -0,0 +1,113 @@
+# coding=utf-8
+# Copyright 2023 Stability and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+""" StableLM β model configuration"""
+
+from transformers import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+STABLE_LM_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+
+
+class StableLMAlphaConfig(PretrainedConfig):
+    r"""
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 50432):
+            Vocabulary size of the StableLM model. Defines the number of different tokens that
+            can be represented by the `inputs_ids` passed when calling [`StableLMAlphaModel`].
+        hidden_size (`int`, *optional*, defaults to 6144):
+            Dimension of the decoder layers and the pooler layer.
+        num_hidden_layers (`int`, *optional*, defaults to 44):
+            Number of hidden layers in the Transformer decoder.
+        num_heads (`int`, *optional*, defaults to 64):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string).
+        rotary_pct (`float`, *optional*, defaults to 0.25):
+            Percentage of hidden dimensions to allocate to rotary embeddings.
+        rotary_emb_base (`int`, *optional*, defaults to 10000)
+            Base for computing rotary embeddings frequency.
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used with.
+            Typically set this to something large just in case (e.g., 512 or 1024 or 2048).
+        initializer_range (`float`, *optional*, defaults to 1e-5):
+            The standard deviation of the truncated_normal_initializer for initializing
+             all weight matrices.
+        norm_eps (`float`, *optional*, defaults to 1e-5):
+            The epsilon used by the normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions
+            (not used by all models). Only relevant if `config.is_decoder=True`.
+        tie_word_embeddings(`bool`, *optional*, defaults to `False`):
+            Whether to tie weight embeddings
+
+    Example:
+
+    ```python
+    >>> from transformers import StableLMAlphaConfig, StableLMAlphaModel
+
+    >>> # Initializing a StableLMAlphaConfig style configuration
+    >>> configuration = StableLMAlphaConfig()
+
+    >>> # Initializing a model (with random weights) from the style configuration
+    >>> model = StableLMAlphaModel(configuration)  # doctest: +SKIP
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config  # doctest: +SKIP
+    ```"""
+    model_type = "stablelm_alpha"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=50_432,
+        hidden_size=2_560,
+        num_hidden_layers=32,
+        num_heads=32,
+        hidden_act="silu",
+        rotary_pct=0.25,
+        rotary_emb_base=10_000,
+        max_position_embeddings=2_048,
+        initializer_range=0.02,
+        norm_eps=1e-5,
+        use_cache=True,
+        bos_token_id=0,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_heads = num_heads
+        self.hidden_act = hidden_act
+        self.rotary_pct = rotary_pct
+        self.rotary_emb_base = rotary_emb_base
+        self.initializer_range = initializer_range
+        self.norm_eps = norm_eps
+        self.use_cache = use_cache
+        self.tie_word_embeddings = tie_word_embeddings
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "transformers_version": "4.30.2"
+}

model-00001-of-00003.safetensors

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

model-00002-of-00003.safetensors

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

model-00003-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0943c68ece0fffa4fc56afaf6980e33c505d36878b6e3ca9741114066bfb5f89
+size 7843649576

model.safetensors.index.json

@@ -0,0 +1,203 @@
+{
+  "metadata": {
+    "total_size": 27557658624
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00003-of-00003.safetensors",
+    "transformer.embed.weight": "model-00001-of-00003.safetensors",
+    "transformer.final_norm.bias": "model-00003-of-00003.safetensors",
+    "transformer.final_norm.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.0.attention.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.0.attention.qkv_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.0.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.0.mlp.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.0.norm.bias": "model-00001-of-00003.safetensors",
+    "transformer.layers.0.norm.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.1.attention.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.1.attention.qkv_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.1.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.1.mlp.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.1.norm.bias": "model-00001-of-00003.safetensors",
+    "transformer.layers.1.norm.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.10.attention.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.10.attention.qkv_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.10.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.10.mlp.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.10.norm.bias": "model-00001-of-00003.safetensors",
+    "transformer.layers.10.norm.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.11.attention.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.11.attention.qkv_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.11.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.11.mlp.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.11.norm.bias": "model-00001-of-00003.safetensors",
+    "transformer.layers.11.norm.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.12.attention.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.12.attention.qkv_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.12.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.12.mlp.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.12.norm.bias": "model-00002-of-00003.safetensors",
+    "transformer.layers.12.norm.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.13.attention.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.13.attention.qkv_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.13.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.13.mlp.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.13.norm.bias": "model-00002-of-00003.safetensors",
+    "transformer.layers.13.norm.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.14.attention.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.14.attention.qkv_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.14.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.14.mlp.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.14.norm.bias": "model-00002-of-00003.safetensors",
+    "transformer.layers.14.norm.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.15.attention.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.15.attention.qkv_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.15.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.15.mlp.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.15.norm.bias": "model-00002-of-00003.safetensors",
+    "transformer.layers.15.norm.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.16.attention.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.16.attention.qkv_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.16.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.16.mlp.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.16.norm.bias": "model-00002-of-00003.safetensors",
+    "transformer.layers.16.norm.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.17.attention.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.17.attention.qkv_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.17.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.17.mlp.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.17.norm.bias": "model-00002-of-00003.safetensors",
+    "transformer.layers.17.norm.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.18.attention.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.18.attention.qkv_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.18.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.18.mlp.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.18.norm.bias": "model-00002-of-00003.safetensors",
+    "transformer.layers.18.norm.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.19.attention.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.19.attention.qkv_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.19.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.19.mlp.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.19.norm.bias": "model-00002-of-00003.safetensors",
+    "transformer.layers.19.norm.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.2.attention.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.2.attention.qkv_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.2.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.2.mlp.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.2.norm.bias": "model-00001-of-00003.safetensors",
+    "transformer.layers.2.norm.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.20.attention.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.20.attention.qkv_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.20.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.20.mlp.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.20.norm.bias": "model-00002-of-00003.safetensors",
+    "transformer.layers.20.norm.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.21.attention.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.21.attention.qkv_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.21.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.21.mlp.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.21.norm.bias": "model-00002-of-00003.safetensors",
+    "transformer.layers.21.norm.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.22.attention.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.22.attention.qkv_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.22.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.22.mlp.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.22.norm.bias": "model-00002-of-00003.safetensors",
+    "transformer.layers.22.norm.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.23.attention.out_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.23.attention.qkv_proj.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.23.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.23.mlp.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.23.norm.bias": "model-00002-of-00003.safetensors",
+    "transformer.layers.23.norm.weight": "model-00002-of-00003.safetensors",
+    "transformer.layers.24.attention.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.24.attention.qkv_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.24.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.24.mlp.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.24.norm.bias": "model-00003-of-00003.safetensors",
+    "transformer.layers.24.norm.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.25.attention.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.25.attention.qkv_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.25.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.25.mlp.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.25.norm.bias": "model-00003-of-00003.safetensors",
+    "transformer.layers.25.norm.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.26.attention.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.26.attention.qkv_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.26.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.26.mlp.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.26.norm.bias": "model-00003-of-00003.safetensors",
+    "transformer.layers.26.norm.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.27.attention.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.27.attention.qkv_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.27.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.27.mlp.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.27.norm.bias": "model-00003-of-00003.safetensors",
+    "transformer.layers.27.norm.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.28.attention.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.28.attention.qkv_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.28.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.28.mlp.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.28.norm.bias": "model-00003-of-00003.safetensors",
+    "transformer.layers.28.norm.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.29.attention.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.29.attention.qkv_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.29.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.29.mlp.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.29.norm.bias": "model-00003-of-00003.safetensors",
+    "transformer.layers.29.norm.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.3.attention.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.3.attention.qkv_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.3.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.3.mlp.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.3.norm.bias": "model-00001-of-00003.safetensors",
+    "transformer.layers.3.norm.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.30.attention.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.30.attention.qkv_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.30.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.30.mlp.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.30.norm.bias": "model-00003-of-00003.safetensors",
+    "transformer.layers.30.norm.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.31.attention.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.31.attention.qkv_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.31.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.31.mlp.out_proj.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.31.norm.bias": "model-00003-of-00003.safetensors",
+    "transformer.layers.31.norm.weight": "model-00003-of-00003.safetensors",
+    "transformer.layers.4.attention.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.4.attention.qkv_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.4.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.4.mlp.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.4.norm.bias": "model-00001-of-00003.safetensors",
+    "transformer.layers.4.norm.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.5.attention.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.5.attention.qkv_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.5.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.5.mlp.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.5.norm.bias": "model-00001-of-00003.safetensors",
+    "transformer.layers.5.norm.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.6.attention.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.6.attention.qkv_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.6.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.6.mlp.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.6.norm.bias": "model-00001-of-00003.safetensors",
+    "transformer.layers.6.norm.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.7.attention.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.7.attention.qkv_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.7.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.7.mlp.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.7.norm.bias": "model-00001-of-00003.safetensors",
+    "transformer.layers.7.norm.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.8.attention.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.8.attention.qkv_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.8.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.8.mlp.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.8.norm.bias": "model-00001-of-00003.safetensors",
+    "transformer.layers.8.norm.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.9.attention.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.9.attention.qkv_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.9.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.9.mlp.out_proj.weight": "model-00001-of-00003.safetensors",
+    "transformer.layers.9.norm.bias": "model-00001-of-00003.safetensors",
+    "transformer.layers.9.norm.weight": "model-00001-of-00003.safetensors"
+  }
+}

modeling_stablelm_alpha.py

@@ -0,0 +1,656 @@
+# coding=utf-8
+# Copyright 2023 Stability AI, EleutherAI, and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 is based off the following work:
+# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
+# https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py
+""" PyTorch StableLM-Alpha model. """
+from typing import Optional, Tuple, Union
+import math
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from .configuration_stablelm_alpha import StableLMAlphaConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
+    """Expands attention_mask from `[batch_size, seq_len]` to `[batch_size, 1, tgt_seq_len, src_seq_len]`."""
+    batch_size, src_len = mask.size()
+    tgt_len = tgt_len if tgt_len is not None else src_len
+
+    expanded_mask = mask[:, None, None, :].expand(batch_size, 1, tgt_len, src_len).to(dtype)
+    inverted_mask = 1.0 - expanded_mask
+
+    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
+
+
+class LayerNorm(nn.LayerNorm):
+    def __init__(self, normalized_shape: torch.Size, bias: bool = True, **kwargs):
+        r"""
+        use_cache (`bool`, default = True): whether to use the bias term.
+        """
+        super().__init__(normalized_shape, **kwargs)
+        if not bias:
+            self.bias = None
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config: StableLMAlphaConfig):
+        super().__init__()
+
+        self.norm = LayerNorm(config.hidden_size, eps=config.norm_eps)
+        self.attention = Attention(config)
+        self.mlp = MLP(config)
+
+    def forward(
+        self,
+        hidden_states: Optional[torch.FloatTensor],
+        attention_mask: Optional[torch.FloatTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+    ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
+        residual = hidden_states
+
+        # Pre-Norm
+        hidden_states = self.norm(hidden_states)
+
+        # Self-Attention
+        attn_output, attn_weights, present_key_value = self.attention(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+        )
+
+        # Feed-forward
+        mlp_output = self.mlp(hidden_states)
+
+        hidden_states = residual + attn_output + mlp_output
+
+        outputs = (hidden_states,)
+        if output_attentions:
+            outputs += (attn_weights,)
+        if use_cache:
+            outputs += (present_key_value,)
+        return outputs  # hidden_states, (optional: attn_weights), (optional: present_key_value)
+
+
+class MLP(nn.Module):
+    def __init__(self, config: StableLMAlphaConfig):
+        super().__init__()
+
+        hidden_size = config.hidden_size
+        multiple_of = 256
+        ff_dim = int(8 * hidden_size / 3)
+        intermediate_size = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of)
+
+        self.gate_proj = torch.nn.Linear(hidden_size, 2 * intermediate_size, bias=False)
+        self.out_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+        self.act = nn.SiLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        ff, ff_gate = self.gate_proj(x).chunk(2, dim=-1)
+        return self.out_proj(ff * self.act(ff_gate))
+
+
+class RotaryEmbedding(torch.nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        max_position_embeddings: int,
+        base: int = 10_000,
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # Build here to make `torch.jit.trace` work.
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
+        )
+
+    def _set_cos_sin_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
+
+    def forward(self, x: torch.Tensor, seq_len: Optional[int] = None):
+        # x: [batch_size, num_heads, seq_len, head_size]
+        if seq_len > self.max_seq_len_cached:
+            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=torch.get_default_dtype())
+        return (
+            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
+            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
+        )
+
+
+def rotate_half(x: torch.Tensor):
+    """Rotates half the hidden dims of the input."""
+    x1, x2 = torch.chunk(x, 2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
+    # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
+    cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
+    sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
+    cos = cos[position_ids].unsqueeze(1)  # [batch_size, 1, seq_len, dim]
+    sin = sin[position_ids].unsqueeze(1)  # [batch_size, 1, seq_len, dim]
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class Attention(nn.Module):
+    def __init__(self, config: StableLMAlphaConfig):
+        super().__init__()
+
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        if self.hidden_size % self.num_heads != 0:
+            raise ValueError(
+                "`hidden_size` is not divisble by the number of attention heads! Make sure to update them"
+            )
+
+        self.qkv_proj = nn.Linear(self.hidden_size, 3 * self.hidden_size, bias=False)
+        self.out_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self._init_rope()
+
+    def _init_rope(self):
+        self.rotary_ndims = int(self.head_dim * self.config.rotary_pct)
+        self.rotary_emb = RotaryEmbedding(
+            self.rotary_ndims,
+            max_position_embeddings=self.config.max_position_embeddings,
+            base=self.config.rotary_emb_base,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        attention_mask: torch.FloatTensor,
+        position_ids: torch.LongTensor,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        has_past_key_value = past_key_value is not None
+
+        # Compute QKV
+        # [batch_size, seq_len, (num_heads * 3 * head_dim)]
+        qkv = self.qkv_proj(hidden_states)
+
+        # [batch_size, seq_len, num_heads, 3 * head_dim]
+        new_qkv_shape = qkv.size()[:-1] + (self.num_heads, 3 * self.head_dim)
+        qkv = qkv.view(*new_qkv_shape)
+
+        # 3 * [batch_size, num_heads, seq_len, head_dim]
+        query = qkv[..., : self.head_dim].permute(0, 2, 1, 3)
+        key = qkv[..., self.head_dim:(2 * self.head_dim)].permute(0, 2, 1, 3)
+        value = qkv[..., (2 * self.head_dim):].permute(0, 2, 1, 3)
+
+        # Compute rotary embeddings on rotary_ndims
+        # [batch_size, num_heads, seq_len, rotary_ndims]
+        query_rot = query[..., :self.rotary_ndims]
+        query_pass = query[..., self.rotary_ndims:]
+        key_rot = key[..., :self.rotary_ndims]
+        key_pass = key[..., self.rotary_ndims:]
+
+        # Compute token offset for rotary embeddings (when decoding)
+        kv_seq_len = key.shape[-2]
+        if has_past_key_value:
+            kv_seq_len += past_key_value[0].shape[-2]
+
+        # Add rotary embeddings to query and key
+        cos, sin = self.rotary_emb(value, seq_len=kv_seq_len)
+        query, key = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
+
+        # Concatenate rotary embeddings with pass-through query and key
+        # [batch_size, num_heads, seq_len, head_dim]
+        query = torch.cat((query, query_pass), dim=-1)
+        key = torch.cat((key, key_pass), dim=-1)
+
+        # Reuse past key-value states
+        if has_past_key_value:
+            key = torch.cat((past_key_value[0], key), dim=2)
+            value = torch.cat((past_key_value[1], value), dim=2)
+        present_key_value = (key, value) if use_cache else None
+
+        # [batch_size, num_heads, seq_len, head_dim]
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+
+        # Compute attention
+        softmax_scale = 1 / math.sqrt(self.head_dim)
+        attn_scores = torch.einsum('bthd,bshd->bhts', query, key * softmax_scale)
+        # Apply the attention mask
+        if attention_mask is not None:
+            attn_scores = attn_scores + attention_mask
+        attn_weights = nn.functional.softmax(attn_scores, dim=-1, dtype=torch.float32).to(query.dtype)
+        attn_output = torch.einsum('bhts,bshd->bthd', attn_weights, value)
+
+        # Merge heads
+        attn_output = attn_output.reshape(attn_output.shape[0], attn_output.shape[1], -1)
+
+        # Final linear projection
+        attn_output = self.out_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, present_key_value
+
+
+def attention_mask_func(attention_scores: torch.Tensor, ltor_mask: torch.Tensor):
+    attention_scores.masked_fill_(~ltor_mask, torch.finfo(attention_scores.dtype).min)
+    return attention_scores
+
+
+class StableLMAlphaPreTrainedModel(PreTrainedModel):
+    """An abstract class to handle weights initialization and a simple interface
+    for downloading and loading pretrained models.
+    """
+
+    config_class = StableLMAlphaConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["DecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+
+    def _init_weights(self, module: nn.Module):
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            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, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def _set_gradient_checkpointing(self, module: nn.Module, value=False):
+        if isinstance(module, StableLMAlphaModel):
+            module.gradient_checkpointing = value
+
+
+def _make_causal_mask(
+    input_ids_shape: torch.Size,
+    dtype: torch.dtype,
+    device: torch.device,
+    past_key_values_length: int = 0
+):
+    """Make causal mask used for bi-directional self-attention."""
+    batch_size, tgt_len = input_ids_shape
+    mask = torch.full((tgt_len, tgt_len), torch.finfo(torch.float16).min, device=device)
+    mask_cond = torch.arange(mask.size(-1), device=device)
+    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
+    mask = mask.to(dtype)
+    if past_key_values_length > 0:
+        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
+    return mask[None, None, :, :].expand(batch_size, 1, tgt_len, tgt_len + past_key_values_length)
+
+
+class StableLMAlphaModel(StableLMAlphaPreTrainedModel):
+    def __init__(self, config: StableLMAlphaConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.final_norm = LayerNorm(config.hidden_size, eps=config.norm_eps)
+
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed
+
+    def set_input_embeddings(self, value: nn.Module):
+        self.embed = value
+
+    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
+    def _prepare_decoder_attention_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_shape: torch.Size,
+        inputs_embeds: torch.Tensor,
+        past_key_values_length: int,
+    ):
+        # Create causal mask
+        # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
+        combined_attention_mask = None
+        if input_shape[-1] > 1:
+            combined_attention_mask = _make_causal_mask(
+                input_shape,
+                inputs_embeds.dtype,
+                device=inputs_embeds.device,
+                past_key_values_length=past_key_values_length,
+            )
+
+        if attention_mask is not None:
+            # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
+            expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
+                inputs_embeds.device
+            )
+            combined_attention_mask = (
+                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+            )
+
+        return combined_attention_mask
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        r"""
+        past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers`
+            with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks.
+            Can be used to speed up decoding. If `past_key_values` are used, the user
+            can optionally input only the last `decoder_input_ids` (those that don't
+            have their past key value states given to this model) of shape `(batch_size, 1)`
+            instead of all `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+        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 = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        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:
+            input_shape = input_ids.size()
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        batch_size, seq_length = input_shape
+
+        if past_key_values is None:
+            past_key_values_length = 0
+            past_key_values = tuple([None] * self.config.num_hidden_layers)
+            seq_length_with_past = seq_length
+        else:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length + past_key_values_length
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device)
+            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+        else:
+            position_ids = position_ids.view(-1, seq_length).long()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed(input_ids)
+
+        # Attention mask.
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
+            )
+        attention_mask = self._prepare_decoder_attention_mask(
+            attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+        )
+
+        hidden_states = inputs_embeds
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+        present_key_values = () if use_cache else None
+
+        for _, (decoder_layer, past_key_value) in enumerate(zip(self.layers, past_key_values)):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # `None` for `use_cache`
+                        return module(*inputs, output_attentions, None)
+
+                    return custom_forward
+
+                outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(decoder_layer),
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    # `None` for `past_key_value`
+                    None,
+                )
+            else:
+                outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = outputs[0]
+
+            if output_attentions:
+                all_attentions = all_attentions + (outputs[1],)
+
+            if use_cache:
+                present_key_values += (outputs[2 if output_attentions else 1],)
+
+        hidden_states = self.final_norm(hidden_states)
+
+        # Add last hidden state
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        present_key_values = present_key_values if use_cache else None
+        if not return_dict:
+            return tuple(v for v in [hidden_states, present_key_values, all_hidden_states, all_attentions] if v is not None)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=present_key_values,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+
+class StableLMAlphaForCausalLM(StableLMAlphaPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: StableLMAlphaConfig):
+        super().__init__(config)
+
+        self.transformer = StableLMAlphaModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Module):
+        self.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, StableLMAlphaForCausalLM, StableLMAlphaConfig
+        >>> import torch
+
+        >>> tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-base-alpha-3b-v2", trust_remote_code=True)
+        >>> config = StableLMAlphaConfig.from_pretrained("stabilityai/stablelm-base-alpha-3b-v2")
+        >>> config.is_decoder = True
+        >>> model = StableLMAlphaForCausalLM.from_pretrained("stabilityai/stablelm-base-alpha-3b-v2", config=config)
+
+        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+        >>> outputs = model(**inputs)
+
+        >>> logits = outputs.logits
+        ```
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.transformer(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        lm_loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
+            # we are doing next-token prediction; shift prediction scores and input ids by one
+            shift_logits = logits[:, :-1, :].contiguous()
+            labels = labels[:, 1:].contiguous()
+            loss_fct = CrossEntropyLoss()
+            lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((lm_loss,) + output) if lm_loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=lm_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        **kwargs
+    ):
+        # Cut decoder_input_ids if past is used
+        if past_key_values and past_key_values[0] is not None:
+            input_ids = input_ids[:, -1:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # Create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_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(
+            {
+                "attention_mask": attention_mask,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+        )
+
+        return model_inputs
+
+    def _reorder_cache(self, past_key_values: torch.Tensor, beam_idx: int):
+        reordered_past = ()
+        for past_key_value in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx) for past_state in past_key_value[:2]) + past_key_value[2:],
+            )
+        return reordered_past
+
+
+StableLMAlphaConfig.register_for_auto_class()
+StableLMAlphaForCausalLM.register_for_auto_class("AutoModelForCausalLM")

special_tokens_map.json

@@ -0,0 +1,5 @@
+{
+    "bos_token": "<|endoftext|>",
+    "eos_token": "<|endoftext|>",
+    "unk_token": "<|endoftext|>"
+}

tokenizer_config.json

@@ -0,0 +1,8 @@
+{
+  "tokenizer_class": "GPTNeoXTokenizer",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "model_max_length": 1000000000000000019884624838656
+}