Add files using upload-large-folder tool

healed/knee0924/keep20/args.json

@@ -0,0 +1,17 @@
+{
+  "student": "outputs/pruned/knee0924/keep20",
+  "teacher": "allenai/OLMoE-1B-7B-0924",
+  "out_dir": "outputs/healed/knee0924/keep20",
+  "steps": 50,
+  "tokens_per_step": 120000,
+  "lr": 3e-05,
+  "weight_decay": 0.1,
+  "seq_len": 2048,
+  "micro_batch": 2,
+  "dataset": "c4",
+  "data_seed": 1224,
+  "eval_seed": 7,
+  "eval_seq": 32,
+  "teacher_device": "cuda:0",
+  "student_device": "cuda:1"
+}

healed/knee0924/keep20/heal_result.json

@@ -0,0 +1,14 @@
+{
+  "student": "outputs/pruned/knee0924/keep20",
+  "pre_heal_ppl": 164.80226135253906,
+  "post_heal_ppl": 33.18513870239258,
+  "steps": 50,
+  "tokens": 6144000,
+  "lr": 3e-05,
+  "eval": {
+    "dataset": "c4",
+    "n_seq": 32,
+    "seq_len": 2048,
+    "seed": 7
+  }
+}

healed/knee0924/keep20/step0050/config.json

@@ -0,0 +1,643 @@
+{
+  "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": [
+    [
+      256,
+      384,
+      640,
+      384,
+      768,
+      1024,
+      640,
+      256,
+      256,
+      256,
+      384,
+      256,
+      256,
+      384,
+      128,
+      256,
+      640,
+      1024,
+      256,
+      256,
+      384,
+      640,
+      768,
+      384,
+      1024,
+      384,
+      512,
+      256
+    ],
+    [
+      512,
+      256,
+      512,
+      512,
+      128,
+      256,
+      640,
+      128,
+      256,
+      1024,
+      128,
+      384,
+      256,
+      256,
+      1024,
+      1024,
+      640,
+      1024,
+      512,
+      512,
+      256,
+      256,
+      256,
+      256,
+      640,
+      384,
+      1024
+    ],
+    [
+      256,
+      768,
+      512,
+      256,
+      896,
+      256,
+      768,
+      256,
+      896,
+      768,
+      640,
+      256,
+      384,
+      256,
+      512,
+      256,
+      384,
+      256,
+      512,
+      256,
+      256,
+      256,
+      512,
+      384,
+      256,
+      256,
+      1024,
+      768
+    ],
+    [
+      512,
+      640,
+      256,
+      256,
+      256,
+      640,
+      128,
+      256,
+      512,
+      256,
+      512,
+      256,
+      640,
+      256,
+      896,
+      512,
+      1024,
+      640,
+      384,
+      256,
+      256,
+      384,
+      512,
+      384,
+      256,
+      256,
+      256,
+      512,
+      256,
+      896
+    ],
+    [
+      256,
+      256,
+      256,
+      256,
+      256,
+      384,
+      640,
+      384,
+      512,
+      384,
+      128,
+      384,
+      256,
+      384,
+      256,
+      384,
+      384,
+      384,
+      256,
+      640,
+      256,
+      256,
+      768,
+      384,
+      512,
+      128,
+      256,
+      128,
+      128,
+      640,
+      256,
+      384,
+      256,
+      640,
+      1024
+    ],
+    [
+      256,
+      512,
+      512,
+      256,
+      640,
+      256,
+      384,
+      256,
+      256,
+      128,
+      384,
+      384,
+      256,
+      256,
+      128,
+      512,
+      128,
+      256,
+      256,
+      256,
+      128,
+      256,
+      256,
+      256,
+      384,
+      256,
+      256,
+      640,
+      256,
+      384,
+      256,
+      256,
+      384,
+      128,
+      384,
+      384,
+      512,
+      512,
+      256,
+      384,
+      256
+    ],
+    [
+      512,
+      128,
+      768,
+      384,
+      384,
+      512,
+      256,
+      512,
+      256,
+      256,
+      256,
+      512,
+      896,
+      384,
+      512,
+      384,
+      384,
+      256,
+      384,
+      384,
+      256,
+      256,
+      384,
+      256,
+      384,
+      384,
+      384,
+      256,
+      384,
+      256,
+      256,
+      256,
+      512,
+      512
+    ],
+    [
+      128,
+      256,
+      640,
+      256,
+      256,
+      384,
+      256,
+      256,
+      256,
+      384,
+      256,
+      256,
+      384,
+      256,
+      512,
+      512,
+      256,
+      768,
+      512,
+      256,
+      128,
+      256,
+      384,
+      256,
+      256,
+      384,
+      640,
+      384,
+      384,
+      384,
+      128,
+      256,
+      256,
+      256,
+      128,
+      256,
+      128,
+      512,
+      384,
+      256
+    ],
+    [
+      256,
+      256,
+      256,
+      256,
+      128,
+      384,
+      384,
+      256,
+      256,
+      384,
+      256,
+      384,
+      384,
+      256,
+      384,
+      768,
+      384,
+      256,
+      512,
+      384,
+      256,
+      256,
+      256,
+      384,
+      512,
+      128,
+      256,
+      384,
+      768,
+      384,
+      256,
+      256,
+      512,
+      256,
+      512,
+      128,
+      256,
+      512
+    ],
+    [
+      256,
+      256,
+      512,
+      256,
+      640,
+      256,
+      384,
+      256,
+      256,
+      896,
+      384,
+      384,
+      256,
+      128,
+      256,
+      256,
+      256,
+      128,
+      512,
+      256,
+      384,
+      384,
+      256,
+      384,
+      384,
+      256,
+      128,
+      384,
+      256,
+      256,
+      384,
+      512,
+      256,
+      512,
+      256,
+      128,
+      384,
+      512,
+      256
+    ],
+    [
+      256,
+      256,
+      384,
+      512,
+      256,
+      256,
+      640,
+      512,
+      128,
+      512,
+      256,
+      512,
+      256,
+      640,
+      384,
+      512,
+      384,
+      384,
+      128,
+      384,
+      256,
+      384,
+      640,
+      384,
+      384,
+      256,
+      256,
+      256,
+      256,
+      384,
+      384,
+      256,
+      256,
+      128,
+      768,
+      128,
+      128
+    ],
+    [
+      384,
+      512,
+      384,
+      384,
+      128,
+      384,
+      256,
+      640,
+      128,
+      128,
+      512,
+      256,
+      256,
+      256,
+      640,
+      512,
+      256,
+      384,
+      384,
+      128,
+      384,
+      256,
+      256,
+      256,
+      256,
+      256,
+      128,
+      256,
+      256,
+      256,
+      256,
+      128,
+      256,
+      256,
+      256,
+      256,
+      512,
+      256,
+      640,
+      384,
+      384
+    ],
+    [
+      384,
+      128,
+      256,
+      384,
+      512,
+      512,
+      512,
+      384,
+      384,
+      256,
+      512,
+      384,
+      128,
+      256,
+      384,
+      512,
+      256,
+      256,
+      256,
+      256,
+      384,
+      128,
+      640,
+      256,
+      256,
+      512,
+      384,
+      640,
+      384,
+      256,
+      512,
+      384,
+      384,
+      384,
+      256,
+      256,
+      128
+    ],
+    [
+      512,
+      512,
+      256,
+      128,
+      384,
+      256,
+      384,
+      640,
+      512,
+      256,
+      384,
+      384,
+      256,
+      384,
+      256,
+      256,
+      256,
+      256,
+      256,
+      256,
+      640,
+      256,
+      512,
+      512,
+      512,
+      384,
+      256,
+      256,
+      128,
+      256,
+      128,
+      384,
+      256,
+      256,
+      384,
+      640,
+      512
+    ],
+    [
+      384,
+      128,
+      256,
+      128,
+      256,
+      384,
+      128,
+      640,
+      128,
+      384,
+      384,
+      512,
+      256,
+      512,
+      384,
+      384,
+      384,
+      512,
+      128,
+      256,
+      256,
+      384,
+      256,
+      512,
+      256,
+      512,
+      640,
+      256,
+      384,
+      128,
+      128,
+      128,
+      512,
+      256,
+      128,
+      640,
+      384,
+      512,
+      256
+    ],
+    [
+      384,
+      640,
+      384,
+      256,
+      640,
+      512,
+      512,
+      512,
+      256,
+      256,
+      384,
+      384,
+      384,
+      768,
+      256,
+      384,
+      512,
+      256,
+      128,
+      384,
+      512,
+      384,
+      256,
+      768,
+      384,
+      256,
+      512,
+      640,
+      256,
+      384,
+      256,
+      256
+    ]
+  ],
+  "glean_metadata": {
+    "base_model": "allenai/OLMoE-1B-7B-0924",
+    "block_size": 128,
+    "criterion": "reap",
+    "dead_experts": 461,
+    "keep_fraction": 0.2,
+    "min_width": 128,
+    "params": 1759223808,
+    "scores": "outputs/scores_0924_c4_2.1M/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": true,
+  "vocab_size": 50304
+}

healed/knee0924/keep20/step0050/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

healed/knee0924/keep20/step0050/generation_config.json

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

healed/knee0924/keep20/step0050/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
+            )

healed/knee0924/keep20/step0050/special_tokens_map.json

@@ -0,0 +1,16 @@
+{
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|padding|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

healed/knee0924/keep20/step0050/tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "|||IP_ADDRESS|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "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": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": null,
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|padding|>",
+  "tokenizer_class": "GPTNeoXTokenizer",
+  "unk_token": null
+}

healed/knee0924/keep20/train_log.jsonl

