Upload OLMo parity BF16 MLX

README.md

@@ -0,0 +1,27 @@
+---
+language: en
+pipeline_tag: text-generation
+tags:
+- mlx
+- meshllm
+- parity
+- same-origin
+- olmo
+library_name: mlx
+license: apache-2.0
+---
+
+# OLMo-7B-Instruct-hf Parity BF16 MLX
+
+Same-origin parity artifact derived from `allenai/OLMo-7B-Instruct-hf`.
+
+This repo contains the high-fidelity `bf16` MLX artifact used for mesh-llm backend parity validation against the corresponding GGUF artifact.
+
+Accepted local validation status:
+
+- Exact prompts: MLX now tracks GGUF on the semantic probe set after prompt-render parity fixes, with shared family-level drift on strict one-word canaries
+- Behavior smoke: 0 flagged prompts out of 80 on the MT-Bench-derived harness
+
+Paired GGUF repo:
+
+- `meshllm/olmo-7b-instruct-hf-parity-f16-gguf`

chat_template.jinja

@@ -0,0 +1,12 @@
+{{ eos_token }}{% for message in messages %}
+{% if message['role'] == 'user' %}
+{{ '<|user|>
+' + message['content'] }}
+{% elif message['role'] == 'assistant' %}
+{{ '<|assistant|>
+'  + message['content'] + eos_token }}
+{% endif %}
+{% if loop.last and add_generation_prompt %}
+{{ '<|assistant|>' }}
+{% endif %}
+{% endfor %}

config.json

@@ -0,0 +1,28 @@
+{
+    "architectures": [
+        "OlmoForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "clip_qkv": null,
+    "eos_token_id": 50279,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 11008,
+    "max_position_embeddings": 2048,
+    "model_file": "olmo_model.py",
+    "model_type": "olmo",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 32,
+    "num_key_value_heads": 32,
+    "pad_token_id": 1,
+    "rms_norm_eps": 1e-05,
+    "rope_scaling": null,
+    "rope_theta": 10000.0,
+    "tie_word_embeddings": false,
+    "torch_dtype": "float32",
+    "transformers_version": "4.40.2",
+    "use_cache": true,
+    "vocab_size": 50304
+}

generation_config.json

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

model-00001-of-00003.safetensors

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

model-00002-of-00003.safetensors

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

model-00003-of-00003.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,234 @@
+{
+    "metadata": {
+        "total_size": 13776191488,
+        "total_parameters": 6888095744
+    },
+    "weight_map": {
+        "lm_head.weight": "model-00003-of-00003.safetensors",
+        "model.embed_tokens.weight": "model-00001-of-00003.safetensors",
+        "model.layers.0.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.0.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.1.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.1.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.10.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.10.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.10.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.10.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.10.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.10.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.10.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.11.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.11.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.11.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.11.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.11.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.11.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.11.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.12.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.12.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.12.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.12.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.12.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.12.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.12.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.13.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.13.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.13.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.13.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.13.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.13.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.13.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.14.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.14.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.14.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.14.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.14.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.14.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.14.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.15.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.15.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.15.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.15.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.15.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.15.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.15.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.16.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.16.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.16.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.16.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.16.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.16.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.16.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.17.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.17.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.17.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.17.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.17.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.17.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.17.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.18.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.18.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.18.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.18.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.18.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.18.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.18.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.19.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.19.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.19.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.19.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.19.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.19.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.19.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.2.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.2.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.20.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.20.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.20.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.20.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.20.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.20.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.20.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.21.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.21.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.21.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.21.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.21.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.21.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.21.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.22.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.22.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.22.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.22.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.22.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.22.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.22.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.23.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.23.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.23.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.23.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.23.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.23.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.23.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.24.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.24.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.24.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.24.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.24.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.24.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.24.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.25.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.25.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.25.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.25.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.25.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.25.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.25.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+        "model.layers.26.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.26.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.26.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.26.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.26.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.26.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.26.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.27.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.27.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.27.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.27.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.27.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.27.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.27.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.28.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.28.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.28.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.28.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.28.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.28.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.28.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.29.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.29.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.29.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.29.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.29.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.29.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.29.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.3.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.3.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.30.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.30.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.30.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.30.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.30.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.30.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.30.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.31.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.31.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.31.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.31.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.31.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.31.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.31.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+        "model.layers.4.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.4.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.5.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.5.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.6.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.6.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.7.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.7.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.8.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.8.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.9.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.9.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.9.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.9.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.9.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+        "model.layers.9.self_attn.v_proj.weight": "model-00001-of-00003.safetensors"
+    }
+}

olmo_model.py

@@ -0,0 +1,148 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx_lm.models.activations import swiglu
+from mlx_lm.models.base import BaseModelArgs, create_attention_mask
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    vocab_size: int
+    num_key_value_heads: int = 0
+    rope_theta: float = 10000.0
+    tie_word_embeddings: bool = False
+    attention_bias: bool = False
+    clip_qkv: Optional[float] = None
+
+    def __post_init__(self):
+        if not self.num_key_value_heads:
+            self.num_key_value_heads = self.num_attention_heads
+        if self.num_key_value_heads != self.num_attention_heads:
+            raise ValueError("Grouped-query attention is not yet implemented for this OLMo staging converter.")
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        dim = args.hidden_size
+        self.n_heads = args.num_attention_heads
+        self.head_dim = dim // self.n_heads
+        self.scale = self.head_dim**-0.5
+        self.clip_qkv = args.clip_qkv
+
+        self.q_proj = nn.Linear(dim, dim, bias=args.attention_bias)
+        self.k_proj = nn.Linear(dim, dim, bias=args.attention_bias)
+        self.v_proj = nn.Linear(dim, dim, bias=args.attention_bias)
+        self.o_proj = nn.Linear(dim, dim, bias=args.attention_bias)
+        self.rope = nn.RoPE(self.head_dim, traditional=False, base=args.rope_theta)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        bsz, seq_len, _ = x.shape
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+        v = self.v_proj(x)
+
+        if self.clip_qkv is not None:
+            q = mx.clip(q, -self.clip_qkv, self.clip_qkv)
+            k = mx.clip(k, -self.clip_qkv, self.clip_qkv)
+            v = mx.clip(v, -self.clip_qkv, self.clip_qkv)
+
+        q = q.reshape(bsz, seq_len, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+        k = k.reshape(bsz, seq_len, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+        v = v.reshape(bsz, seq_len, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            q = self.rope(q, offset=cache.offset)
+            k = self.rope(k, offset=cache.offset)
+            k, v = cache.update_and_fetch(k, v)
+        else:
+            q = self.rope(q)
+            k = self.rope(k)
+
+        out = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask)
+        out = out.transpose(0, 2, 1, 3).reshape(bsz, seq_len, -1)
+        return self.o_proj(out)
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        dim = args.hidden_size
+        hidden = args.intermediate_size
+        self.gate_proj = nn.Linear(dim, hidden, bias=False)
+        self.up_proj = nn.Linear(dim, hidden, bias=False)
+        self.down_proj = nn.Linear(hidden, dim, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        dim = args.hidden_size
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args)
+        self.input_layernorm = nn.LayerNorm(dim, affine=False)
+        self.post_attention_layernorm = nn.LayerNorm(dim, affine=False)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = x + self.self_attn(self.input_layernorm(x), mask, cache)
+        return h + self.mlp(self.post_attention_layernorm(h))
+
+
+class InnerModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [DecoderLayer(args) for _ in range(args.num_hidden_layers)]
+        self.norm = nn.LayerNorm(args.hidden_size, affine=False)
+
+    def __call__(self, inputs: mx.array, cache=None):
+        h = self.embed_tokens(inputs)
+        if cache is None:
+            cache = [None] * len(self.layers)
+        mask = create_attention_mask(h, cache[0])
+        for layer, layer_cache in zip(self.layers, cache):
+            h = layer(h, mask, layer_cache)
+        h = self.norm(h)
+        return h, cache
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.model_type = args.model_type
+        self.model = InnerModel(args)
+        self.args = args
+        self.tie_word_embeddings = args.tie_word_embeddings
+        if not self.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(self, inputs: mx.array, cache=None):
+        h, cache = self.model(inputs, cache)
+        if self.tie_word_embeddings:
+            return self.model.embed_tokens.as_linear(h), cache
+        return self.lm_head(h), cache
+
+    @property
+    def layers(self):
+        return self.model.layers

tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|padding|>",
+  "tokenizer_class": "GPTNeoXTokenizer",
+  "trim_offsets": true,
+  "unk_token": null
+}