Upload folder using huggingface_hub

README.md

@@ -30,23 +30,25 @@ trained with a knowledge cutoff of **January 2017**, from the
 ## Usage
 
 ```python
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 model = AutoModelForCausalLM.from_pretrained(
     "dogtooth/open-lm-1b-201701",
     dtype=torch.bfloat16,
     device_map="auto",
+    trust_remote_code=True,
 )
+tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
 ```
 
 ## Conversion Notes
 
-- Converted from the original Open LM `.pt` checkpoint to HuggingFace `LlamaForCausalLM` format.
+- Converted from the original Open LM `.pt` checkpoint to a custom `OpenLMForCausalLM` format.
+- Uses **LayerNorm** (not RMSNorm) to match the original Open LM training.
+- Includes **QK norm** (LayerNorm on Q and K projections before attention).
 - Architecture dimensions are auto-detected from checkpoint weights.
-- The original model uses **QK norm** (RMSNorm on Q and K projections), which is not natively
-  supported by HF Llama. QK norm weights are dropped during conversion. For exact numerical
-  equivalence, use the [open_lm](https://github.com/mlfoundations/open_lm) library.
+- Requires `trust_remote_code=True` when loading.
 
 ## Citation
 

config.json

@@ -1,9 +1,13 @@
 {
   "architectures": [
-    "LlamaForCausalLM"
+    "OpenLMForCausalLM"
   ],
   "attention_bias": false,
   "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "modeling_open_lm_hf.OpenLMConfig",
+    "AutoModelForCausalLM": "modeling_open_lm_hf.OpenLMForCausalLM"
+  },
   "bos_token_id": 1,
   "dtype": "float32",
   "eos_token_id": 2,
@@ -14,11 +18,12 @@
   "intermediate_size": 5632,
   "max_position_embeddings": 2048,
   "mlp_bias": false,
-  "model_type": "llama",
+  "model_type": "open_lm",
   "num_attention_heads": 16,
   "num_hidden_layers": 24,
   "num_key_value_heads": 16,
   "pretraining_tp": 1,
+  "qk_norm": true,
   "rms_norm_eps": 1e-05,
   "rope_scaling": null,
   "rope_theta": 10000.0,

model-00001-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:382e04d8fb9e3dd2d511c2bc0d56216bfe53763e46d9ab13de9d2aaaf7014aa6
-size 4985313960
+oid sha256:013ff505ba55eaa060105291399de41b923fd62cc5ee28c7f27c1c72c21762bd
+size 4985679304

model-00002-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d01d139603008cf5eb4f1df2c1276eecb4925885b8c956888801e33c06025072
-size 773891960
+oid sha256:d78d215b3f01a75a49518123fc5aa2322bd94e4c56549234ec584bcaaaa3e1eb
+size 773925168

model.safetensors.index.json

@@ -1,7 +1,7 @@
 {
   "metadata": {
-    "total_parameters": 1439795200,
-    "total_size": 5759180800
+    "total_parameters": 1439893504,
+    "total_size": 5759574016
   },
   "weight_map": {
     "lm_head.weight": "model-00002-of-00002.safetensors",
@@ -11,8 +11,10 @@
     "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.0.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.1.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -20,8 +22,10 @@
     "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.1.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.10.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -29,8 +33,10 @@
     "model.layers.10.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.10.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.10.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.10.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.10.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.10.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.10.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.11.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -38,8 +44,10 @@
     "model.layers.11.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.11.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.11.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.11.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.11.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.11.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.11.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.12.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -47,8 +55,10 @@
     "model.layers.12.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.12.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.12.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.12.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.12.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.12.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.12.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.13.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -56,8 +66,10 @@
     "model.layers.13.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.13.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.13.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.13.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.13.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.13.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.13.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.14.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -65,8 +77,10 @@
     "model.layers.14.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.14.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.14.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.14.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.14.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.14.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.14.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.15.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -74,8 +88,10 @@
     "model.layers.15.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.15.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.15.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.15.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.15.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.15.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.15.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.16.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -83,8 +99,10 @@
     "model.layers.16.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.16.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.16.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.16.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.16.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.16.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.16.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.17.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -92,8 +110,10 @@
     "model.layers.17.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.17.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.17.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.17.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.17.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.17.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.17.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.18.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -101,8 +121,10 @@
     "model.layers.18.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.18.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.18.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.18.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.18.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.18.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.18.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.19.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -110,8 +132,10 @@
     "model.layers.19.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.19.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.19.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.19.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.19.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.19.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.19.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.2.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -119,8 +143,10 @@
     "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.2.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.20.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -128,8 +154,10 @@
     "model.layers.20.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.20.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.20.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.20.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.20.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.20.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.20.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.21.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -137,8 +165,10 @@
     "model.layers.21.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.21.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.21.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.21.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.21.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.21.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.21.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.22.input_layernorm.weight": "model-00002-of-00002.safetensors",
@@ -146,8 +176,10 @@
     "model.layers.22.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
     "model.layers.22.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
     "model.layers.22.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.22.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
     "model.layers.22.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.22.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.22.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
     "model.layers.22.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.22.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.23.input_layernorm.weight": "model-00002-of-00002.safetensors",
@@ -155,8 +187,10 @@
     "model.layers.23.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
     "model.layers.23.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
     "model.layers.23.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.23.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
     "model.layers.23.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
     "model.layers.23.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.23.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
     "model.layers.23.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
     "model.layers.23.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
     "model.layers.3.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -164,8 +198,10 @@
     "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.3.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.4.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -173,8 +209,10 @@
     "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.4.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.5.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -182,8 +220,10 @@
     "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.5.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.6.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -191,8 +231,10 @@
     "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.6.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.7.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -200,8 +242,10 @@
     "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.7.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.8.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -209,8 +253,10 @@
     "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.8.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.9.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -218,8 +264,10 @@
     "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.9.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.9.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.9.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.9.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.9.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.9.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
     "model.norm.weight": "model-00002-of-00002.safetensors"

modeling_open_lm_hf.py

@@ -0,0 +1,122 @@
+"""
+Custom HuggingFace model for Open LM checkpoints.
+
+Open LM uses LayerNorm (not RMSNorm) and QK norm, which standard
+LlamaForCausalLM does not support. This module provides:
+  - OpenLMConfig: LlamaConfig subclass with qk_norm flag
+  - OpenLMForCausalLM: LlamaForCausalLM subclass with LayerNorm + QK norm
+
+Usage:
+    model = AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True)
+"""
+
+from typing import Callable, Optional
+
+import torch
+import torch.nn as nn
+from transformers import LlamaConfig, LlamaForCausalLM
+from transformers.models.llama.modeling_llama import (
+    ALL_ATTENTION_FUNCTIONS,
+    LlamaAttention,
+    LlamaRMSNorm,
+    apply_rotary_pos_emb,
+    eager_attention_forward,
+)
+
+try:
+    from typing import Unpack
+    from transformers.utils.generic import TransformersKwargs
+except ImportError:
+    pass
+
+from transformers.cache_utils import Cache
+
+
+class OpenLMConfig(LlamaConfig):
+    model_type = "open_lm"
+
+    def __init__(self, qk_norm: bool = True, **kwargs):
+        super().__init__(**kwargs)
+        self.qk_norm = qk_norm
+
+
+class OpenLMAttention(LlamaAttention):
+    """LlamaAttention with QK norm applied before reshape (matching Open LM)."""
+
+    def __init__(self, config: OpenLMConfig, layer_idx: int):
+        super().__init__(config, layer_idx)
+        if getattr(config, "qk_norm", False):
+            self.q_norm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps, bias=False)
+            self.k_norm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps, bias=False)
+        else:
+            self.q_norm = nn.Identity()
+            self.k_norm = nn.Identity()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        # QK norm applied to flat projected vectors BEFORE reshape (matches Open LM)
+        query_states = self.q_norm(self.q_proj(hidden_states)).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states)).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx, cache_kwargs
+            )
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class OpenLMForCausalLM(LlamaForCausalLM):
+    """LlamaForCausalLM with LayerNorm (instead of RMSNorm) and QK norm support."""
+
+    config_class = OpenLMConfig
+
+    def __init__(self, config: OpenLMConfig):
+        super().__init__(config)
+
+        # Replace all LlamaRMSNorm with nn.LayerNorm(bias=False)
+        eps = config.rms_norm_eps
+        hidden_size = config.hidden_size
+
+        self.model.norm = nn.LayerNorm(hidden_size, eps=eps, bias=False)
+        for layer in self.model.layers:
+            layer.input_layernorm = nn.LayerNorm(hidden_size, eps=eps, bias=False)
+            layer.post_attention_layernorm = nn.LayerNorm(hidden_size, eps=eps, bias=False)
+
+            # Replace attention module with QK norm version
+            layer.self_attn = OpenLMAttention(config, layer.self_attn.layer_idx)
+
+        # Re-run post_init to tie weights etc.
+        self.post_init()