@@ -0,0 +1,50 @@
+{"step": 1, "kl": 2.799445072809855, "tokens": 122880, "cumulative_loss_tokens": 122880, "secs": 45.8}
+{"step": 2, "kl": 2.00822594165802, "tokens": 122880, "cumulative_loss_tokens": 245760, "secs": 36.7}
+{"step": 3, "kl": 1.7884283105532328, "tokens": 122880, "cumulative_loss_tokens": 368640, "secs": 36.7}
+{"step": 4, "kl": 1.7781060735384624, "tokens": 122880, "cumulative_loss_tokens": 491520, "secs": 36.7}
+{"step": 5, "kl": 1.6532676378885904, "tokens": 122880, "cumulative_loss_tokens": 614400, "secs": 36.7}
+{"step": 6, "kl": 1.4948806603749594, "tokens": 122880, "cumulative_loss_tokens": 737280, "secs": 36.7}
+{"step": 7, "kl": 1.4933756669362386, "tokens": 122880, "cumulative_loss_tokens": 860160, "secs": 36.7}
+{"step": 8, "kl": 1.5488511085510255, "tokens": 122880, "cumulative_loss_tokens": 983040, "secs": 36.7}
+{"step": 9, "kl": 1.4010032296180726, "tokens": 122880, "cumulative_loss_tokens": 1105920, "secs": 36.7}
+{"step": 10, "kl": 1.3937324325243632, "tokens": 122880, "cumulative_loss_tokens": 1228800, "secs": 36.7}
+{"step": 11, "kl": 1.393221026659012, "tokens": 122880, "cumulative_loss_tokens": 1351680, "secs": 36.7}
+{"step": 12, "kl": 1.2561392446358999, "tokens": 122880, "cumulative_loss_tokens": 1474560, "secs": 36.7}
+{"step": 13, "kl": 1.3050465305646262, "tokens": 122880, "cumulative_loss_tokens": 1597440, "secs": 36.7}
+{"step": 14, "kl": 1.244993543624878, "tokens": 122880, "cumulative_loss_tokens": 1720320, "secs": 36.7}
+{"step": 15, "kl": 1.2129537065823872, "tokens": 122880, "cumulative_loss_tokens": 1843200, "secs": 36.7}
+{"step": 16, "kl": 1.3954785863558452, "tokens": 122880, "cumulative_loss_tokens": 1966080, "secs": 36.8}
+{"step": 17, "kl": 1.1894520342350006, "tokens": 122880, "cumulative_loss_tokens": 2088960, "secs": 36.7}
+{"step": 18, "kl": 1.2797086536884308, "tokens": 122880, "cumulative_loss_tokens": 2211840, "secs": 36.7}
+{"step": 19, "kl": 1.1934045394261679, "tokens": 122880, "cumulative_loss_tokens": 2334720, "secs": 36.6}
+{"step": 20, "kl": 1.2127291063467662, "tokens": 122880, "cumulative_loss_tokens": 2457600, "secs": 36.7}
+{"step": 21, "kl": 1.2141900857289631, "tokens": 122880, "cumulative_loss_tokens": 2580480, "secs": 36.6}
+{"step": 22, "kl": 1.2126166780789693, "tokens": 122880, "cumulative_loss_tokens": 2703360, "secs": 36.7}
+{"step": 23, "kl": 1.188023897012075, "tokens": 122880, "cumulative_loss_tokens": 2826240, "secs": 36.7}
+{"step": 24, "kl": 1.220291304588318, "tokens": 122880, "cumulative_loss_tokens": 2949120, "secs": 36.7}
+{"step": 25, "kl": 1.1050480604171753, "tokens": 122880, "cumulative_loss_tokens": 3072000, "secs": 36.7}
+{"step": 26, "kl": 1.107023400068283, "tokens": 122880, "cumulative_loss_tokens": 3194880, "secs": 36.7}
+{"step": 27, "kl": 1.1752450664838154, "tokens": 122880, "cumulative_loss_tokens": 3317760, "secs": 36.7}
+{"step": 28, "kl": 1.1123643855253855, "tokens": 122880, "cumulative_loss_tokens": 3440640, "secs": 36.7}
+{"step": 29, "kl": 1.1148907899856568, "tokens": 122880, "cumulative_loss_tokens": 3563520, "secs": 36.7}
+{"step": 30, "kl": 1.1799089233080546, "tokens": 122880, "cumulative_loss_tokens": 3686400, "secs": 36.7}
+{"step": 31, "kl": 1.1139978369077046, "tokens": 122880, "cumulative_loss_tokens": 3809280, "secs": 36.7}
+{"step": 32, "kl": 1.1341905176639557, "tokens": 122880, "cumulative_loss_tokens": 3932160, "secs": 36.6}
+{"step": 33, "kl": 1.0690260032812755, "tokens": 122880, "cumulative_loss_tokens": 4055040, "secs": 36.7}
+{"step": 34, "kl": 0.9891984204451243, "tokens": 122880, "cumulative_loss_tokens": 4177920, "secs": 36.7}
+{"step": 35, "kl": 1.0958283106486002, "tokens": 122880, "cumulative_loss_tokens": 4300800, "secs": 36.7}
+{"step": 36, "kl": 1.0473421196142831, "tokens": 122880, "cumulative_loss_tokens": 4423680, "secs": 36.6}
+{"step": 37, "kl": 1.1448211173216503, "tokens": 122880, "cumulative_loss_tokens": 4546560, "secs": 36.7}
+{"step": 38, "kl": 1.0664462268352508, "tokens": 122880, "cumulative_loss_tokens": 4669440, "secs": 36.6}
+{"step": 39, "kl": 1.0221477329730988, "tokens": 122880, "cumulative_loss_tokens": 4792320, "secs": 36.7}
+{"step": 40, "kl": 1.082367899020513, "tokens": 122880, "cumulative_loss_tokens": 4915200, "secs": 36.7}
+{"step": 41, "kl": 1.0801438073317209, "tokens": 122880, "cumulative_loss_tokens": 5038080, "secs": 36.6}
+{"step": 42, "kl": 1.0414667387803396, "tokens": 122880, "cumulative_loss_tokens": 5160960, "secs": 36.6}
+{"step": 43, "kl": 1.0454433143138886, "tokens": 122880, "cumulative_loss_tokens": 5283840, "secs": 36.6}
+{"step": 44, "kl": 0.9927415390809377, "tokens": 122880, "cumulative_loss_tokens": 5406720, "secs": 36.6}
+{"step": 45, "kl": 1.0122422436873117, "tokens": 122880, "cumulative_loss_tokens": 5529600, "secs": 36.7}
+{"step": 46, "kl": 1.0627652982870737, "tokens": 122880, "cumulative_loss_tokens": 5652480, "secs": 36.6}
+{"step": 47, "kl": 1.0868095854918163, "tokens": 122880, "cumulative_loss_tokens": 5775360, "secs": 36.6}
+{"step": 48, "kl": 1.0355110168457031, "tokens": 122880, "cumulative_loss_tokens": 5898240, "secs": 36.7}
+{"step": 49, "kl": 1.031413996219635, "tokens": 122880, "cumulative_loss_tokens": 6021120, "secs": 36.8}
+{"step": 50, "kl": 1.0871112326780954, "tokens": 122880, "cumulative_loss_tokens": 6144000, "secs": 36.7}

healed/knee0924/keep25/args.json

@@ -0,0 +1,17 @@
+{
+  "student": "outputs/pruned/knee0924/keep25",
+  "teacher": "allenai/OLMoE-1B-7B-0924",
+  "out_dir": "outputs/healed/knee0924/keep25",
+  "steps": 50,
+  "tokens_per_step": 120000,
+  "lr": 3e-05,
+  "weight_decay": 0.1,
+  "seq_len": 2048,
+  "micro_batch": 2,
+  "dataset": "c4",
+  "data_seed": 1224,
+  "eval_seed": 7,
+  "eval_seq": 32,
+  "teacher_device": "cuda:0",
+  "student_device": "cuda:1"
+}

healed/knee0924/keep25/heal_result.json

@@ -0,0 +1,14 @@
+{
+  "student": "outputs/pruned/knee0924/keep25",
+  "pre_heal_ppl": 105.4487533569336,
+  "post_heal_ppl": 28.708465576171875,
+  "steps": 50,
+  "tokens": 6144000,
+  "lr": 3e-05,
+  "eval": {
+    "dataset": "c4",
+    "n_seq": 32,
+    "seq_len": 2048,
+    "seed": 7
+  }
+}

healed/knee0924/keep25/step0050/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
+            )

healed/knee0924/keep25/step0050/special_tokens_map.json

