off-policy step-150 release checkpoint

README.md

@@ -0,0 +1,31 @@
+---
+license: apache-2.0
+base_model: allenai/OLMoE-1B-7B-0125-Instruct
+tags: [moe, pruning, olmoe, glean]
+---
+
+# winnow-olmoe-math-keep25
+
+Channel-level pruned + healed [OLMoE-1B-7B-0125-Instruct](https://huggingface.co/allenai/OLMoE-1B-7B-0125-Instruct).
+Requires `trust_remote_code=True` (ragged variable-width experts).
+
+| | |
+|---|---|
+| Base | allenai/OLMoE-1B-7B-0125-Instruct |
+| Params | 2.09B |
+| Experts fully deleted | 442/1024 |
+| Keep fraction | 0.25 |
+| Criterion | channel-level REAP, per-layer budgets, block 128, min width 128 |
+| Calibration | Dolmino-math (`scores_0125inst_dolmino-math`) |
+| Final forward top-128 KL | 0.088 |
+
+## Healing
+
+Off-policy forward-KL distillation against cached top-128 teacher targets
+(`dolci_math_curated_opd_top128`), 150 steps, 120k loss tokens/step (18M total), AdamW8bit,
+lr 3e-5, wd 0.1, 10 warmup steps, grad clip 1.0, chat frames, max seq len 2048, seed 1223.
+
+All six models in this release share this recipe and are matched on optimizer
+steps and loss tokens. Note they are **not** matched on data exposure: the math
+cache is 6.48M unique tokens (~2.8 epochs at 150 steps) while the general cache
+is 22.3M (~0.81 epochs).

chat_template.jinja

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

config.json

@@ -0,0 +1,662 @@
+{
+  "architectures": [
+    "PrunedOlmoeForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_pruned_olmoe.PrunedOlmoeConfig",
+    "AutoModelForCausalLM": "modeling_pruned_olmoe.PrunedOlmoeForCausalLM"
+  },
+  "clip_qkv": null,
+  "dtype": "bfloat16",
+  "eos_token_id": 50279,
+  "expert_widths": [
+    [
+      768,
+      512,
+      640,
+      512,
+      896,
+      640,
+      640,
+      512,
+      256,
+      256,
+      384,
+      384,
+      256,
+      512,
+      640,
+      384,
+      256,
+      256,
+      384,
+      896,
+      1024,
+      384,
+      256,
+      384,
+      896,
+      256,
+      384,
+      128,
+      1024,
+      512,
+      640,
+      512
+    ],
+    [
+      256,
+      512,
+      512,
+      384,
+      256,
+      256,
+      384,
+      896,
+      768,
+      256,
+      256,
+      896,
+      256,
+      512,
+      384,
+      384,
+      1024,
+      1024,
+      512,
+      1024,
+      384,
+      256,
+      256,
+      1024,
+      384,
+      384,
+      256,
+      768,
+      384,
+      1024,
+      256,
+      256
+    ],
+    [
+      384,
+      768,
+      128,
+      384,
+      640,
+      512,
+      256,
+      1024,
+      768,
+      640,
+      256,
+      896,
+      512,
+      256,
+      384,
+      256,
+      768,
+      256,
+      640,
+      512,
+      640,
+      512,
+      128,
+      256,
+      640,
+      384,
+      128,
+      384,
+      256,
+      256,
+      1024,
+      256,
+      512,
+      512,
+      256
+    ],
+    [
+      256,
+      128,
+      512,
+      640,
+      512,
+      384,
+      640,
+      256,
+      256,
+      896,
+      640,
+      512,
+      1024,
+      512,
+      896,
+      384,
+      384,
+      256,
+      384,
+      384,
+      384,
+      128,
+      640,
+      640,
+      768,
+      256,
+      768,
+      512,
+      1024,
+      384,
+      1024
+    ],
+    [
+      256,
+      256,
+      512,
+      256,
+      640,
+      256,
+      384,
+      512,
+      640,
+      512,
+      128,
+      1024,
+      256,
+      640,
+      256,
+      128,
+      640,
+      384,
+      384,
+      256,
+      384,
+      384,
+      768,
+      128,
+      384,
+      256,
+      256,
+      512,
+      128,
+      256,
+      512,
+      384,
+      512,
+      768,
+      256,
+      896,
+      384,
+      896
+    ],
+    [
+      1024,
+      256,
+      768,
+      768,
+      384,
+      768,
+      256,
+      768,
+      512,
+      768,
+      512,
+      512,
+      640,
+      384,
+      256,
+      512,
+      512,
+      256,
+      256,
+      512,
+      128,
+      384,
+      256,
+      128,
+      384,
+      384,
+      512,
+      256,
+      128,
+      256,
+      640,
+      128,
+      768,
+      256,
+      768,
+      384
+    ],
+    [
+      256,
+      256,
+      512,
+      256,
+      640,
+      768,
+      768,
+      768,
+      256,
+      256,
+      256,
+      768,
+      1024,
+      640,
+      256,
+      768,
+      640,
+      256,
+      256,
+      384,
+      256,
+      384,
+      512,
+      384,
+      384,
+      640,
+      256,
+      1024,
+      256,
+      768,
+      768,
+      768
+    ],
+    [
+      768,
+      1024,
+      768,
+      512,
+      384,
+      512,
+      384,
+      128,
+      512,
+      256,
+      256,
+      640,
+      384,
+      640,
+      384,
+      512,
+      256,
+      384,
+      384,
+      384,
+      512,
+      640,
+      256,
+      512,
+      128,
+      640,
+      512,
+      256,
+      384,
+      512,
+      384,
+      640,
+      256,
+      768,
+      512
+    ],
+    [
+      128,
+      512,
+      128,
+      256,
+      384,
+      384,
+      512,
+      384,
+      384,
+      896,
+      512,
+      512,
+      256,
+      512,
+      384,
+      128,
+      640,
+      640,
+      256,
+      640,
+      384,
+      384,
+      384,
+      384,
+      896,
+      384,
+      384,
+      384,
+      256,
+      384,
+      128,
+      256,
+      768,
+      896,
+      384,
+      640,
+      640
+    ],
+    [
+      256,
+      256,
+      384,
+      1024,
+      128,
+      768,
+      256,
+      256,
+      384,
+      512,
+      384,
+      256,
+      512,
+      384,
+      256,
+      512,
+      256,
+      384,
+      384,
+      256,
+      640,
+      256,
+      512,
+      512,
+      512,
+      384,
+      384,
+      256,
+      256,
+      384,
+      256,
+      384,
+      896,
+      384,
+      512,
+      640,
+      128,
+      384,
+      384,
+      512
+    ],
+    [
+      384,
+      256,
+      512,
+      384,
+      256,
+      256,
+      384,
+      256,
+      384,
+      256,
+      768,
+      256,
+      384,
+      512,
+      256,
+      640,
+      512,
+      384,
+      512,
+      640,
+      256,
+      512,
+      128,
+      128,
+      256,
+      384,
+      512,
+      512,
+      256,
+      512,
+      512,
+      256,
+      512,
+      256,
+      384,
+      640,
+      640,
+      768,
+      256,
+      512
+    ],
+    [
+      384,
+      256,
+      512,
+      256,
+      640,
+      512,
+      128,
+      384,
+      256,
+      512,
+      384,
+      256,
+      128,
+      384,
+      640,
+      256,
+      384,
+      512,
+      512,
+      640,
+      128,
+      512,
+      256,
+      256,
+      384,
+      128,
+      128,
+      256,
+      512,
+      256,
+      1024,
+      256,
+      384,
+      256,
+      256,
+      512,
+      512,
+      384,
+      768,
+      512,
+      768
+    ],
+    [
+      128,
+      128,
+      384,
+      384,
+      384,
+      512,
+      896,
+      640,
+      256,
+      256,
+      384,
+      512,
+      640,
+      640,
+      256,
+      256,
+      256,
+      128,
+      256,
+      384,
+      384,
+      384,
+      768,
+      640,
+      512,
+      128,
+      384,
+      384,
+      512,
+      768,
+      640,
+      640,
+      128,
+      256,
+      640,
+      128,
+      640,
+      128,
+      256,
+      128,
+      256
+    ],
+    [
+      640,
+      768,
+      640,
+      768,
+      640,
+      256,
+      384,
+      768,
+      256,
+      128,
+      128,
+      384,
+      512,
+      384,
+      512,
+      512,
+      512,
+      640,
+      512,
+      384,
+      512,
+      256,
+      768,
+      256,
+      384,
+      256,
+      640,
+      256,
+      256,
+      128,
+      384,
+      384,
+      640,
+      384,
+      768,
+      256,
+      128
+    ],
+    [
+      384,
+      384,
+      256,
+      256,
+      128,
+      256,
+      896,
+      896,
+      768,
+      512,
+      640,
+      384,
+      512,
+      768,
+      256,
+      384,
+      256,
+      128,
+      640,
+      640,
+      512,
+      384,
+      512,
+      768,
+      768,
+      512,
+      640,
+      512,
+      512,
+      384,
+      128,
+      384,
+      128,
+      512,
+      128,
+      256
+    ],
+    [
+      256,
+      768,
+      512,
+      640,
+      512,
+      384,
+      512,
+      256,
+      640,
+      256,
+      384,
+      896,
+      256,
+      512,
+      128,
+      256,
+      256,
+      256,
+      384,
+      640,
+      256,
+      640,
+      256,
+      640,
+      256,
+      128,
+      256,
+      512,
+      512,
+      128,
+      896,
+      512,
+      384,
+      512,
+      640,
+      384,
+      256,
+      256,
+      256
+    ]
+  ],
+  "glean_metadata": {
+    "base_model": "allenai/OLMoE-1B-7B-0125-Instruct",
+    "block_size": 128,
+    "criterion": "reap",
+    "dead_experts": 442,
+    "keep_fraction": 0.25,
+    "min_width": 128,
+    "params": 2086418432,
+    "scores": "outputs/scores_0125inst_dolmino-math/scores.pt"
+  },
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 1024,
+  "max_position_embeddings": 4096,
+  "model_type": "pruned_olmoe",
+  "norm_topk_prob": false,
+  "num_attention_heads": 16,
+  "num_experts": 64,
+  "num_experts_per_tok": 8,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 16,
+  "output_router_logits": false,
+  "pad_token_id": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "router_aux_loss_coef": 0.01,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.6",
+  "use_cache": false,
+  "vocab_size": 50304
+}

configuration_pruned_olmoe.py

@@ -0,0 +1,27 @@
+"""Configuration for GLEAN-pruned OLMoE: variable-width, variable-count experts.
+"""
+
+from transformers.models.olmoe.configuration_olmoe import OlmoeConfig
+
+
+class PrunedOlmoeConfig(OlmoeConfig):
+    """OlmoeConfig plus a per-(layer, expert) width table.
+
+    ``expert_widths[l]`` lists the SwiGLU intermediate width of each surviving
+    expert in decoder layer ``l``, in expert order. Lists are ragged: layers
+    may keep different numbers of experts (deleted experts simply don't
+    appear — the router in layer ``l`` has ``len(expert_widths[l])`` rows),
+    and each width may differ (multiples of the GEMM block size, 128, for
+    variable-MegaBlocks execution). ``None`` means an unpruned model
+    (uniform ``num_experts`` × ``intermediate_size``).
+
+    The inherited ``num_experts`` / ``intermediate_size`` keep their ORIGINAL
+    (pre-pruning) values for provenance; the width table is authoritative for
+    the built architecture.
+    """
+
+    model_type = "pruned_olmoe"
+
+    def __init__(self, expert_widths: list[list[int]] | None = None, **kwargs):
+        super().__init__(**kwargs)
+        self.expert_widths = expert_widths

generation_config.json

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

model.safetensors

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

modeling_pruned_olmoe.py

@@ -0,0 +1,66 @@
+"""GLEAN-pruned OLMoE: HF-loadable model with ragged (variable-width) experts.
+
+Pattern follows hbfreed/variable-flex-olmo's PrunedFlexOlmoForCausalLM
+(docs/recon/prior-work-hbfreed.md), generalized from one scalar width to a
+per-(layer, expert) width table: ``super().__init__`` builds the uniform
+architecture from the config, then every MoE block is rebuilt to its pruned
+shape — surviving experts only, each at its own width, router sliced to
+match — so the state dict aligns exactly with what
+``glean.prune.prune_channels_global`` leaves behind.
+
+Caveat: ``output_router_logits=True`` (the load-balancing aux loss) assumes a
+uniform ``config.num_experts`` and is unsupported on ragged models.
+"""
+
+import torch.nn as nn
+from transformers.activations import ACT2FN
+from transformers.models.olmoe.modeling_olmoe import OlmoeForCausalLM
+
+from .configuration_pruned_olmoe import PrunedOlmoeConfig
+
+
+class RaggedOlmoeMLP(nn.Module):
+    """OlmoeMLP with an explicit intermediate width (SwiGLU, no biases)."""
+
+    def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+        self.act_fn = ACT2FN[hidden_act]
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class PrunedOlmoeForCausalLM(OlmoeForCausalLM):
+    """OLMoE with per-layer surviving-expert lists at per-expert widths."""
+
+    config_class = PrunedOlmoeConfig
+
+    def __init__(self, config: PrunedOlmoeConfig):
+        super().__init__(config)
+        widths_table = getattr(config, "expert_widths", None)
+        if widths_table is None:
+            return  # unpruned: plain OLMoE
+        if len(widths_table) != len(self.model.layers):
+            raise ValueError(
+                f"expert_widths has {len(widths_table)} rows but the model has "
+                f"{len(self.model.layers)} decoder layers"
+            )
+        for layer, widths in zip(self.model.layers, widths_table):
+            if any(w <= 0 for w in widths):
+                raise ValueError("expert_widths must list surviving experts only (>0)")
+            block = layer.mlp
+            if len(widths) < block.top_k:
+                raise ValueError(
+                    f"a layer keeps {len(widths)} experts < top_k={block.top_k}"
+                )
+            block.num_experts = len(widths)
+            block.gate = nn.Linear(config.hidden_size, len(widths), bias=False)
+            block.experts = nn.ModuleList(
+                RaggedOlmoeMLP(config.hidden_size, w, config.hidden_act)
+                for w in widths
+            )

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "|||IP_ADDRESS|||",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "|||IP_ADDRESS|||",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,247 @@
+{
+  "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
+    },
+    "50277": {
+      "content": "|||EMAIL_ADDRESS|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50278": {
+      "content": "|||PHONE_NUMBER|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50279": {
+      "content": "|||IP_ADDRESS|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50280": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "|||IP_ADDRESS|||",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "|||IP_ADDRESS|||",
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "tokenizer_class": "GPTNeoXTokenizer",
+  "unk_token": null
+}