Upload OLMoForCausalLM

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,56 @@
+{
+  "_name_or_path": "/home/itay.itzhak/projects/proj2/finetuning/open-instruct/output/allenai/tulu-v2-sft-mixture_allenai/OLMo-7B_lora_r128_alpha256_LR2e-5_seed_2/merged",
+  "activation_type": "swiglu",
+  "alibi": false,
+  "alibi_bias_max": 8.0,
+  "architectures": [
+    "OLMoForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "attention_layer_norm": false,
+  "attention_layer_norm_with_affine": false,
+  "auto_map": {
+    "AutoConfig": "configuration_olmo.OLMoConfig",
+    "AutoModelForCausalLM": "modeling_olmo.OLMoForCausalLM",
+    "AutoTokenizer": [
+      "allenai/OLMo-7B--tokenization_olmo_fast.OLMoTokenizerFast",
+      "allenai/OLMo-7B--tokenization_olmo_fast.OLMoTokenizerFast"
+    ]
+  },
+  "bias_for_layer_norm": false,
+  "block_group_size": 1,
+  "block_type": "sequential",
+  "clip_qkv": null,
+  "d_model": 4096,
+  "embedding_dropout": 0.0,
+  "embedding_size": 50304,
+  "eos_token_id": 50279,
+  "flash_attention": true,
+  "include_bias": false,
+  "init_cutoff_factor": null,
+  "init_device": "meta",
+  "init_fn": "mitchell",
+  "init_std": 0.02,
+  "layer_norm_eps": 1e-05,
+  "layer_norm_type": "default",
+  "layer_norm_with_affine": false,
+  "max_sequence_length": 2048,
+  "mlp_hidden_size": 22016,
+  "mlp_ratio": 4,
+  "model_type": "hf_olmo",
+  "multi_query_attention": false,
+  "n_heads": 32,
+  "n_kv_heads": null,
+  "n_layers": 32,
+  "pad_token_id": 1,
+  "precision": "amp_bf16",
+  "residual_dropout": 0.0,
+  "rope": true,
+  "rope_full_precision": true,
+  "scale_logits": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.42.4",
+  "use_cache": true,
+  "vocab_size": 50280,
+  "weight_tying": false
+}

configuration_olmo.py

@@ -0,0 +1,43 @@
+"""
+OLMo configuration
+"""
+
+from transformers import AutoConfig, PretrainedConfig
+from transformers.utils import logging
+
+from olmo.config import ModelConfig
+
+logger = logging.get_logger(__name__)
+
+
+class OLMoConfig(PretrainedConfig):
+    model_type = "hf_olmo"
+    keys_to_ignore_at_inference = ["past_key_values"]  # TODO: confirm
+
+    def __init__(self, use_cache: bool = False, **kwargs):
+        model_config = ModelConfig()
+        all_kwargs = model_config.asdict()
+        all_kwargs.update(kwargs)
+        all_kwargs.update({"use_cache": use_cache})
+        all_kwargs.update(
+            {"architectures": all_kwargs.get("architectures", ["OLMoForCausalLM"]) or ["OLMoForCausalLM"]}
+        )
+        super().__init__(**all_kwargs)
+
+    @property
+    def num_attention_heads(self):
+        return self.n_heads
+
+    @property
+    def num_hidden_layers(self):
+        return self.n_layers
+
+    @property
+    def hidden_size(self):
+        return self.d_model
+
+
+# Register the config class so that it is available for transformer pipelines, auto-loading etc.
+# OLMo is integrated directly in transformers from v4.40.0 onwards, but the version in transformers
+# may not support the newest architectures we create.
+AutoConfig.register("hf_olmo", OLMoConfig)

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 50279,
+  "pad_token_id": 1,
+  "transformers_version": "4.42.4"
+}

model-00001-of-00006.safetensors

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

model-00002-of-00006.safetensors

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

model-00003-of-00006.safetensors

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

model-00004-of-00006.safetensors

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

model-00005-of-00006.safetensors

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