@@ -0,0 +1,16 @@
+{
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|padding|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

healed/knee0924/keep25/train_log.jsonl

@@ -0,0 +1,50 @@
+{"step": 1, "kl": 2.313830228646596, "tokens": 122880, "cumulative_loss_tokens": 122880, "secs": 47.4}
+{"step": 2, "kl": 1.6554232875506083, "tokens": 122880, "cumulative_loss_tokens": 245760, "secs": 37.5}
+{"step": 3, "kl": 1.4781606793403625, "tokens": 122880, "cumulative_loss_tokens": 368640, "secs": 37.4}
+{"step": 4, "kl": 1.4840904911359152, "tokens": 122880, "cumulative_loss_tokens": 491520, "secs": 37.4}
+{"step": 5, "kl": 1.3798327247301738, "tokens": 122880, "cumulative_loss_tokens": 614400, "secs": 37.4}
+{"step": 6, "kl": 1.2483799874782562, "tokens": 122880, "cumulative_loss_tokens": 737280, "secs": 37.4}
+{"step": 7, "kl": 1.2335848788420358, "tokens": 122880, "cumulative_loss_tokens": 860160, "secs": 37.5}
+{"step": 8, "kl": 1.2804628968238831, "tokens": 122880, "cumulative_loss_tokens": 983040, "secs": 37.4}
+{"step": 9, "kl": 1.1768523673216502, "tokens": 122880, "cumulative_loss_tokens": 1105920, "secs": 37.4}
+{"step": 10, "kl": 1.1638859351476034, "tokens": 122880, "cumulative_loss_tokens": 1228800, "secs": 37.4}
+{"step": 11, "kl": 1.1788139998912812, "tokens": 122880, "cumulative_loss_tokens": 1351680, "secs": 37.4}
+{"step": 12, "kl": 1.0587308764457704, "tokens": 122880, "cumulative_loss_tokens": 1474560, "secs": 37.4}
+{"step": 13, "kl": 1.096179978052775, "tokens": 122880, "cumulative_loss_tokens": 1597440, "secs": 37.4}
+{"step": 14, "kl": 1.0411964118480683, "tokens": 122880, "cumulative_loss_tokens": 1720320, "secs": 37.4}
+{"step": 15, "kl": 1.0206178704897562, "tokens": 122880, "cumulative_loss_tokens": 1843200, "secs": 37.4}
+{"step": 16, "kl": 1.197655294338862, "tokens": 122880, "cumulative_loss_tokens": 1966080, "secs": 37.5}
+{"step": 17, "kl": 0.9969692786534627, "tokens": 122880, "cumulative_loss_tokens": 2088960, "secs": 37.4}
+{"step": 18, "kl": 1.0813468813896179, "tokens": 122880, "cumulative_loss_tokens": 2211840, "secs": 37.4}
+{"step": 19, "kl": 1.001629074414571, "tokens": 122880, "cumulative_loss_tokens": 2334720, "secs": 37.3}
+{"step": 20, "kl": 1.0257816811402638, "tokens": 122880, "cumulative_loss_tokens": 2457600, "secs": 37.4}
+{"step": 21, "kl": 1.0294838051001232, "tokens": 122880, "cumulative_loss_tokens": 2580480, "secs": 37.3}
+{"step": 22, "kl": 1.0281385282675426, "tokens": 122880, "cumulative_loss_tokens": 2703360, "secs": 37.4}
+{"step": 23, "kl": 1.0075202882289886, "tokens": 122880, "cumulative_loss_tokens": 2826240, "secs": 37.4}
+{"step": 24, "kl": 1.0307145178318025, "tokens": 122880, "cumulative_loss_tokens": 2949120, "secs": 37.4}
+{"step": 25, "kl": 0.9371936043103536, "tokens": 122880, "cumulative_loss_tokens": 3072000, "secs": 37.4}
+{"step": 26, "kl": 0.9309814016024271, "tokens": 122880, "cumulative_loss_tokens": 3194880, "secs": 37.4}
+{"step": 27, "kl": 1.0037646392981212, "tokens": 122880, "cumulative_loss_tokens": 3317760, "secs": 37.4}
+{"step": 28, "kl": 0.9431975464026133, "tokens": 122880, "cumulative_loss_tokens": 3440640, "secs": 37.4}
+{"step": 29, "kl": 0.9441798885663351, "tokens": 122880, "cumulative_loss_tokens": 3563520, "secs": 37.4}
+{"step": 30, "kl": 1.0116740624109903, "tokens": 122880, "cumulative_loss_tokens": 3686400, "secs": 37.4}
+{"step": 31, "kl": 0.9417168378829956, "tokens": 122880, "cumulative_loss_tokens": 3809280, "secs": 37.3}
+{"step": 32, "kl": 0.969733864068985, "tokens": 122880, "cumulative_loss_tokens": 3932160, "secs": 37.3}
+{"step": 33, "kl": 0.9051402866840362, "tokens": 122880, "cumulative_loss_tokens": 4055040, "secs": 37.4}
+{"step": 34, "kl": 0.8359194397926331, "tokens": 122880, "cumulative_loss_tokens": 4177920, "secs": 37.4}
+{"step": 35, "kl": 0.9334291179974874, "tokens": 122880, "cumulative_loss_tokens": 4300800, "secs": 37.4}
+{"step": 36, "kl": 0.8888086318969727, "tokens": 122880, "cumulative_loss_tokens": 4423680, "secs": 37.3}
+{"step": 37, "kl": 0.9763996124267578, "tokens": 122880, "cumulative_loss_tokens": 4546560, "secs": 37.4}
+{"step": 38, "kl": 0.9056931614875794, "tokens": 122880, "cumulative_loss_tokens": 4669440, "secs": 37.4}
+{"step": 39, "kl": 0.8712732523679734, "tokens": 122880, "cumulative_loss_tokens": 4792320, "secs": 37.4}
+{"step": 40, "kl": 0.923673681418101, "tokens": 122880, "cumulative_loss_tokens": 4915200, "secs": 37.4}
+{"step": 41, "kl": 0.9259252826372782, "tokens": 122880, "cumulative_loss_tokens": 5038080, "secs": 37.4}
+{"step": 42, "kl": 0.8857697010040283, "tokens": 122880, "cumulative_loss_tokens": 5160960, "secs": 37.4}
+{"step": 43, "kl": 0.8929120381673177, "tokens": 122880, "cumulative_loss_tokens": 5283840, "secs": 37.3}
+{"step": 44, "kl": 0.8391634007294972, "tokens": 122880, "cumulative_loss_tokens": 5406720, "secs": 37.3}
+{"step": 45, "kl": 0.8613568643728892, "tokens": 122880, "cumulative_loss_tokens": 5529600, "secs": 37.3}
+{"step": 46, "kl": 0.9118083119392395, "tokens": 122880, "cumulative_loss_tokens": 5652480, "secs": 37.3}
+{"step": 47, "kl": 0.9288128594557444, "tokens": 122880, "cumulative_loss_tokens": 5775360, "secs": 37.4}
+{"step": 48, "kl": 0.8719660659631093, "tokens": 122880, "cumulative_loss_tokens": 5898240, "secs": 37.4}
+{"step": 49, "kl": 0.8787553608417511, "tokens": 122880, "cumulative_loss_tokens": 6021120, "secs": 37.4}
+{"step": 50, "kl": 0.9383302330970764, "tokens": 122880, "cumulative_loss_tokens": 6144000, "secs": 37.3}

healed/knee0924/keep30/args.json

@@ -0,0 +1,17 @@
+{
+  "student": "outputs/pruned/knee0924/keep30",
+  "teacher": "allenai/OLMoE-1B-7B-0924",
+  "out_dir": "outputs/healed/knee0924/keep30",
+  "steps": 50,
+  "tokens_per_step": 120000,
+  "lr": 3e-05,
+  "weight_decay": 0.1,
+  "seq_len": 2048,
+  "micro_batch": 2,
+  "dataset": "c4",
+  "data_seed": 1224,
+  "eval_seed": 7,
+  "eval_seq": 32,
+  "teacher_device": "cuda:0",
+  "student_device": "cuda:1"
+}

healed/knee0924/keep30/heal_result.json

@@ -0,0 +1,14 @@
+{
+  "student": "outputs/pruned/knee0924/keep30",
+  "pre_heal_ppl": 69.09920501708984,
+  "post_heal_ppl": 25.705188751220703,
+  "steps": 50,
+  "tokens": 6144000,
+  "lr": 3e-05,
+  "eval": {
+    "dataset": "c4",
+    "n_seq": 32,
+    "seq_len": 2048,
+    "seed": 7
+  }
+}

healed/knee0924/keep30/train_log.jsonl

@@ -0,0 +1,50 @@
+{"step": 1, "kl": 1.9134400963783265, "tokens": 122880, "cumulative_loss_tokens": 122880, "secs": 48.5}
+{"step": 2, "kl": 1.3438369194666544, "tokens": 122880, "cumulative_loss_tokens": 245760, "secs": 38.1}
+{"step": 3, "kl": 1.1803164184093475, "tokens": 122880, "cumulative_loss_tokens": 368640, "secs": 38.1}
+{"step": 4, "kl": 1.2290376603603363, "tokens": 122880, "cumulative_loss_tokens": 491520, "secs": 38.1}
+{"step": 5, "kl": 1.1476012349128724, "tokens": 122880, "cumulative_loss_tokens": 614400, "secs": 38.1}
+{"step": 6, "kl": 1.0419771254062653, "tokens": 122880, "cumulative_loss_tokens": 737280, "secs": 38.1}
+{"step": 7, "kl": 1.0220232804616292, "tokens": 122880, "cumulative_loss_tokens": 860160, "secs": 38.1}
+{"step": 8, "kl": 1.0815871198972067, "tokens": 122880, "cumulative_loss_tokens": 983040, "secs": 38.1}
+{"step": 9, "kl": 0.9827445109685262, "tokens": 122880, "cumulative_loss_tokens": 1105920, "secs": 38.1}
+{"step": 10, "kl": 0.9709443469842275, "tokens": 122880, "cumulative_loss_tokens": 1228800, "secs": 38.1}
+{"step": 11, "kl": 0.991158950328827, "tokens": 122880, "cumulative_loss_tokens": 1351680, "secs": 38.1}
+{"step": 12, "kl": 0.8874583780765534, "tokens": 122880, "cumulative_loss_tokens": 1474560, "secs": 38.1}
+{"step": 13, "kl": 0.9268575588862101, "tokens": 122880, "cumulative_loss_tokens": 1597440, "secs": 38.1}
+{"step": 14, "kl": 0.8714903970559438, "tokens": 122880, "cumulative_loss_tokens": 1720320, "secs": 38.1}
+{"step": 15, "kl": 0.8550567487875621, "tokens": 122880, "cumulative_loss_tokens": 1843200, "secs": 38.1}
+{"step": 16, "kl": 1.032442420721054, "tokens": 122880, "cumulative_loss_tokens": 1966080, "secs": 38.1}
+{"step": 17, "kl": 0.8429429848988851, "tokens": 122880, "cumulative_loss_tokens": 2088960, "secs": 38.1}
+{"step": 18, "kl": 0.9245830039183299, "tokens": 122880, "cumulative_loss_tokens": 2211840, "secs": 38.1}
+{"step": 19, "kl": 0.8454516013463338, "tokens": 122880, "cumulative_loss_tokens": 2334720, "secs": 38.1}
+{"step": 20, "kl": 0.8760805626710256, "tokens": 122880, "cumulative_loss_tokens": 2457600, "secs": 38.1}
+{"step": 21, "kl": 0.8805038313070933, "tokens": 122880, "cumulative_loss_tokens": 2580480, "secs": 38.0}
+{"step": 22, "kl": 0.8789269506931305, "tokens": 122880, "cumulative_loss_tokens": 2703360, "secs": 38.1}
+{"step": 23, "kl": 0.861937924226125, "tokens": 122880, "cumulative_loss_tokens": 2826240, "secs": 38.1}
+{"step": 24, "kl": 0.8894271711508434, "tokens": 122880, "cumulative_loss_tokens": 2949120, "secs": 38.1}
+{"step": 25, "kl": 0.8030505537986755, "tokens": 122880, "cumulative_loss_tokens": 3072000, "secs": 38.1}
+{"step": 26, "kl": 0.7937345544497172, "tokens": 122880, "cumulative_loss_tokens": 3194880, "secs": 38.1}
+{"step": 27, "kl": 0.8633080859978993, "tokens": 122880, "cumulative_loss_tokens": 3317760, "secs": 38.1}
+{"step": 28, "kl": 0.8080259402592976, "tokens": 122880, "cumulative_loss_tokens": 3440640, "secs": 38.1}
+{"step": 29, "kl": 0.8120014548301697, "tokens": 122880, "cumulative_loss_tokens": 3563520, "secs": 38.1}
+{"step": 30, "kl": 0.8738385081291199, "tokens": 122880, "cumulative_loss_tokens": 3686400, "secs": 38.1}
+{"step": 31, "kl": 0.8093875706195831, "tokens": 122880, "cumulative_loss_tokens": 3809280, "secs": 38.1}
+{"step": 32, "kl": 0.8427861034870148, "tokens": 122880, "cumulative_loss_tokens": 3932160, "secs": 38.0}
+{"step": 33, "kl": 0.7789631843566894, "tokens": 122880, "cumulative_loss_tokens": 4055040, "secs": 38.1}
+{"step": 34, "kl": 0.7124047259489695, "tokens": 122880, "cumulative_loss_tokens": 4177920, "secs": 38.1}
+{"step": 35, "kl": 0.8032700498898824, "tokens": 122880, "cumulative_loss_tokens": 4300800, "secs": 38.1}
+{"step": 36, "kl": 0.7631795922915141, "tokens": 122880, "cumulative_loss_tokens": 4423680, "secs": 38.1}
+{"step": 37, "kl": 0.84235600233078, "tokens": 122880, "cumulative_loss_tokens": 4546560, "secs": 38.1}
+{"step": 38, "kl": 0.7787195940812429, "tokens": 122880, "cumulative_loss_tokens": 4669440, "secs": 38.1}
+{"step": 39, "kl": 0.7503240784009297, "tokens": 122880, "cumulative_loss_tokens": 4792320, "secs": 38.1}
+{"step": 40, "kl": 0.8020394384860993, "tokens": 122880, "cumulative_loss_tokens": 4915200, "secs": 38.1}
+{"step": 41, "kl": 0.8019351383050283, "tokens": 122880, "cumulative_loss_tokens": 5038080, "secs": 38.1}
+{"step": 42, "kl": 0.7664479116598765, "tokens": 122880, "cumulative_loss_tokens": 5160960, "secs": 38.1}
+{"step": 43, "kl": 0.7694141308466593, "tokens": 122880, "cumulative_loss_tokens": 5283840, "secs": 38.0}
+{"step": 44, "kl": 0.7206377108891805, "tokens": 122880, "cumulative_loss_tokens": 5406720, "secs": 38.1}
+{"step": 45, "kl": 0.7483313242594402, "tokens": 122880, "cumulative_loss_tokens": 5529600, "secs": 38.1}
+{"step": 46, "kl": 0.7960013389587403, "tokens": 122880, "cumulative_loss_tokens": 5652480, "secs": 38.0}
+{"step": 47, "kl": 0.8123511930306753, "tokens": 122880, "cumulative_loss_tokens": 5775360, "secs": 38.1}
+{"step": 48, "kl": 0.7524132609367371, "tokens": 122880, "cumulative_loss_tokens": 5898240, "secs": 38.1}
+{"step": 49, "kl": 0.760091096162796, "tokens": 122880, "cumulative_loss_tokens": 6021120, "secs": 38.1}
+{"step": 50, "kl": 0.8231359044710795, "tokens": 122880, "cumulative_loss_tokens": 6144000, "secs": 38.0}

healed/knee0924/keep40/args.json

@@ -0,0 +1,17 @@
+{
+  "student": "outputs/pruned/knee0924/keep40",
+  "teacher": "allenai/OLMoE-1B-7B-0924",
+  "out_dir": "outputs/healed/knee0924/keep40",
+  "steps": 50,
+  "tokens_per_step": 120000,
+  "lr": 3e-05,
+  "weight_decay": 0.1,
+  "seq_len": 2048,
+  "micro_batch": 2,
+  "dataset": "c4",
+  "data_seed": 1224,
+  "eval_seed": 7,
+  "eval_seq": 32,
+  "teacher_device": "cuda:0",
+  "student_device": "cuda:1"
+}

healed/knee0924/keep40/heal_result.json

@@ -0,0 +1,14 @@
+{
+  "student": "outputs/pruned/knee0924/keep40",
+  "pre_heal_ppl": 34.65605163574219,
+  "post_heal_ppl": 21.057231903076172,
+  "steps": 50,
+  "tokens": 6144000,
+  "lr": 3e-05,
+  "eval": {
+    "dataset": "c4",
+    "n_seq": 32,
+    "seq_len": 2048,
+    "seed": 7
+  }
+}

healed/knee0924/keep40/train_log.jsonl

@@ -0,0 +1,50 @@
+{"step": 1, "kl": 1.203008109331131, "tokens": 122880, "cumulative_loss_tokens": 122880, "secs": 49.6}
+{"step": 2, "kl": 0.8977614363034566, "tokens": 122880, "cumulative_loss_tokens": 245760, "secs": 40.4}
+{"step": 3, "kl": 0.7635675986607869, "tokens": 122880, "cumulative_loss_tokens": 368640, "secs": 40.4}
+{"step": 4, "kl": 0.8397774795691172, "tokens": 122880, "cumulative_loss_tokens": 491520, "secs": 40.3}
+{"step": 5, "kl": 0.7817749043305715, "tokens": 122880, "cumulative_loss_tokens": 614400, "secs": 40.4}
+{"step": 6, "kl": 0.7146296739578247, "tokens": 122880, "cumulative_loss_tokens": 737280, "secs": 40.4}
+{"step": 7, "kl": 0.7007134864727657, "tokens": 122880, "cumulative_loss_tokens": 860160, "secs": 40.4}
+{"step": 8, "kl": 0.7662253022193909, "tokens": 122880, "cumulative_loss_tokens": 983040, "secs": 40.4}
+{"step": 9, "kl": 0.6819654544194539, "tokens": 122880, "cumulative_loss_tokens": 1105920, "secs": 40.3}
+{"step": 10, "kl": 0.6813271423180898, "tokens": 122880, "cumulative_loss_tokens": 1228800, "secs": 40.5}
+{"step": 11, "kl": 0.7043846547603607, "tokens": 122880, "cumulative_loss_tokens": 1351680, "secs": 40.5}
+{"step": 12, "kl": 0.6235980540513992, "tokens": 122880, "cumulative_loss_tokens": 1474560, "secs": 42.6}
+{"step": 13, "kl": 0.6587181915839513, "tokens": 122880, "cumulative_loss_tokens": 1597440, "secs": 40.7}
+{"step": 14, "kl": 0.608491646250089, "tokens": 122880, "cumulative_loss_tokens": 1720320, "secs": 40.2}
+{"step": 15, "kl": 0.6108660946289698, "tokens": 122880, "cumulative_loss_tokens": 1843200, "secs": 40.4}
+{"step": 16, "kl": 0.7790463368097941, "tokens": 122880, "cumulative_loss_tokens": 1966080, "secs": 40.4}
+{"step": 17, "kl": 0.6061649690071742, "tokens": 122880, "cumulative_loss_tokens": 2088960, "secs": 40.6}
+{"step": 18, "kl": 0.6803837349017461, "tokens": 122880, "cumulative_loss_tokens": 2211840, "secs": 40.4}
+{"step": 19, "kl": 0.6089968343575796, "tokens": 122880, "cumulative_loss_tokens": 2334720, "secs": 40.2}
+{"step": 20, "kl": 0.6348114997148514, "tokens": 122880, "cumulative_loss_tokens": 2457600, "secs": 40.4}
+{"step": 21, "kl": 0.6464616119861603, "tokens": 122880, "cumulative_loss_tokens": 2580480, "secs": 40.3}
+{"step": 22, "kl": 0.6489827265342076, "tokens": 122880, "cumulative_loss_tokens": 2703360, "secs": 40.3}
+{"step": 23, "kl": 0.6233947813510895, "tokens": 122880, "cumulative_loss_tokens": 2826240, "secs": 40.3}
+{"step": 24, "kl": 0.6494492749373119, "tokens": 122880, "cumulative_loss_tokens": 2949120, "secs": 40.3}
+{"step": 25, "kl": 0.5768886228402456, "tokens": 122880, "cumulative_loss_tokens": 3072000, "secs": 40.3}
+{"step": 26, "kl": 0.566129986445109, "tokens": 122880, "cumulative_loss_tokens": 3194880, "secs": 40.3}
+{"step": 27, "kl": 0.6275245815515518, "tokens": 122880, "cumulative_loss_tokens": 3317760, "secs": 40.3}
+{"step": 28, "kl": 0.5870765805244446, "tokens": 122880, "cumulative_loss_tokens": 3440640, "secs": 40.3}
+{"step": 29, "kl": 0.5843763609727224, "tokens": 122880, "cumulative_loss_tokens": 3563520, "secs": 40.3}
+{"step": 30, "kl": 0.639552503824234, "tokens": 122880, "cumulative_loss_tokens": 3686400, "secs": 40.2}
+{"step": 31, "kl": 0.5818809111913045, "tokens": 122880, "cumulative_loss_tokens": 3809280, "secs": 40.3}
+{"step": 32, "kl": 0.6239005188147227, "tokens": 122880, "cumulative_loss_tokens": 3932160, "secs": 40.3}
+{"step": 33, "kl": 0.5610700309276581, "tokens": 122880, "cumulative_loss_tokens": 4055040, "secs": 40.3}
+{"step": 34, "kl": 0.5062820067008337, "tokens": 122880, "cumulative_loss_tokens": 4177920, "secs": 40.3}
+{"step": 35, "kl": 0.584183860818545, "tokens": 122880, "cumulative_loss_tokens": 4300800, "secs": 40.3}
+{"step": 36, "kl": 0.550735863049825, "tokens": 122880, "cumulative_loss_tokens": 4423680, "secs": 40.3}
+{"step": 37, "kl": 0.6168364971876145, "tokens": 122880, "cumulative_loss_tokens": 4546560, "secs": 40.3}
+{"step": 38, "kl": 0.5639737645785013, "tokens": 122880, "cumulative_loss_tokens": 4669440, "secs": 40.2}
+{"step": 39, "kl": 0.5465332046151161, "tokens": 122880, "cumulative_loss_tokens": 4792320, "secs": 40.3}
+{"step": 40, "kl": 0.5914001047611237, "tokens": 122880, "cumulative_loss_tokens": 4915200, "secs": 40.3}
+{"step": 41, "kl": 0.5894514630238216, "tokens": 122880, "cumulative_loss_tokens": 5038080, "secs": 40.3}
+{"step": 42, "kl": 0.5617780417203904, "tokens": 122880, "cumulative_loss_tokens": 5160960, "secs": 40.2}
+{"step": 43, "kl": 0.5585460325082143, "tokens": 122880, "cumulative_loss_tokens": 5283840, "secs": 40.2}
+{"step": 44, "kl": 0.5199935048818588, "tokens": 122880, "cumulative_loss_tokens": 5406720, "secs": 40.2}
+{"step": 45, "kl": 0.5476068059603373, "tokens": 122880, "cumulative_loss_tokens": 5529600, "secs": 40.3}
+{"step": 46, "kl": 0.5946912695964177, "tokens": 122880, "cumulative_loss_tokens": 5652480, "secs": 40.2}
+{"step": 47, "kl": 0.6040740340948105, "tokens": 122880, "cumulative_loss_tokens": 5775360, "secs": 40.2}
+{"step": 48, "kl": 0.5480124910672506, "tokens": 122880, "cumulative_loss_tokens": 5898240, "secs": 40.3}
+{"step": 49, "kl": 0.5551868160565694, "tokens": 122880, "cumulative_loss_tokens": 6021120, "secs": 40.3}
+{"step": 50, "kl": 0.616097284356753, "tokens": 122880, "cumulative_loss_tokens": 6144000, "secs": 40.2}

healed/knee0924/keep50/args.json

@@ -0,0 +1,17 @@
+{
+  "student": "outputs/pruned/knee0924/keep50",
+  "teacher": "allenai/OLMoE-1B-7B-0924",
+  "out_dir": "outputs/healed/knee0924/keep50",
+  "steps": 50,
+  "tokens_per_step": 120000,
+  "lr": 3e-05,
+  "weight_decay": 0.1,
+  "seq_len": 2048,
+  "micro_batch": 2,
+  "dataset": "c4",
+  "data_seed": 1224,
+  "eval_seed": 7,
+  "eval_seq": 32,
+  "teacher_device": "cuda:0",
+  "student_device": "cuda:1"
+}

healed/knee0924/keep50/heal_result.json

@@ -0,0 +1,14 @@
+{
+  "student": "outputs/pruned/knee0924/keep50",
+  "pre_heal_ppl": 25.590167999267578,
+  "post_heal_ppl": 18.254446029663086,
+  "steps": 50,
+  "tokens": 6144000,
+  "lr": 3e-05,
+  "eval": {
+    "dataset": "c4",
+    "n_seq": 32,
+    "seq_len": 2048,
+    "seed": 7
+  }
+}

healed/knee0924/keep50/train_log.jsonl

@@ -0,0 +1,50 @@
+{"step": 1, "kl": 0.8699087758858999, "tokens": 122880, "cumulative_loss_tokens": 122880, "secs": 52.5}
+{"step": 2, "kl": 0.6575553546349208, "tokens": 122880, "cumulative_loss_tokens": 245760, "secs": 42.0}
+{"step": 3, "kl": 0.5402249326308568, "tokens": 122880, "cumulative_loss_tokens": 368640, "secs": 42.1}
+{"step": 4, "kl": 0.6123988221089045, "tokens": 122880, "cumulative_loss_tokens": 491520, "secs": 42.0}
+{"step": 5, "kl": 0.5700773159662883, "tokens": 122880, "cumulative_loss_tokens": 614400, "secs": 42.0}
+{"step": 6, "kl": 0.5183772613604863, "tokens": 122880, "cumulative_loss_tokens": 737280, "secs": 42.0}
+{"step": 7, "kl": 0.5071838061014812, "tokens": 122880, "cumulative_loss_tokens": 860160, "secs": 42.0}
+{"step": 8, "kl": 0.5582418620586396, "tokens": 122880, "cumulative_loss_tokens": 983040, "secs": 42.0}
+{"step": 9, "kl": 0.4857077270746231, "tokens": 122880, "cumulative_loss_tokens": 1105920, "secs": 42.0}
+{"step": 10, "kl": 0.4897372633218765, "tokens": 122880, "cumulative_loss_tokens": 1228800, "secs": 42.0}
+{"step": 11, "kl": 0.5120153983434041, "tokens": 122880, "cumulative_loss_tokens": 1351680, "secs": 42.1}
+{"step": 12, "kl": 0.4449257473150889, "tokens": 122880, "cumulative_loss_tokens": 1474560, "secs": 42.0}
+{"step": 13, "kl": 0.4799923410018285, "tokens": 122880, "cumulative_loss_tokens": 1597440, "secs": 42.0}
+{"step": 14, "kl": 0.42896502117315927, "tokens": 122880, "cumulative_loss_tokens": 1720320, "secs": 42.0}
+{"step": 15, "kl": 0.4418807705243429, "tokens": 122880, "cumulative_loss_tokens": 1843200, "secs": 42.0}
+{"step": 16, "kl": 0.601253342628479, "tokens": 122880, "cumulative_loss_tokens": 1966080, "secs": 42.1}
+{"step": 17, "kl": 0.4330762306849162, "tokens": 122880, "cumulative_loss_tokens": 2088960, "secs": 42.1}
+{"step": 18, "kl": 0.5030014336109161, "tokens": 122880, "cumulative_loss_tokens": 2211840, "secs": 42.0}
+{"step": 19, "kl": 0.44243644972642265, "tokens": 122880, "cumulative_loss_tokens": 2334720, "secs": 42.0}
+{"step": 20, "kl": 0.4637090047200521, "tokens": 122880, "cumulative_loss_tokens": 2457600, "secs": 42.0}
+{"step": 21, "kl": 0.4785055696964264, "tokens": 122880, "cumulative_loss_tokens": 2580480, "secs": 42.0}
+{"step": 22, "kl": 0.47858787775039674, "tokens": 122880, "cumulative_loss_tokens": 2703360, "secs": 42.0}
+{"step": 23, "kl": 0.4540826499462128, "tokens": 122880, "cumulative_loss_tokens": 2826240, "secs": 42.0}
+{"step": 24, "kl": 0.4782628983259201, "tokens": 122880, "cumulative_loss_tokens": 2949120, "secs": 42.0}
+{"step": 25, "kl": 0.416959140698115, "tokens": 122880, "cumulative_loss_tokens": 3072000, "secs": 42.0}
+{"step": 26, "kl": 0.40522719323635104, "tokens": 122880, "cumulative_loss_tokens": 3194880, "secs": 42.0}
+{"step": 27, "kl": 0.4650170256694158, "tokens": 122880, "cumulative_loss_tokens": 3317760, "secs": 42.0}
+{"step": 28, "kl": 0.43529969255129497, "tokens": 122880, "cumulative_loss_tokens": 3440640, "secs": 42.0}
+{"step": 29, "kl": 0.4287110875050227, "tokens": 122880, "cumulative_loss_tokens": 3563520, "secs": 42.0}
+{"step": 30, "kl": 0.4738962093989054, "tokens": 122880, "cumulative_loss_tokens": 3686400, "secs": 42.1}
+{"step": 31, "kl": 0.4250258584817251, "tokens": 122880, "cumulative_loss_tokens": 3809280, "secs": 42.0}
+{"step": 32, "kl": 0.47182846466700235, "tokens": 122880, "cumulative_loss_tokens": 3932160, "secs": 42.0}
+{"step": 33, "kl": 0.4137675185998281, "tokens": 122880, "cumulative_loss_tokens": 4055040, "secs": 42.0}
+{"step": 34, "kl": 0.3596928596496582, "tokens": 122880, "cumulative_loss_tokens": 4177920, "secs": 42.1}
+{"step": 35, "kl": 0.4286925872166952, "tokens": 122880, "cumulative_loss_tokens": 4300800, "secs": 42.0}
+{"step": 36, "kl": 0.4013708233833313, "tokens": 122880, "cumulative_loss_tokens": 4423680, "secs": 42.0}
+{"step": 37, "kl": 0.456011638045311, "tokens": 122880, "cumulative_loss_tokens": 4546560, "secs": 42.1}
+{"step": 38, "kl": 0.41011486351490023, "tokens": 122880, "cumulative_loss_tokens": 4669440, "secs": 42.0}
+{"step": 39, "kl": 0.4057838886976242, "tokens": 122880, "cumulative_loss_tokens": 4792320, "secs": 42.1}
+{"step": 40, "kl": 0.44491428236166636, "tokens": 122880, "cumulative_loss_tokens": 4915200, "secs": 42.1}
+{"step": 41, "kl": 0.44118417501449586, "tokens": 122880, "cumulative_loss_tokens": 5038080, "secs": 42.1}
+{"step": 42, "kl": 0.4141279463966688, "tokens": 122880, "cumulative_loss_tokens": 5160960, "secs": 42.1}
+{"step": 43, "kl": 0.41535414357980094, "tokens": 122880, "cumulative_loss_tokens": 5283840, "secs": 42.0}
+{"step": 44, "kl": 0.37733205656210583, "tokens": 122880, "cumulative_loss_tokens": 5406720, "secs": 42.1}
+{"step": 45, "kl": 0.40215973754723866, "tokens": 122880, "cumulative_loss_tokens": 5529600, "secs": 42.1}
+{"step": 46, "kl": 0.44426603813966115, "tokens": 122880, "cumulative_loss_tokens": 5652480, "secs": 41.9}
+{"step": 47, "kl": 0.45498027503490446, "tokens": 122880, "cumulative_loss_tokens": 5775360, "secs": 42.1}
+{"step": 48, "kl": 0.3959922303756078, "tokens": 122880, "cumulative_loss_tokens": 5898240, "secs": 42.1}
+{"step": 49, "kl": 0.405680842200915, "tokens": 122880, "cumulative_loss_tokens": 6021120, "secs": 42.1}
+{"step": 50, "kl": 0.46550515393416086, "tokens": 122880, "cumulative_loss_tokens": 6144000, "secs": 42.0}

healed/mixdistill_smoke/step0006/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 %}

