feat: Add OLMoE-1B-7B-0924 support

litgpt/config.py

@@ -1064,6 +1064,44 @@ def check_indicator_and_length(
         copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
         configs.append(copy)
 
+####################
+# Allen AI OLMoE
+####################
+olmoe = [
+    # https://huggingface.co/allenai/OLMoE-1B-7B-0924/blob/main/config.json
+    dict(
+        name="OLMoE-1B-7B-0924{}",
+        hf_config=dict(org="allenai", name="OLMoE-1B-7B-0924{}"),
+        block_size=4096,
+        vocab_size=50304,
+        padded_vocab_size=50304,
+        n_layer=16,
+        n_head=16,
+        n_embd=2048,
+        n_query_groups=8,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        norm_class_name="RMSNorm",
+        norm_eps=1e-5,
+        mlp_class_name="LLaMAMoE",
+        intermediate_size=2048,  # dense fallback (required by Config.__post_init__)
+        moe_intermediate_size=1024,  # per-expert hidden dim
+        rope_base=10000,
+        n_expert=64,
+        n_expert_per_token=8,
+        norm_topk_prob=False,
+        norm_qk=True,
+        norm_qk_type="olmo2",
+    ),
+]
+for c in olmoe:
+    for kind in ("", "-Instruct", "-SFT"):
+        copy = deepcopy(c)
+        copy["name"] = c["name"].format(kind)
+        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
+        configs.append(copy)
+
 ###############
 # Google Gemma
 ###############
@@ -2022,6 +2060,7 @@ def check_indicator_and_length(
         rope_base=1000000,
         n_expert=8,
         n_expert_per_token=2,
+        norm_topk_prob=True,
     ),
     # https://huggingface.co/mistralai/Mixtral-8x22B-Instruct-v0.1/blob/main/config.json
     dict(
@@ -2043,6 +2082,7 @@ def check_indicator_and_length(
         rope_base=1000000,
         n_expert=8,
         n_expert_per_token=2,
+        norm_topk_prob=True,
     ),
 ]
 for c in mistral:
@@ -2880,6 +2920,7 @@ def check_indicator_and_length(
         norm_qk=True,
         n_expert=128,
         n_expert_per_token=8,
+        norm_topk_prob=True,
     ),
     # https://huggingface.co/Qwen/Qwen3-30B-A3B-Base/blob/main/config.json
     dict(
@@ -2905,6 +2946,7 @@ def check_indicator_and_length(
         norm_qk=True,
         n_expert=128,
         n_expert_per_token=8,
+        norm_topk_prob=True,
     ),
     # https://huggingface.co/Qwen/Qwen3-235B-A22B/blob/main/config.json
     dict(
@@ -2930,6 +2972,7 @@ def check_indicator_and_length(
         norm_qk=True,
         n_expert=128,
         n_expert_per_token=8,
+        norm_topk_prob=True,
     ),
 ]
 configs.extend(qwen_3_moe)
@@ -2959,6 +3002,7 @@ def check_indicator_and_length(
         norm_qk=True,
         n_expert=128,
         n_expert_per_token=8,
+        norm_topk_prob=True,
     ),
     # https://huggingface.co/Qwen/Qwen3-30B-A3B-Thinking-2507/blob/main/config.json
     dict(
@@ -2984,6 +3028,7 @@ def check_indicator_and_length(
         norm_qk=True,
         n_expert=128,
         n_expert_per_token=8,
+        norm_topk_prob=True,
     ),
     # https://huggingface.co/Qwen/Qwen3-4B-Thinking-2507/blob/main/config.json
     dict(

litgpt/model.py

@@ -866,8 +866,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = x.view(-1, C)  # (B*T, C)
         if not self.config.n_expert_groups:
             router = self.gate(x)  # (B*T, n_expert)
-            probs, indices = torch.topk(router, self.config.n_expert_per_token)  # (B*T, n_expert_per_token)
-            probs = probs.softmax(dim=1, dtype=torch.float).to(dtype=x.dtype)
+            probs = F.softmax(router, dim=1, dtype=torch.float)
+            probs, indices = torch.topk(probs, self.config.n_expert_per_token)  # (B*T, n_expert_per_token)
+            if self.config.norm_topk_prob:
+                probs = probs / probs.sum(dim=-1, keepdim=True)
+            probs = probs.to(dtype=x.dtype)
         else:
             probs, indices = self.gate(x)
         if self.config.routed_scaling_factor != 1.0:

litgpt/prompts.py

@@ -413,6 +413,19 @@ def __init__(self):
         )
 
 
+class OLMoE(ChatML):
+    """Prompt style for OLMoE-1B-7B-0924-Instruct (ChatML template).
+
+    The instruct model was fine-tuned with the ChatML template:
+        <|im_start|>system\\n{sys}<|im_end|>
+        <|im_start|>user\\n{prompt}<|im_end|>
+        <|im_start|>assistant\\n
+    """
+
+    def __init__(self):
+        super().__init__("You are a helpful assistant.")
+
+
 # Maps prompt style names to PromptStyle classes
 prompt_styles: dict[str, type[PromptStyle]] = {
     # Dataset-specific prompt styles
@@ -446,6 +459,7 @@ def __init__(self):
     "qwen3": Qwen3,
     "smollm2": SmolLM2,
     "salamandra": Salamandra,
+    "olmoe": OLMoE,
 }
 
 
@@ -510,6 +524,8 @@ def model_name_to_prompt_style(model_name: str) -> PromptStyle:
         return SmolLM2()
     if re.search(r"salamandra-.*-instruct", model_name):
         return Salamandra()
+    if re.search(r"OLMoE.*-Instruct", model_name):
+        return OLMoE()
     return Default()
 
 

litgpt/scripts/convert_hf_checkpoint.py

@@ -638,6 +638,141 @@ def copy_weights_olmo2(
                 pbar.update(progress_per_file)
 
 
+def copy_weights_olmoe(
+    config: Config,
+    qkv_weights: dict[int, list[NotYetLoadedTensor | None]],
+    state_dict: dict[str, torch.Tensor],
+    hf_weights: dict[str, torch.Tensor | NotYetLoadedTensor],
+    saver: incremental_save | None = None,
+    dtype: torch.dtype | None = None,
+    pbar: tqdm | None = None,
+    progress_per_file: float | None = None,
+    debug_mode: bool | None = False,
+) -> None:
+    """Convert OLMoE-1B-7B-0924 HF weights to LitGPT format.
+
+    HF OLMoE stores per-expert weights individually:
+      model.layers.{L}.mlp.experts.{E}.gate_proj.weight
+      model.layers.{L}.mlp.experts.{E}.up_proj.weight
+      model.layers.{L}.mlp.experts.{E}.down_proj.weight
+
+    These map to LitGPT's individual expert Linear weights:
+      transformer.h.{L}.mlp.experts.{E}.fc_1.weight  (gate)
+      transformer.h.{L}.mlp.experts.{E}.fc_2.weight  (up)
+      transformer.h.{L}.mlp.experts.{E}.proj.weight  (down)
+
+    HF OLMoE also has per-head q_norm / k_norm weights that map to
+    LitGPT's attn.norm_q / attn.norm_k.
+    """
+    weight_map = {
+        "model.embed_tokens.weight": "transformer.wte.weight",
+        "model.layers.{}.input_layernorm.weight": "transformer.h.{}.norm_1.weight",
+        "model.layers.{}.self_attn.q_proj.weight": None,  # merged into qkv below
+        "model.layers.{}.self_attn.k_proj.weight": None,
+        "model.layers.{}.self_attn.v_proj.weight": None,
+        "model.layers.{}.self_attn.q_norm.weight": "transformer.h.{}.attn.norm_q.weight",
+        "model.layers.{}.self_attn.k_norm.weight": "transformer.h.{}.attn.norm_k.weight",
+        "model.layers.{}.self_attn.o_proj.weight": "transformer.h.{}.attn.proj.weight",
+        "model.layers.{}.self_attn.rotary_emb.inv_freq": None,
+        "model.layers.{}.post_attention_layernorm.weight": "transformer.h.{}.norm_2.weight",
+        "model.layers.{}.mlp.gate.weight": "transformer.h.{}.mlp.gate.weight",
+        "model.norm.weight": "transformer.ln_f.weight",
+        "lm_head.weight": "lm_head.weight",
+    }
+
+    if progress_per_file is not None:
+        progress_per_file = progress_per_file / max(1, len(hf_weights) + len(qkv_weights))
+
+    for from_name, param in hf_weights.items():
+        param = load_param(param, from_name, dtype, verbose=debug_mode)
+
+        # ── per-expert gate_proj → fc_1 ──
+        if ".mlp.experts." in from_name and from_name.endswith("gate_proj.weight"):
+            layer_idx, expert_idx = from_name.split(".")[2], from_name.split(".")[5]
+            to_key = f"transformer.h.{layer_idx}.mlp.experts.{expert_idx}.fc_1.weight"
+            state_dict[to_key] = param
+            if progress_per_file is not None:
+                pbar.update(progress_per_file)
+            continue
+
+        # ── per-expert up_proj → fc_2 ──
+        if ".mlp.experts." in from_name and from_name.endswith("up_proj.weight"):
+            layer_idx, expert_idx = from_name.split(".")[2], from_name.split(".")[5]
+            to_key = f"transformer.h.{layer_idx}.mlp.experts.{expert_idx}.fc_2.weight"
+            state_dict[to_key] = param
+            if progress_per_file is not None:
+                pbar.update(progress_per_file)
+            continue
+
+        # ── per-expert down_proj → proj ──
+        if ".mlp.experts." in from_name and from_name.endswith("down_proj.weight"):
+            layer_idx, expert_idx = from_name.split(".")[2], from_name.split(".")[5]
+            to_key = f"transformer.h.{layer_idx}.mlp.experts.{expert_idx}.proj.weight"
+            state_dict[to_key] = param
+            if progress_per_file is not None:
+                pbar.update(progress_per_file)
+            continue
+
+        # ── fused gate_up_proj: shape (n_expert, 2*intermediate, n_embd) ──
+        if from_name.endswith("mlp.experts.gate_up_proj"):
+            name_template, *ids = layer_template(from_name, num_matches=1)
+            layer_idx = ids[0]
+            gate, up = param.chunk(2, dim=1)
+            for expert_idx in range(config.n_expert):
+                fc1_key = f"transformer.h.{layer_idx}.mlp.experts.{expert_idx}.fc_1.weight"
+                fc2_key = f"transformer.h.{layer_idx}.mlp.experts.{expert_idx}.fc_2.weight"
+                state_dict[fc1_key] = gate[expert_idx]
+                state_dict[fc2_key] = up[expert_idx]
+            if progress_per_file is not None:
+                pbar.update(progress_per_file)
+            continue
+
+        # ── fused down_proj: shape (n_expert, n_embd, moe_intermediate_size) ──
+        if from_name.endswith("mlp.experts.down_proj"):
+            name_template, *ids = layer_template(from_name, num_matches=1)
+            layer_idx = ids[0]
+            for expert_idx in range(config.n_expert):
+                proj_key = f"transformer.h.{layer_idx}.mlp.experts.{expert_idx}.proj.weight"
+                state_dict[proj_key] = param[expert_idx]
+            if progress_per_file is not None:
+                pbar.update(progress_per_file)
+            continue
+
+        # ── q / k / v projections: accumulate for qkv merge ──
+        name_template, *ids = layer_template(from_name, num_matches=1)
+        if any(w in from_name for w in ("q_proj", "k_proj", "v_proj")):
+            qkv = qkv_weights.setdefault(ids[0], defaultdict(dict))
+            weight_name, weight_type = from_name.split(".")[-2:]
+            qkv[weight_type][weight_name] = param
+
+        to_name = weight_map.get(name_template)
+        if to_name is None:
+            if progress_per_file is not None:
+                pbar.update(progress_per_file)
+            continue
+
+        to_name = to_name.format(*ids)
+        if saver is not None:
+            param = saver.store_early(param)
+        state_dict[to_name] = param
+
+        if progress_per_file is not None:
+            pbar.update(progress_per_file)
+
+    for i in list(qkv_weights):
+        for weight_type in list(qkv_weights[i]):
+            qkv = qkv_weights[i][weight_type]
+            if len(qkv) != 3:
+                continue
+            q = load_param(qkv["q_proj"], f"layer {i} q {weight_type}", dtype, verbose=debug_mode)
+            k = load_param(qkv["k_proj"], f"layer {i} k {weight_type}", dtype, verbose=debug_mode)
+            v = load_param(qkv["v_proj"], f"layer {i} v {weight_type}", dtype, verbose=debug_mode)
+            state_dict[f"transformer.h.{i}.attn.qkv.{weight_type}"] = torch.cat((q, k, v))
+            del qkv_weights[i][weight_type]
+            if progress_per_file is not None:
+                pbar.update(progress_per_file)
+
+
 def copy_weights_qwen_3(
     config: Config,
     qkv_weights: dict[int, list[NotYetLoadedTensor | None]],
@@ -814,6 +949,9 @@ def convert_hf_checkpoint(
         # holder to reconstitute the split q, k, v
         qkv_weights = {}
         copy_fn = partial(copy_weights_qwen_2_5, config, qkv_weights)
+    elif model_name.lower().startswith("olmoe"):
+        qkv_weights = {}
+        copy_fn = partial(copy_weights_olmoe, config, qkv_weights)
     elif model_name.lower().startswith("olmo-2-"):
         # holder to reconstitute the split q, k, v
         qkv_weights = {}

litgpt/scripts/convert_lit_checkpoint.py

@@ -450,6 +450,67 @@ def copy_weights_olmo2(
             state_dict[to_name] = param
 
 
+def copy_weights_olmoe(
+    config: Config,
+    state_dict: dict[str, torch.Tensor],
+    lit_weights: dict[str, torch.Tensor | NotYetLoadedTensor],
+    untie_weights: bool = False,
+    saver: incremental_save | None = None,
+) -> None:
+    """Convert OLMoE-1B-7B-0924 LitGPT weights back to HuggingFace format.
+
+    Splits the merged qkv tensor into separate q_proj / k_proj / v_proj
+    and maps expert weights back to OLMoE HF naming.
+    """
+    weight_map = {
+        "transformer.wte.weight": "model.embed_tokens.weight",
+        "transformer.h.{}.norm_1.weight": "model.layers.{}.input_layernorm.weight",
+        "transformer.h.{}.attn.proj.weight": "model.layers.{}.self_attn.o_proj.weight",
+        "transformer.h.{}.norm_2.weight": "model.layers.{}.post_attention_layernorm.weight",
+        "transformer.h.{}.mlp.gate.weight": "model.layers.{}.mlp.gate.weight",
+        "transformer.h.{}.attn.norm_q.weight": "model.layers.{}.self_attn.q_norm.weight",
+        "transformer.h.{}.attn.norm_k.weight": "model.layers.{}.self_attn.k_norm.weight",
+        # expert weights — num_matches=2 to capture layer_idx AND expert_idx
+        "transformer.h.{}.mlp.experts.{}.fc_1.weight": "model.layers.{}.mlp.experts.{}.gate_proj.weight",
+        "transformer.h.{}.mlp.experts.{}.fc_2.weight": "model.layers.{}.mlp.experts.{}.up_proj.weight",
+        "transformer.h.{}.mlp.experts.{}.proj.weight": "model.layers.{}.mlp.experts.{}.down_proj.weight",
+        "transformer.ln_f.weight": "model.norm.weight",
+        "lm_head.weight": "lm_head.weight",
+    }
+
+    for from_name, param in lit_weights.items():
+        if from_name == "lm_head.weight" and untie_weights:
+            continue
+        name_template, *ids = layer_template(from_name, num_matches=2)
+        param = load_param(param, from_name, None)
+
+        if from_name.endswith(".attn.qkv.weight"):
+            layer_idx = ids[0]
+            to_names = (
+                f"model.layers.{layer_idx}.self_attn.q_proj.weight",
+                f"model.layers.{layer_idx}.self_attn.k_proj.weight",
+                f"model.layers.{layer_idx}.self_attn.v_proj.weight",
+            )
+            params = param.split(
+                (
+                    config.n_head * config.head_size,
+                    config.n_query_groups * config.head_size,
+                    config.n_query_groups * config.head_size,
+                )
+            )
+        else:
+            to_name = weight_map.get(name_template)
+            if to_name is None:
+                continue
+            to_names = (to_name.format(*ids),)
+            params = (param,)
+
+        for to_name, p in zip(to_names, params):
+            if saver is not None:
+                p = saver.store_early(p)
+            state_dict[to_name] = p
+
+
 def copy_weights_qwen_3(
     config: Config,
     state_dict: dict[str, torch.Tensor],
@@ -562,6 +623,8 @@ def convert_lit_checkpoint(checkpoint_dir: Path, output_dir: Path) -> None:
         copy_fn = partial(copy_weights_phi, config)
     elif config.name.lower().startswith(("qwen2.5", "qwq")):
         copy_fn = partial(copy_weights_qwen_2_5, config)
+    elif config.name.lower().startswith("olmoe"):
+        copy_fn = partial(copy_weights_olmoe, config)
     elif config.name.lower().startswith("olmo-2-"):
         copy_fn = partial(copy_weights_olmo2, config)
     elif config.name.lower().startswith("qwen3"):