model-00006-of-00006.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,137 @@
+{
+  "metadata": {
+    "total_size": 27552382976
+  },
+  "weight_map": {
+    "model.transformer.blocks.0.att_proj.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.0.attn_out.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.0.ff_out.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.0.ff_proj.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.1.att_proj.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.1.attn_out.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.1.ff_out.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.1.ff_proj.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.10.att_proj.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.10.attn_out.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.10.ff_out.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.10.ff_proj.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.11.att_proj.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.11.attn_out.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.11.ff_out.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.11.ff_proj.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.12.att_proj.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.12.attn_out.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.12.ff_out.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.12.ff_proj.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.13.att_proj.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.13.attn_out.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.13.ff_out.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.13.ff_proj.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.14.att_proj.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.14.attn_out.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.14.ff_out.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.14.ff_proj.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.15.att_proj.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.15.attn_out.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.15.ff_out.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.15.ff_proj.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.16.att_proj.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.16.attn_out.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.16.ff_out.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.16.ff_proj.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.17.att_proj.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.17.attn_out.weight": "model-00003-of-00006.safetensors",
+    "model.transformer.blocks.17.ff_out.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.17.ff_proj.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.18.att_proj.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.18.attn_out.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.18.ff_out.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.18.ff_proj.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.19.att_proj.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.19.attn_out.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.19.ff_out.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.19.ff_proj.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.2.att_proj.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.2.attn_out.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.2.ff_out.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.2.ff_proj.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.20.att_proj.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.20.attn_out.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.20.ff_out.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.20.ff_proj.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.21.att_proj.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.21.attn_out.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.21.ff_out.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.21.ff_proj.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.22.att_proj.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.22.attn_out.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.22.ff_out.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.22.ff_proj.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.23.att_proj.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.23.attn_out.weight": "model-00004-of-00006.safetensors",
+    "model.transformer.blocks.23.ff_out.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.23.ff_proj.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.24.att_proj.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.24.attn_out.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.24.ff_out.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.24.ff_proj.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.25.att_proj.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.25.attn_out.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.25.ff_out.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.25.ff_proj.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.26.att_proj.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.26.attn_out.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.26.ff_out.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.26.ff_proj.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.27.att_proj.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.27.attn_out.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.27.ff_out.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.27.ff_proj.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.28.att_proj.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.28.attn_out.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.28.ff_out.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.28.ff_proj.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.29.att_proj.weight": "model-00006-of-00006.safetensors",
+    "model.transformer.blocks.29.attn_out.weight": "model-00005-of-00006.safetensors",
+    "model.transformer.blocks.29.ff_out.weight": "model-00006-of-00006.safetensors",
+    "model.transformer.blocks.29.ff_proj.weight": "model-00006-of-00006.safetensors",
+    "model.transformer.blocks.3.att_proj.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.3.attn_out.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.3.ff_out.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.3.ff_proj.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.30.att_proj.weight": "model-00006-of-00006.safetensors",
+    "model.transformer.blocks.30.attn_out.weight": "model-00006-of-00006.safetensors",
+    "model.transformer.blocks.30.ff_out.weight": "model-00006-of-00006.safetensors",
+    "model.transformer.blocks.30.ff_proj.weight": "model-00006-of-00006.safetensors",
+    "model.transformer.blocks.31.att_proj.weight": "model-00006-of-00006.safetensors",
+    "model.transformer.blocks.31.attn_out.weight": "model-00006-of-00006.safetensors",
+    "model.transformer.blocks.31.ff_out.weight": "model-00006-of-00006.safetensors",
+    "model.transformer.blocks.31.ff_proj.weight": "model-00006-of-00006.safetensors",
+    "model.transformer.blocks.4.att_proj.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.4.attn_out.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.4.ff_out.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.4.ff_proj.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.5.att_proj.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.5.attn_out.weight": "model-00001-of-00006.safetensors",
+    "model.transformer.blocks.5.ff_out.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.5.ff_proj.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.6.att_proj.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.6.attn_out.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.6.ff_out.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.6.ff_proj.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.7.att_proj.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.7.attn_out.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.7.ff_out.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.7.ff_proj.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.8.att_proj.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.8.attn_out.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.8.ff_out.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.8.ff_proj.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.9.att_proj.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.9.attn_out.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.9.ff_out.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.blocks.9.ff_proj.weight": "model-00002-of-00006.safetensors",
+    "model.transformer.ff_out.weight": "model-00006-of-00006.safetensors",
+    "model.transformer.wte.weight": "model-00001-of-00006.safetensors"
+  }
+}

modeling_olmo.py