healed/mixdistill_smoke/step0006/config.json

@@ -0,0 +1,887 @@
+{
+  "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": [
+    [
+      1024,
+      384,
+      256,
+      256,
+      768,
+      1024,
+      128,
+      768,
+      384,
+      128,
+      1024,
+      384,
+      384,
+      128,
+      640,
+      896,
+      896,
+      256,
+      768,
+      128,
+      768,
+      768,
+      768,
+      768,
+      640,
+      896,
+      256,
+      128,
+      384,
+      896,
+      256,
+      768,
+      512,
+      640,
+      768,
+      896,
+      1024,
+      384,
+      384,
+      512,
+      384,
+      256,
+      768,
+      640,
+      896,
+      896,
+      512,
+      768,
+      512,
+      384,
+      1024,
+      896,
+      896,
+      768,
+      128
+    ],
+    [
+      768,
+      896,
+      896,
+      256,
+      896,
+      384,
+      768,
+      768,
+      1024,
+      384,
+      512,
+      1024,
+      768,
+      768,
+      1024,
+      640,
+      128,
+      768,
+      384,
+      768,
+      128,
+      256,
+      512,
+      384,
+      896,
+      256,
+      1024,
+      1024,
+      896,
+      1024,
+      896,
+      768,
+      256,
+      768,
+      128,
+      1024,
+      768,
+      768,
+      640,
+      384,
+      1024,
+      384,
+      640,
+      896,
+      1024,
+      512,
+      512,
+      512,
+      640
+    ],
+    [
+      768,
+      1024,
+      512,
+      384,
+      384,
+      768,
+      256,
+      384,
+      1024,
+      896,
+      640,
+      256,
+      1024,
+      1024,
+      1024,
+      640,
+      1024,
+      896,
+      512,
+      768,
+      128,
+      768,
+      1024,
+      768,
+      256,
+      1024,
+      896,
+      1024,
+      896,
+      512,
+      640,
+      384,
+      896,
+      384,
+      512,
+      640,
+      256,
+      768,
+      256,
+      768,
+      768,
+      640,
+      1024,
+      640,
+      896,
+      768,
+      768,
+      256
+    ],
+    [
+      640,
+      640,
+      896,
+      512,
+      640,
+      896,
+      896,
+      384,
+      896,
+      896,
+      384,
+      896,
+      640,
+      256,
+      640,
+      1024,
+      1024,
+      768,
+      1024,
+      896,
+      1024,
+      768,
+      896,
+      256,
+      512,
+      768,
+      1024,
+      256,
+      768,
+      512,
+      256,
+      640,
+      1024,
+      1024,
+      512,
+      1024,
+      768,
+      256,
+      1024,
+      384,
+      896,
+      1024,
+      896,
+      1024,
+      384
+    ],
+    [
+      256,
+      640,
+      640,
+      896,
+      768,
+      896,
+      768,
+      768,
+      896,
+      896,
+      1024,
+      256,
+      512,
+      1024,
+      640,
+      896,
+      512,
+      512,
+      384,
+      384,
+      256,
+      384,
+      384,
+      896,
+      896,
+      768,
+      640,
+      896,
+      768,
+      1024,
+      512,
+      640,
+      512,
+      640,
+      896,
+      512,
+      512,
+      384,
+      640,
+      896,
+      896,
+      896,
+      1024,
+      640,
+      1024,
+      640,
+      1024
+    ],
+    [
+      1024,
+      512,
+      1024,
+      1024,
+      640,
+      896,
+      640,
+      1024,
+      896,
+      384,
+      1024,
+      128,
+      896,
+      768,
+      1024,
+      768,
+      640,
+      896,
+      768,
+      640,
+      512,
+      896,
+      512,
+      640,
+      256,
+      768,
+      640,
+      768,
+      384,
+      896,
+      512,
+      512,
+      256,
+      512,
+      896,
+      256,
+      384,
+      640,
+      512,
+      640,
+      896,
+      512,
+      1024,
+      256,
+      768,
+      1024,
+      768,
+      256,
+      256
+    ],
+    [
+      640,
+      896,
+      1024,
+      896,
+      1024,
+      1024,
+      1024,
+      512,
+      256,
+      256,
+      1024,
+      768,
+      512,
+      768,
+      1024,
+      1024,
+      1024,
+      384,
+      512,
+      1024,
+      512,
+      1024,
+      128,
+      640,
+      640,
+      896,
+      768,
+      128,
+      256,
+      256,
+      256,
+      256,
+      896,
+      1024,
+      1024,
+      384,
+      896,
+      256,
+      896,
+      640,
+      1024,
+      384,
+      640,
+      256,
+      1024,
+      1024,
+      1024
+    ],
+    [
+      1024,
+      384,
+      1024,
+      1024,
+      256,
+      128,
+      256,
+      384,
+      256,
+      384,
+      896,
+      768,
+      896,
+      896,
+      512,
+      896,
+      640,
+      384,
+      384,
+      896,
+      768,
+      384,
+      896,
+      768,
+      768,
+      512,
+      896,
+      768,
+      768,
+      896,
+      128,
+      896,
+      512,
+      256,
+      768,
+      128,
+      384,
+      256,
+      896,
+      896,
+      384,
+      768,
+      512,
+      640,
+      256,
+      768,
+      640,
+      896,
+      384,
+      512,
+      1024,
+      768,
+      384
+    ],
+    [
+      512,
+      768,
+      512,
+      256,
+      128,
+      640,
+      384,
+      640,
+      768,
+      896,
+      640,
+      768,
+      256,
+      384,
+      1024,
+      896,
+      256,
+      896,
+      512,
+      256,
+      896,
+      768,
+      256,
+      896,
+      896,
+      384,
+      896,
+      640,
+      768,
+      512,
+      768,
+      768,
+      1024,
+      768,
+      640,
+      768,
+      384,
+      256,
+      512,
+      896,
+      128,
+      384,
+      256,
+      768,
+      384,
+      256,
+      1024,
+      1024,
+      896,
+      256,
+      1024,
+      256,
+      128,
+      896
+    ],
+    [
+      640,
+      640,
+      896,
+      256,
+      1024,
+      512,
+      1024,
+      768,
+      384,
+      512,
+      256,
+      768,
+      896,
+      768,
+      512,
+      768,
+      768,
+      640,
+      384,
+      768,
+      512,
+      768,
+      768,
+      512,
+      768,
+      128,
+      896,
+      512,
+      768,
+      1024,
+      128,
+      384,
+      768,
+      768,
+      768,
+      384,
+      512,
+      640,
+      768,
+      512,
+      768,
+      1024,
+      640,
+      896,
+      256,
+      1024,
+      384,
+      768,
+      768,
+      768
+    ],
+    [
+      896,
+      512,
+      896,
+      768,
+      384,
+      384,
+      768,
+      512,
+      768,
+      512,
+      1024,
+      640,
+      896,
+      896,
+      256,
+      640,
+      1024,
+      256,
+      896,
+      128,
+      128,
+      128,
+      768,
+      896,
+      384,
+      896,
+      512,
+      896,
+      384,
+      256,
+      640,
+      640,
+      896,
+      768,
+      640,
+      256,
+      896,
+      896,
+      512,
+      128,
+      896,
+      256,
+      256,
+      640,
+      896,
+      896,
+      128,
+      1024,
+      256,
+      384,
+      1024,
+      640,
+      896
+    ],
+    [
+      768,
+      384,
+      640,
+      896,
+      256,
+      128,
+      384,
+      896,
+      128,
+      128,
+      896,
+      256,
+      384,
+      896,
+      512,
+      768,
+      768,
+      512,
+      512,
+      768,
+      896,
+      640,
+      768,
+      896,
+      896,
+      896,
+      640,
+      896,
+      640,
+      512,
+      896,
+      256,
+      512,
+      128,
+      512,
+      384,
+      768,
+      768,
+      1024,
+      256,
+      768,
+      256,
+      768,
+      512,
+      640,
+      1024,
+      128,
+      896,
+      896,
+      896,
+      896,
+      1024
+    ],
+    [
+      512,
+      384,
+      768,
+      640,
+      640,
+      768,
+      1024,
+      896,
+      512,
+      256,
+      640,
+      768,
+      640,
+      896,
+      128,
+      256,
+      896,
+      1024,
+      256,
+      640,
+      512,
+      256,
+      128,
+      512,
+      256,
+      640,
+      768,
+      768,
+      128,
+      128,
+      768,
+      640,
+      1024,
+      1024,
+      768,
+      512,
+      896,
+      768,
+      896,
+      1024,
+      896,
+      896,
+      512,
+      640,
+      1024,
+      512,
+      1024,
+      512,
+      512,
+      512,
+      768
+    ],
+    [
+      896,
+      896,
+      1024,
+      1024,
+      896,
+      128,
+      768,
+      256,
+      1024,
+      256,
+      1024,
+      640,
+      384,
+      256,
+      256,
+      512,
+      768,
+      896,
+      512,
+      768,
+      384,
+      1024,
+      896,
+      896,
+      1024,
+      896,
+      768,
+      896,
+      640,
+      1024,
+      512,
+      896,
+      512,
+      1024,
+      512,
+      512,
+      256,
+      256,
+      256,
+      512,
+      768,
+      128,
+      384,
+      512,
+      896,
+      896,
+      1024,
+      256,
+      384,
+      640
+    ],
+    [
+      896,
+      640,
+      384,
+      512,
+      256,
+      640,
+      1024,
+      384,
+      1024,
+      1024,
+      768,
+      256,
+      1024,
+      768,
+      512,
+      896,
+      256,
+      1024,
+      768,
+      768,
+      768,
+      384,
+      384,
+      256,
+      1024,
+      384,
+      384,
+      384,
+      896,
+      768,
+      640,
+      768,
+      512,
+      896,
+      896,
+      896,
+      896,
+      256,
+      384,
+      128,
+      1024,
+      896,
+      256,
+      256,
+      768,
+      640,
+      896,
+      384,
+      768,
+      512,
+      640
+    ],
+    [
+      896,
+      1024,
+      768,
+      1024,
+      896,
+      256,
+      768,
+      128,
+      128,
+      768,
+      512,
+      896,
+      384,
+      768,
+      1024,
+      256,
+      768,
+      768,
+      256,
+      512,
+      512,
+      640,
+      512,
+      256,
+      768,
+      896,
+      384,
+      1024,
+      640,
+      1024,
+      512,
+      512,
+      384,
+      512,
+      512,
+      1024,
+      384,
+      896,
+      768,
+      384,
+      384,
+      128,
+      384,
+      1024,
+      896,
+      640,
+      768,
+      768,
+      256,
+      640,
+      512,
+      640,
+      384
+    ]
+  ],
+  "glean_metadata": {
+    "base_model": "allenai/OLMoE-1B-7B-0125-Instruct",
+    "block_size": 128,
+    "criterion": "reap",
+    "dead_experts": 217,
+    "keep_fraction": 0.5,
+    "min_width": 128,
+    "params": 3697491968,
+    "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
+}

healed/mixdistill_smoke/step0006/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

healed/mixdistill_smoke/step0006/generation_config.json

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

healed/mixdistill_smoke/step0006/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
+            )

