diff --git a/litgpt/config.py b/litgpt/config.py index bdba7eeac7..5088a398c5 100644 --- a/litgpt/config.py +++ b/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( diff --git a/litgpt/model.py b/litgpt/model.py index 541860ab5b..73becf5ecb 100644 --- a/litgpt/model.py +++ b/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: diff --git a/litgpt/prompts.py b/litgpt/prompts.py index 1998bc7eeb..0e8d9621e4 100644 --- a/litgpt/prompts.py +++ b/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() diff --git a/litgpt/scripts/convert_hf_checkpoint.py b/litgpt/scripts/convert_hf_checkpoint.py index 7fc96f7e55..cb033284da 100644 --- a/litgpt/scripts/convert_hf_checkpoint.py +++ b/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 = {} diff --git a/litgpt/scripts/convert_lit_checkpoint.py b/litgpt/scripts/convert_lit_checkpoint.py index a016615d73..c508aa330a 100644 --- a/litgpt/scripts/convert_lit_checkpoint.py +++ b/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"): diff --git a/tests/test_model.py b/tests/test_model.py index 8d0cf21d5e..250b1ec442 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -30,6 +30,7 @@ from transformers.models.mixtral import MixtralConfig, MixtralForCausalLM from transformers.models.olmo import OlmoConfig, OlmoForCausalLM from transformers.models.olmo2 import Olmo2Config, Olmo2ForCausalLM +from transformers.models.olmoe import OlmoeConfig, OlmoeForCausalLM from transformers.models.qwen2 import Qwen2Config, Qwen2ForCausalLM from transformers.models.qwen3 import Qwen3Config, Qwen3ForCausalLM from transformers.models.qwen3_moe import Qwen3MoeConfig, Qwen3MoeForCausalLM @@ -44,6 +45,7 @@ copy_weights_gpt_neox, copy_weights_hf_llama, copy_weights_olmo2, + copy_weights_olmoe, copy_weights_phi, copy_weights_qwen_2_5, copy_weights_qwen_3, @@ -697,6 +699,61 @@ def test_against_olmo2(model_name, device, dtype): torch.testing.assert_close(ours_y, theirs_y) +@torch.inference_mode() +def test_against_hf_olmoe(): + """End-to-end numerical equivalence test: HF OLMoE → LitGPT → same logits.""" + device = torch.device("cpu") + dtype = torch.float32 + + # Use a tiny version of the real OLMoE-1B-7B-0924 config so it runs fast. + # n_expert is shrunk to 4 (real model has 64) for speed; everything else + # matches the actual architecture (GQA 16/8, RMSNorm, SwiGLU MoE). + ours_config = Config.from_name( + "OLMoE-1B-7B-0924", + padded_vocab_size=10000, + n_layer=2, + n_embd=32, + n_head=8, + n_query_groups=4, + intermediate_size=32, # dense fallback (not used in MoE blocks) + moe_intermediate_size=16, # per-expert hidden dim + n_expert=4, + n_expert_per_token=2, + ) + T = 5 + + theirs_config = OlmoeConfig( + vocab_size=ours_config.padded_vocab_size, + hidden_size=ours_config.n_embd, + num_hidden_layers=ours_config.n_layer, + num_attention_heads=ours_config.n_head, + num_key_value_heads=ours_config.n_query_groups, + intermediate_size=ours_config.moe_intermediate_size, + max_position_embeddings=T, + rms_norm_eps=ours_config.norm_eps, + rope_theta=ours_config.rope_base, + num_experts=ours_config.n_expert, + num_experts_per_tok=ours_config.n_expert_per_token, + attention_bias=ours_config.bias, + ) + + theirs_model = OlmoeForCausalLM(theirs_config).to(device) + theirs_state_dict = theirs_model.state_dict() + + state_dict = {} + copy_weights_olmoe(ours_config, {}, state_dict, theirs_state_dict) + + ours_model = GPT(ours_config).to(device) + keys = ours_model.load_state_dict(state_dict, strict=False) + assert not keys.unexpected_keys + + x = torch.tensor([[9856, 23, 491, 1536, 304], [23, 345, 65, 123, 321]], dtype=torch.int32, device=device) + assert x.size(1) == T + ours_y = ours_model(x) + theirs_y = theirs_model(x)["logits"].to(dtype) + torch.testing.assert_close(ours_y, theirs_y) + + @torch.inference_mode() @pytest.mark.parametrize( ("device", "dtype"),