@@ -0,0 +1,228 @@
+import logging
+from dataclasses import fields
+from typing import List, Optional, Tuple, Union
+
+import torch
+from transformers import PreTrainedModel
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.auto import AutoModelForCausalLM
+
+from olmo.config import ModelConfig
+from olmo.model import OLMo
+
+from .configuration_olmo import OLMoConfig
+
+log = logging.getLogger(__name__)
+
+
+def create_model_config_from_pretrained_config(config: OLMoConfig):
+    """
+    Utility function
+    """
+
+    kwargs = {}
+    for field in fields(ModelConfig):
+        kwargs[field.name] = getattr(config, field.name)
+
+    model_config = ModelConfig(**kwargs)
+    return model_config
+
+
+class OLMoForCausalLM(PreTrainedModel):
+    """
+    Extremely barebones HF model wrapper.
+    """
+
+    config_class = OLMoConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["OLMoBlock"]
+
+    def __init__(self, config: OLMoConfig, model: Optional[OLMo] = None, init_params: bool = False):
+        super().__init__(config)
+
+        if not model:
+            model_config = create_model_config_from_pretrained_config(config)
+            # Initialize model (always on CPU to start with so we don't run out of GPU memory).
+            model_config.init_device = "cpu"
+            self.model = OLMo(model_config, init_params=init_params)
+        else:
+            self.model = model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        attention_bias: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[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,
+        cache_position: Optional[
+            Cache
+        ] = None,  # This is a hack mitigation of an issue in transformers `4.39.x` https://github.com/huggingface/transformers/issues/29426
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        if use_cache is None:
+            use_cache = self.config.use_cache
+
+        if output_attentions:
+            raise ValueError("output_attentions is not yet supported in OLMo")
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model.forward(
+            input_ids=input_ids,
+            input_embeddings=inputs_embeds,
+            attention_mask=attention_mask,
+            attention_bias=attention_bias,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_hidden_states=output_hidden_states,
+        )
+
+        logits = outputs.logits
+        hidden_states = outputs.hidden_states
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = torch.nn.CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.embedding_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.attn_key_values,
+            hidden_states=hidden_states,
+        )
+
+    def can_generate(self) -> bool:
+        return True
+
+    def prepare_inputs_for_generation(
+        self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple]] = None, **kwargs
+    ):
+        if past_key_values:
+            # This is because we want the model to only process the last generated token.
+            input_ids = input_ids[:, -1:]
+        model_inputs = {"input_ids": input_ids, "past_key_values": past_key_values}
+
+        model_inputs.update(kwargs)
+        model_inputs["use_cache"] = kwargs.pop("use_cache", self.config.use_cache)
+        return model_inputs
+
+    # TODO: these are required to make the implementation complete.
+    # def resize_position_embeddings(self, new_num_position_embeddings: int):
+    #     pass
+    #
+    # def get_position_embeddings(self) -> Union[nn.Embedding, Tuple[nn.Embedding]]:
+    #     pass
+    #
+    # def _reorder_cache(self, past_key_values, beam_idx):
+    #     pass
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.model.transformer.wte
+
+    def set_input_embeddings(self, value: torch.nn.Module):
+        self.model.transformer.wte = value
+
+    def get_output_embeddings(self):
+        if self.config.weight_tying:
+            return self.model.transformer.wte
+        else:
+            return self.model.transformer.ff_out
+
+    def set_output_embeddings(self, value: torch.nn.Module):
+        if self.config.weight_tying:
+            self.model.transformer.wte = value
+        else:
+            self.model.transformer.ff_out = value
+
+    def tie_weights(self):
+        """
+        This function is intentionally left as a no-op.
+
+        Weight tying is handled as follows:
+        - When the model is initialized, the `ff_out` layer is conditionally defined based on the `weight_tying` configuration.
+        See: `if not config.weight_tying: self.transformer.update(...)` in `olmo/model.py`.
+        - When computing logits, the `wte` weights are used directly if `weight_tying` is enabled.
+        See: `if self.config.weight_tying: logits = F.linear(x, self.transformer.wte.weight, None)` in the `forward` method.
+
+        Therefore, there is no need to explicitly tie the weights in this function.
+        """
+        pass
+
+    def resize_token_embeddings(
+        self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
+    ) -> torch.nn.Embedding:
+        """
+        Resizes input token embeddings matrix of the model if `new_num_tokens != config.embedding_size`.
+
+        Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method.
+
+        Arguments:
+            new_num_tokens (`int`, *optional*):
+                The new number of tokens in the embedding matrix. Increasing the size will add newly initialized
+                vectors at the end. Reducing the size will remove vectors from the end. If not provided or `None`, just
+                returns a pointer to the input tokens `torch.nn.Embedding` module of the model without doing anything.
+            pad_to_multiple_of (`int`, *optional*):
+                If set will pad the embedding matrix to a multiple of the provided value. If `new_num_tokens` is set to
+                `None` will just pad the embedding to a multiple of `pad_to_multiple_of`.
+
+                This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
+                `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. For more
+                details about this, or help on choosing the correct value for resizing, refer to this guide:
+                https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc
+
+        Return:
+            `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model.
+
+        Note:
+            This method differs from the base class implementation by resizing the `embedding_size` attribute of the
+            model configuration instead of the `vocab_size`. It also includes a warning if the resized `embedding_size`
+            is less than the `vocab_size`. In OLMo, `embedding_size` refers to the dimensionality of the model's token
+            embeddings, while `vocab_size` refers to the number of unique tokens in the vocabulary.
+        """
+        model_embeds = self._resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
+        if new_num_tokens is None and pad_to_multiple_of is None:
+            return model_embeds
+
+        # Update base model and current model config
+        self.config.embedding_size = model_embeds.weight.shape[0]
+        self.model.config.embedding_size = model_embeds.weight.shape[0]
+
+        # Check if the embedding size is less than the vocab size
+        if self.config.embedding_size < self.config.vocab_size:
+            warning_message = (
+                f"Resizing token embeddings to size {self.config.embedding_size}, which is less than the vocab size "
+                f"{self.config.vocab_size} defined in the model configuration. Make sure your tokenizer's vocabulary "
+                "size is less than or equal to the new token embedding size."
+            )
+            log.warning(warning_message)
+
+        # Tie weights again if needed
+        self.tie_weights()
+
+        return model_embeds
+
+
+# Register the model so that it is available for transformer pipelines, auto-loading, etc.
+# OLMo is integrated directly in transformers from v4.40.0 onwards, but the version in transformers
+# may not support the newest architectures we create.
+AutoModelForCausalLM.register(OLMoConfig, OLMoForCausalLM)