healed/mixdistill_smoke/step0006/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
+  }
+}

healed/mixdistill_smoke/step0006/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
+}

healed/opd_warm_unleashed/step0300/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 %}

healed/opd_warm_unleashed/step0300/config.json

@@ -0,0 +1,887 @@
+{
+  "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": [
+    [
+      1024,
+      384,
+      256,
+      256,
+      768,
+      1024,
+      128,
+      768,
+      384,
+      128,
+      1024,
+      384,
+      384,
+      128,
+      640,
+      896,
+      896,
+      256,
+      768,
+      128,
+      768,
+      768,
+      768,
+      768,
+      640,
+      896,
+      256,
+      128,
+      384,
+      896,
+      256,
+      768,
+      512,
+      640,
+      768,
+      896,
+      1024,
+      384,
+      384,
+      512,
+      384,
+      256,
+      768,
+      640,
+      896,
+      896,
+      512,
+      768,
+      512,
+      384,
+      1024,
+      896,
+      896,
+      768,
+      128
+    ],
+    [
+      768,
+      896,
+      896,
+      256,
+      896,
+      384,
+      768,
+      768,
+      1024,
+      384,
+      512,
+      1024,
+      768,
+      768,
+      1024,
+      640,
+      128,
+      768,
+      384,
+      768,
+      128,
+      256,
+      512,
+      384,
+      896,
+      256,
+      1024,
+      1024,
+      896,
+      1024,
+      896,
+      768,
+      256,
+      768,
+      128,
+      1024,
+      768,
+      768,
+      640,
+      384,
+      1024,
+      384,
+      640,
+      896,
+      1024,
+      512,
+      512,
+      512,
+      640
+    ],
+    [
+      768,
+      1024,
+      512,
+      384,
+      384,
+      768,
+      256,
+      384,
+      1024,
+      896,
+      640,
+      256,
+      1024,
+      1024,
+      1024,
+      640,
+      1024,
+      896,
+      512,
+      768,
+      128,
+      768,
+      1024,
+      768,
+      256,
+      1024,
+      896,
+      1024,
+      896,
+      512,
+      640,
+      384,
+      896,
+      384,
+      512,
+      640,
+      256,
+      768,
+      256,
+      768,
+      768,
+      640,
+      1024,
+      640,
+      896,
+      768,
+      768,
+      256
+    ],
+    [
+      640,
+      640,
+      896,
+      512,
+      640,
+      896,
+      896,
+      384,
+      896,
+      896,
+      384,
+      896,
+      640,
+      256,
+      640,
+      1024,
+      1024,
+      768,
+      1024,
+      896,
+      1024,
+      768,
+      896,
+      256,
+      512,
+      768,
+      1024,
+      256,
+      768,
+      512,
+      256,
+      640,
+      1024,
+      1024,
+      512,
+      1024,
+      768,
+      256,
+      1024,
+      384,
+      896,
+      1024,
+      896,
+      1024,
+      384
+    ],
+    [
+      256,
+      640,
+      640,
+      896,
+      768,
+      896,
+      768,
+      768,
+      896,
+      896,
+      1024,
+      256,
+      512,
+      1024,
+      640,
+      896,
+      512,
+      512,
+      384,
+      384,
+      256,
+      384,
+      384,
+      896,
+      896,
+      768,
+      640,
+      896,
+      768,
+      1024,
+      512,
+      640,
+      512,
+      640,
+      896,
+      512,
+      512,
+      384,
+      640,
+      896,
+      896,
+      896,
+      1024,
+      640,
+      1024,
+      640,
+      1024
+    ],
+    [
+      1024,
+      512,
+      1024,
+      1024,
+      640,
+      896,
+      640,
+      1024,
+      896,
+      384,
+      1024,
+      128,
+      896,
+      768,
+      1024,
+      768,
+      640,
+      896,
+      768,
+      640,
+      512,
+      896,
+      512,
+      640,
+      256,
+      768,
+      640,
+      768,
+      384,
+      896,
+      512,
+      512,
+      256,
+      512,
+      896,
+      256,
+      384,
+      640,
+      512,
+      640,
+      896,
+      512,
+      1024,
+      256,
+      768,
+      1024,
+      768,
+      256,
+      256
+    ],
+    [
+      640,
+      896,
+      1024,
+      896,
+      1024,
+      1024,
+      1024,
+      512,
+      256,
+      256,
+      1024,
+      768,
+      512,
+      768,
+      1024,
+      1024,
+      1024,
+      384,
+      512,
+      1024,
+      512,
+      1024,
+      128,
+      640,
+      640,
+      896,
+      768,
+      128,
+      256,
+      256,
+      256,
+      256,
+      896,
+      1024,
+      1024,
+      384,
+      896,
+      256,
+      896,
+      640,
+      1024,
+      384,
+      640,
+      256,
+      1024,
+      1024,
+      1024
+    ],
+    [
+      1024,
+      384,
+      1024,
+      1024,
+      256,
+      128,
+      256,
+      384,
+      256,
+      384,
+      896,
+      768,
+      896,
+      896,
+      512,
+      896,
+      640,
+      384,
+      384,
+      896,
+      768,
+      384,
+      896,
+      768,
+      768,
+      512,
+      896,
+      768,
+      768,
+      896,
+      128,
+      896,
+      512,
+      256,
+      768,
+      128,
+      384,
+      256,
+      896,
+      896,
+      384,
+      768,
+      512,
+      640,
+      256,
+      768,
+      640,
+      896,
+      384,
+      512,
+      1024,
+      768,
+      384
+    ],
+    [
+      512,
+      768,
+      512,
+      256,
+      128,
+      640,
+      384,
+      640,
+      768,
+      896,
+      640,
+      768,
+      256,
+      384,
+      1024,
+      896,
+      256,
+      896,
+      512,
+      256,
+      896,
+      768,
+      256,
+      896,
+      896,
+      384,
+      896,
+      640,
+      768,
+      512,
+      768,
+      768,
+      1024,
+      768,
+      640,
+      768,
+      384,
+      256,
+      512,
+      896,
+      128,
+      384,
+      256,
+      768,
+      384,
+      256,
+      1024,
+      1024,
+      896,
+      256,
+      1024,
+      256,
+      128,
+      896
+    ],
+    [
+      640,
+      640,
+      896,
+      256,
+      1024,
+      512,
+      1024,
+      768,
+      384,
+      512,
+      256,
+      768,
+      896,
+      768,
+      512,
+      768,
+      768,
+      640,
+      384,
+      768,
+      512,
+      768,
+      768,
+      512,
+      768,
+      128,
+      896,
+      512,
+      768,
+      1024,
+      128,
+      384,
+      768,
+      768,
+      768,
+      384,
+      512,
+      640,
+      768,
+      512,
+      768,
+      1024,
+      640,
+      896,
+      256,
+      1024,
+      384,
+      768,
+      768,
+      768
+    ],
+    [
+      896,
+      512,
+      896,
+      768,
+      384,
+      384,
+      768,
+      512,
+      768,
+      512,
+      1024,
+      640,
+      896,
+      896,
+      256,
+      640,
+      1024,
+      256,
+      896,
+      128,
+      128,
+      128,
+      768,
+      896,
+      384,
+      896,
+      512,
+      896,
+      384,
+      256,
+      640,
+      640,
+      896,
+      768,
+      640,
+      256,
+      896,
+      896,
+      512,
+      128,
+      896,
+      256,
+      256,
+      640,
+      896,
+      896,
+      128,
+      1024,
+      256,
+      384,
+      1024,
+      640,
+      896
+    ],
+    [
+      768,
+      384,
+      640,
+      896,
+      256,
+      128,
+      384,
+      896,
+      128,
+      128,
+      896,
+      256,
+      384,
+      896,
+      512,
+      768,
+      768,
+      512,
+      512,
+      768,
+      896,
+      640,
+      768,
+      896,
+      896,
+      896,
+      640,
+      896,
+      640,
+      512,
+      896,
+      256,
+      512,
+      128,
+      512,
+      384,
+      768,
+      768,
+      1024,
+      256,
+      768,
+      256,
+      768,
+      512,
+      640,
+      1024,
+      128,
+      896,
+      896,
+      896,
+      896,
+      1024
+    ],
+    [
+      512,
+      384,
+      768,
+      640,
+      640,
+      768,
+      1024,
+      896,
+      512,
+      256,
+      640,
+      768,
+      640,
+      896,
+      128,
+      256,
+      896,
+      1024,
+      256,
+      640,
+      512,
+      256,
+      128,
+      512,
+      256,
+      640,
+      768,
+      768,
+      128,
+      128,
+      768,
+      640,
+      1024,
+      1024,
+      768,
+      512,
+      896,
+      768,
+      896,
+      1024,
+      896,
+      896,
+      512,
+      640,
+      1024,
+      512,
+      1024,
+      512,
+      512,
+      512,
+      768
+    ],
+    [
+      896,
+      896,
+      1024,
+      1024,
+      896,
+      128,
+      768,
+      256,
+      1024,
+      256,
+      1024,
+      640,
+      384,
+      256,
+      256,
+      512,
+      768,
+      896,
+      512,
+      768,
+      384,
+      1024,
+      896,
+      896,
+      1024,
+      896,
+      768,
+      896,
+      640,
+      1024,
+      512,
+      896,
+      512,
+      1024,
+      512,
+      512,
+      256,
+      256,
+      256,
+      512,
+      768,
+      128,
+      384,
+      512,
+      896,
+      896,
+      1024,
+      256,
+      384,
+      640
+    ],
+    [
+      896,
+      640,
+      384,
+      512,
+      256,
+      640,
+      1024,
+      384,
+      1024,
+      1024,
+      768,
+      256,
+      1024,
+      768,
+      512,
+      896,
+      256,
+      1024,
+      768,
+      768,
+      768,
+      384,
+      384,
+      256,
+      1024,
+      384,
+      384,
+      384,
+      896,
+      768,
+      640,
+      768,
+      512,
+      896,
+      896,
+      896,
+      896,
+      256,
+      384,
+      128,
+      1024,
+      896,
+      256,
+      256,
+      768,
+      640,
+      896,
+      384,
+      768,
+      512,
+      640
+    ],
+    [
+      896,
+      1024,
+      768,
+      1024,
+      896,
+      256,
+      768,
+      128,
+      128,
+      768,
+      512,
+      896,
+      384,
+      768,
+      1024,
+      256,
+      768,
+      768,
+      256,
+      512,
+      512,
+      640,
+      512,
+      256,
+      768,
+      896,
+      384,
+      1024,
+      640,
+      1024,
+      512,
+      512,
+      384,
+      512,
+      512,
+      1024,
+      384,
+      896,
+      768,
+      384,
+      384,
+      128,
+      384,
+      1024,
+      896,
+      640,
+      768,
+      768,
+      256,
+      640,
+      512,
+      640,
+      384
+    ]
+  ],
+  "glean_metadata": {
+    "base_model": "allenai/OLMoE-1B-7B-0125-Instruct",
+    "block_size": 128,
+    "criterion": "reap",
+    "dead_experts": 217,
+    "keep_fraction": 0.5,
+    "min_width": 128,
+    "params": 3697491968,
+    "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
+}

healed/opd_warm_unleashed/step0300/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

healed/opd_warm_unleashed/step0300/generation_config.json

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

healed/opd_warm_unleashed/step0300/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
+            )

