Upload folder using huggingface_hub

config.json

@@ -1,8 +1,28 @@
 {
+  "apply_qk_norm": true,
   "architectures": [
-    "OpenLMModel"
+    "OpenLMforCausalLM"
   ],
+  "attn_func": null,
+  "auto_map": {
+    "AutoConfig": "configuration_openlm.OpenLMConfig",
+    "AutoModel": "modeling_openlm.OpenLMModel",
+    "AutoModelForCausalLM": "modeling_openlm.OpenLMforCausalLM"
+  },
+  "dim": 2560,
+  "ffn_type": "swiglu",
   "model_type": "openlm",
+  "moe_capacity_factor": 1.25,
+  "moe_expert_model_parallelism": false,
+  "moe_freq": 0,
+  "moe_loss_weight": 0.1,
+  "moe_num_experts": null,
+  "moe_top_k": 2,
+  "moe_weight_parallelism": false,
+  "n_heads": 32,
+  "n_layers": 32,
+  "norm_eps": 1e-05,
+  "norm_type": null,
   "params": null,
   "params_args_dict": {
     "apply_qk_norm": true,
@@ -31,6 +51,11 @@
     "vocab_size": 50432,
     "weight_tying": false
   },
-  "torch_dtype": "float32",
-  "transformers_version": "4.40.0"
+  "positional_embedding_type": "rotary",
+  "post_embed_norm": false,
+  "seq_len": 2048,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.40.0",
+  "vocab_size": 50432,
+  "weight_tying": false
 }

configuration_openlm.py

@@ -0,0 +1 @@
+from open_lm.utils.transformers.hf_config import OpenLMConfig

modeling_openlm.py

@@ -0,0 +1,203 @@
+from argparse import Namespace
+from torch.utils.checkpoint import checkpoint
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from open_lm.utils.transformers.hf_config import OpenLMConfig
+from open_lm.model import Transformer, create_params
+from open_lm.attention import get_attn_func, xformers_attn, torch_attn
+from open_lm.norms import get_norm_class
+import torch
+import torch.nn as nn
+from typing import Union, Tuple, Optional, List
+import os
+
+
+class OpenLMModel(PreTrainedModel):
+    config_class = OpenLMConfig
+
+    def __init__(self, config, **kwargs):
+        # This has to be done before init as it sets makes sure hf config is correct
+        if hasattr(config, "params"):
+            params = config.params
+        else:
+            params_args_dict = config.params_args_dict
+            if not params_args_dict.get("norm_type"):
+                params_args_dict["norm_type"] = get_norm_class(params_args_dict["model_norm"])
+            if not params_args_dict.get("attn_func"):
+                params_args_dict["attn_func"] = get_attn_func(
+                    params_args_dict["attn_name"], 
+                    params_args_dict["attn_activation"],
+                    params_args_dict["attn_seq_scalar"],
+                    params_args_dict["attn_seq_scalar_alpha"]
+                    )
+            params = create_params(Namespace(**config.params_args_dict))
+        config.set_params(params)
+        super().__init__(config, **kwargs)
+
+        self.supports_gradient_checkpointing = True
+        self.model = Transformer(params)
+
+    @property
+    def gradient_checkpointing(self):
+        return self.model.grad_checkpointing
+
+    @gradient_checkpointing.setter
+    def gradient_checkpointing(self, value):
+        self.model.grad_checkpointing = value
+
+    def forward(self, input_ids=None, inputs_embeds=None, **kwargs):
+        return self.model(input_ids=input_ids, inputs_embeds=inputs_embeds, **kwargs)
+
+
+class OpenLMforCausalLM(OpenLMModel):
+    _keys_to_ignore_on_load_missing = [r"lm_head.weight"]
+
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
+        self.lm_head = None
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.tok_embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.tok_embeddings = value
+
+    def get_output_embeddings(self):
+        return self.model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        raise NotImplementedError
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        Returns:
+        Example:
+        ```python
+        >>> from transformers import AutoTokenizer, OpenLlamaForCausalLM
+        >>> model = OpenLlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+        >>> prompt = "Hey, are you consciours? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you."
+        ```"""
+        assert position_ids is None, "Position IDs are not supported"
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        logits, _, past_key_values = self.model(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            attention_mask=attention_mask,
+        )
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, shift_logits.size(-1))
+            shift_labels = shift_labels.view(-1).to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        output = CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values, loss=loss)
+        return output
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values is not None:
+            past_length = past_key_values[0][0].shape[1]
+
+            # Some generation methods already pass only the last input ID
+            if input_ids.shape[1] > past_length:
+                remove_prefix_length = past_length
+            else:
+                # Default to old behavior: keep only final ID
+                remove_prefix_length = input_ids.shape[1] - 1
+
+            input_ids = input_ids[:, remove_prefix_length:]
+
+        # 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,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_cache = ()
+        for layer_past in past_key_values:
+            reordered_cache += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
+        return reordered_cache
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
+        if (
+            os.path.isdir(pretrained_model_name_or_path)
+            and kwargs.get("config", None) is not None
+            and getattr(kwargs["config"], "checkpoint_file", None) is not None
+        ):
+            # Setting torch default dtype
+            torch_dtype = getattr(kwargs["config"], "torch_dtype", None)
+            if isinstance(torch_dtype, str):
+                torch_dtype = getattr(torch, torch_dtype)
+            if torch_dtype is not None:
+                torch.set_default_dtype(torch_dtype)
+
+            print("Loading checkpoint from directory")
+            checkpoint_path = kwargs["config"].checkpoint_file
+            checkpoint = torch.load(checkpoint_path)
+
+            state_dict = checkpoint["state_dict"]
+            state_dict = {x.replace("module.", ""): y for x, y in state_dict.items()}
+            state_dict = {f"model.{x}": y for x, y in state_dict.items()}
+
+            return super().from_pretrained(None, state_dict=state_dict, **kwargs)
+        elif os.path.isdir(pretrained_model_name_or_path):
+            # Load from a PyTorch checkpoint
+            print("Loading checkpoint from directory")
+            checkpoint_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
+            state_dict = torch.load(checkpoint_path)
+
+            # state_dict = {x.replace("module.", ""): y for x, y in state_dict.items()}
+            state_dict = {f"model.{x}" if "model."  not in x else x: y for x, y in state_dict.items()}
+
+            return super().from_pretrained(pretrained_model_name_or_path, state_dict=state_dict, **kwargs)
+        else:
+            return super().from_pretrained(pretrained_model_name_or_path, **kwargs)

pytorch_model.bin

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

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,214 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|padding|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50254": {
+      "content": "                        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50255": {
+      "content": "                       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50256": {
+      "content": "                      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50257": {
+      "content": "                     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50258": {
+      "content": "                    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50259": {
+      "content": "                   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50260": {
+      "content": "                  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50261": {
+      "content": "                 ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50262": {
+      "content": "                ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50263": {
+      "content": "               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50264": {
+      "content": "              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50265": {
+      "content": "             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50266": {
+      "content": "            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50267": {
+      "content": "           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50268": {
+      "content": "          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50269": {
+      "content": "         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50270": {
+      "content": "        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50271": {
+      "content": "       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50272": {
+      "content": "      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50273": {
+      "content": "     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50274": {
+      "content": "    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50275": {
+      "content": "   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50276": {
+      "content": "  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": null,
+  "tokenizer_class": "GPTNeoXTokenizer",
+  "unk_token": "<|endoftext|>"
+}