healed/opd_warm_unleashed/step0300/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
+  }
+}

healed/opd_warm_unleashed/step0300/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
+}

healed/opd_warm_unleashed/step0350/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 %}

healed/opd_warm_unleashed/step0350/config.json

@@ -0,0 +1,887 @@
+{
+  "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": [
+    [
+      1024,
+      384,
+      256,
+      256,
+      768,
+      1024,
+      128,
+      768,
+      384,
+      128,
+      1024,
+      384,
+      384,
+      128,
+      640,
+      896,
+      896,
+      256,
+      768,
+      128,
+      768,
+      768,
+      768,
+      768,
+      640,
+      896,
+      256,
+      128,
+      384,
+      896,
+      256,
+      768,
+      512,
+      640,
+      768,
+      896,
+      1024,
+      384,
+      384,
+      512,
+      384,
+      256,
+      768,
+      640,
+      896,
+      896,
+      512,
+      768,
+      512,
+      384,
+      1024,
+      896,
+      896,
+      768,
+      128
+    ],
+    [
+      768,
+      896,
+      896,
+      256,
+      896,
+      384,
+      768,
+      768,
+      1024,
+      384,
+      512,
+      1024,
+      768,
+      768,
+      1024,
+      640,
+      128,
+      768,
+      384,
+      768,
+      128,
+      256,
+      512,
+      384,
+      896,
+      256,
+      1024,
+      1024,
+      896,
+      1024,
+      896,
+      768,
+      256,
+      768,
+      128,
+      1024,
+      768,
+      768,
+      640,
+      384,
+      1024,
+      384,
+      640,
+      896,
+      1024,
+      512,
+      512,
+      512,
+      640
+    ],
+    [
+      768,
+      1024,
+      512,
+      384,
+      384,
+      768,
+      256,
+      384,
+      1024,
+      896,
+      640,
+      256,
+      1024,
+      1024,
+      1024,
+      640,
+      1024,
+      896,
+      512,
+      768,
+      128,
+      768,
+      1024,
+      768,
+      256,
+      1024,
+      896,
+      1024,
+      896,
+      512,
+      640,
+      384,
+      896,
+      384,
+      512,
+      640,
+      256,
+      768,
+      256,
+      768,
+      768,
+      640,
+      1024,
+      640,
+      896,
+      768,
+      768,
+      256
+    ],
+    [
+      640,
+      640,
+      896,
+      512,
+      640,
+      896,
+      896,
+      384,
+      896,
+      896,
+      384,
+      896,
+      640,
+      256,
+      640,
+      1024,
+      1024,
+      768,
+      1024,
+      896,
+      1024,
+      768,
+      896,
+      256,
+      512,
+      768,
+      1024,
+      256,
+      768,
+      512,
+      256,
+      640,
+      1024,
+      1024,
+      512,
+      1024,
+      768,
+      256,
+      1024,
+      384,
+      896,
+      1024,
+      896,
+      1024,
+      384
+    ],
+    [
+      256,
+      640,
+      640,
+      896,
+      768,
+      896,
+      768,
+      768,
+      896,
+      896,
+      1024,
+      256,
+      512,
+      1024,
+      640,
+      896,
+      512,
+      512,
+      384,
+      384,
+      256,
+      384,
+      384,
+      896,
+      896,
+      768,
+      640,
+      896,
+      768,
+      1024,
+      512,
+      640,
+      512,
+      640,
+      896,
+      512,
+      512,
+      384,
+      640,
+      896,
+      896,
+      896,
+      1024,
+      640,
+      1024,
+      640,
+      1024
+    ],
+    [
+      1024,
+      512,
+      1024,
+      1024,
+      640,
+      896,
+      640,
+      1024,
+      896,
+      384,
+      1024,
+      128,
+      896,
+      768,
+      1024,
+      768,
+      640,
+      896,
+      768,
+      640,
+      512,
+      896,
+      512,
+      640,
+      256,
+      768,
+      640,
+      768,
+      384,
+      896,
+      512,
+      512,
+      256,
+      512,
+      896,
+      256,
+      384,
+      640,
+      512,
+      640,
+      896,
+      512,
+      1024,
+      256,
+      768,
+      1024,
+      768,
+      256,
+      256
+    ],
+    [
+      640,
+      896,
+      1024,
+      896,
+      1024,
+      1024,
+      1024,
+      512,
+      256,
+      256,
+      1024,
+      768,
+      512,
+      768,
+      1024,
+      1024,
+      1024,
+      384,
+      512,
+      1024,
+      512,
+      1024,
+      128,
+      640,
+      640,
+      896,
+      768,
+      128,
+      256,
+      256,
+      256,
+      256,
+      896,
+      1024,
+      1024,
+      384,
+      896,
+      256,
+      896,
+      640,
+      1024,
+      384,
+      640,
+      256,
+      1024,
+      1024,
+      1024
+    ],
+    [
+      1024,
+      384,
+      1024,
+      1024,
+      256,
+      128,
+      256,
+      384,
+      256,
+      384,
+      896,
+      768,
+      896,
+      896,
+      512,
+      896,
+      640,
+      384,
+      384,
+      896,
+      768,
+      384,
+      896,
+      768,
+      768,
+      512,
+      896,
+      768,
+      768,
+      896,
+      128,
+      896,
+      512,
+      256,
+      768,
+      128,
+      384,
+      256,
+      896,
+      896,
+      384,
+      768,
+      512,
+      640,
+      256,
+      768,
+      640,
+      896,
+      384,
+      512,
+      1024,
+      768,
+      384
+    ],
+    [
+      512,
+      768,
+      512,
+      256,
+      128,
+      640,
+      384,
+      640,
+      768,
+      896,
+      640,
+      768,
+      256,
+      384,
+      1024,
+      896,
+      256,
+      896,
+      512,
+      256,
+      896,
+      768,
+      256,
+      896,
+      896,
+      384,
+      896,
+      640,
+      768,
+      512,
+      768,
+      768,
+      1024,
+      768,
+      640,
+      768,
+      384,
+      256,
+      512,
+      896,
+      128,
+      384,
+      256,
+      768,
+      384,
+      256,
+      1024,
+      1024,
+      896,
+      256,
+      1024,
+      256,
+      128,
+      896
+    ],
+    [
+      640,
+      640,
+      896,
+      256,
+      1024,
+      512,
+      1024,
+      768,
+      384,
+      512,
+      256,
+      768,
+      896,
+      768,
+      512,
+      768,
+      768,
+      640,
+      384,
+      768,
+      512,
+      768,
+      768,
+      512,
+      768,
+      128,
+      896,
+      512,
+      768,
+      1024,
+      128,
+      384,
+      768,
+      768,
+      768,
+      384,
+      512,
+      640,
+      768,
+      512,
+      768,
+      1024,
+      640,
+      896,
+      256,
+      1024,
+      384,
+      768,
+      768,
+      768
+    ],
+    [
+      896,
+      512,
+      896,
+      768,
+      384,
+      384,
+      768,
+      512,
+      768,
+      512,
+      1024,
+      640,
+      896,
+      896,
+      256,
+      640,
+      1024,
+      256,
+      896,
+      128,
+      128,
+      128,
+      768,
+      896,
+      384,
+      896,
+      512,
+      896,
+      384,
+      256,
+      640,
+      640,
+      896,
+      768,
+      640,
+      256,
+      896,
+      896,
+      512,
+      128,
+      896,
+      256,
+      256,
+      640,
+      896,
+      896,
+      128,
+      1024,
+      256,
+      384,
+      1024,
+      640,
+      896
+    ],
+    [
+      768,
+      384,
+      640,
+      896,
+      256,
+      128,
+      384,
+      896,
+      128,
+      128,
+      896,
+      256,
+      384,
+      896,
+      512,
+      768,
+      768,
+      512,
+      512,
+      768,
+      896,
+      640,
+      768,
+      896,
+      896,
+      896,
+      640,
+      896,
+      640,
+      512,
+      896,
+      256,
+      512,
+      128,
+      512,
+      384,
+      768,
+      768,
+      1024,
+      256,
+      768,
+      256,
+      768,
+      512,
+      640,
+      1024,
+      128,
+      896,
+      896,
+      896,
+      896,
+      1024
+    ],
+    [
+      512,
+      384,
+      768,
+      640,
+      640,
+      768,
+      1024,
+      896,
+      512,
+      256,
+      640,
+      768,
+      640,
+      896,
+      128,
+      256,
+      896,
+      1024,
+      256,
+      640,
+      512,
+      256,
+      128,
+      512,
+      256,
+      640,
+      768,
+      768,
+      128,
+      128,
+      768,
+      640,
+      1024,
+      1024,
+      768,
+      512,
+      896,
+      768,
+      896,
+      1024,
+      896,
+      896,
+      512,
+      640,
+      1024,
+      512,
+      1024,
+      512,
+      512,
+      512,
+      768
+    ],
+    [
+      896,
+      896,
+      1024,
+      1024,
+      896,
+      128,
+      768,
+      256,
+      1024,
+      256,
+      1024,
+      640,
+      384,
+      256,
+      256,
+      512,
+      768,
+      896,
+      512,
+      768,
+      384,
+      1024,
+      896,
+      896,
+      1024,
+      896,
+      768,
+      896,
+      640,
+      1024,
+      512,
+      896,
+      512,
+      1024,
+      512,
+      512,
+      256,
+      256,
+      256,
+      512,
+      768,
+      128,
+      384,
+      512,
+      896,
+      896,
+      1024,
+      256,
+      384,
+      640
+    ],
+    [
+      896,
+      640,
+      384,
+      512,
+      256,
+      640,
+      1024,
+      384,
+      1024,
+      1024,
+      768,
+      256,
+      1024,
+      768,
+      512,
+      896,
+      256,
+      1024,
+      768,
+      768,
+      768,
+      384,
+      384,
+      256,
+      1024,
+      384,
+      384,
+      384,
+      896,
+      768,
+      640,
+      768,
+      512,
+      896,
+      896,
+      896,
+      896,
+      256,
+      384,
+      128,
+      1024,
+      896,
+      256,
+      256,
+      768,
+      640,
+      896,
+      384,
+      768,
+      512,
+      640
+    ],
+    [
+      896,
+      1024,
+      768,
+      1024,
+      896,
+      256,
+      768,
+      128,
+      128,
+      768,
+      512,
+      896,
+      384,
+      768,
+      1024,
+      256,
+      768,
+      768,
+      256,
+      512,
+      512,
+      640,
+      512,
+      256,
+      768,
+      896,
+      384,
+      1024,
+      640,
+      1024,
+      512,
+      512,
+      384,
+      512,
+      512,
+      1024,
+      384,
+      896,
+      768,
+      384,
+      384,
+      128,
+      384,
+      1024,
+      896,
+      640,
+      768,
+      768,
+      256,
+      640,
+      512,
+      640,
+      384
+    ]
+  ],
+  "glean_metadata": {
+    "base_model": "allenai/OLMoE-1B-7B-0125-Instruct",
+    "block_size": 128,
+    "criterion": "reap",
+    "dead_experts": 217,
+    "keep_fraction": 0.5,
+    "min_width": 128,
+    "params": 3697491968,
+    "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
+}

healed/opd_warm_unleashed/step0350/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

healed/opd_warm_unleashed/step0350/generation_config.json

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

healed/opd_warm_unleashed/step0350/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
+            )

healed/opd_warm_unleashed/step0350/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
+  }
+}