Request for Official HF → OLMoE Checkpoint Conversion Script

[aditi184]

Hi OLMo team,

I'm currently working on converting a Hugging Face model (allenai/OLMoE-1B-7B-0924-Instruct) into OLMo/OLMoE's pretraining checkpoint format to resume pretraining. While I was able to convert the model weights, I encountered issues with missing metadata.json, which seems to be a critical component for loading the checkpoint in restore_checkpoint().

After examining the metadata.json generated when training OLMo from scratch, I realized that it contains non-trivial fields and additional metadata about the architecture and checkpoint format. Reconstructing this file seems error-prone, and I couldn't find documentation or scripts for safely performing this conversion.

Would it be possible for you to provide an official Hugging Face → OLMo checkpoint conversion script, or any guidance on the exact format required for metadata.json? This would be extremely helpful for me as I am working on a project that requires continual pretraining of the OLMoE- Instruct model using OLMo from existing HF checkpoints.

Alternately, it would also be helpful if you can share the OLMoE checkpoint that is accepted by the pretraining script (https://github.com/allenai/OLMo/blob/Muennighoff/MoE/scripts/train.py)

Thanks for your help!

[Muennighoff]

I realized that it contains non-trivial fields and additional metadata about the architecture and checkpoint format

Can you share which you are referring to? Maybe you could also share the script you already wrote? I'm not sure atm if such a script exists but tagging some people in case they know cc @soldni @swj0419

[aditi184]

Hey @Muennighoff, I meant to load the pretrained checkpoint to continually pretrain it futher, the model is needed to be in the OLMo class format right?

I wrote the code to convert it to that class using the attached script.

Run script using python convert_hf_to_olmo.py --hf_model_path allenai/OLMoE-1B-7B-0924-Instruct --output_path /home/mila/k/khandela/scratch/ai2-llm/checkpoints/OLMoE/base --tokenizer_path allenai/OLMoE-1B-7B-0924

import os
import torch
import yaml
import argparse
from pathlib import Path
from transformers import OlmoeForCausalLM, AutoTokenizer

def convert_hf_to_olmo(hf_model_path, output_path, tokenizer_path=None):
    os.makedirs(output_path, exist_ok=True)
    print(f"Loading Hugging Face model from {hf_model_path}")
    model = OlmoeForCausalLM.from_pretrained(hf_model_path)
    state_dict = model.state_dict()
    
    olmo_state_dict = {}
    
    print("Reformatting weights to OLMo pretraining format...")
    for key, value in state_dict.items():
        new_key = key.replace("model.layers", "transformer.blocks").replace(".self_attn.", ".att_proj.")
        new_key = new_key.replace("q_proj", "q_proj").replace("k_proj", "k_proj").replace("v_proj", "v_proj")
        new_key = new_key.replace("o_proj", "attn_out").replace("gate_proj", "ffn.router.layer")
        new_key = new_key.replace("up_proj", "ffn.experts.mlp.v1").replace("down_proj", "ffn.experts.mlp.w2")
        new_key = new_key.replace("input_layernorm", "attn_norm").replace("post_attention_layernorm", "ff_norm")
        
        olmo_state_dict[new_key] = value
    
    torch.save(olmo_state_dict, os.path.join(output_path, "model.pt"))
    print(f"Converted model saved to {output_path}/model.pt")
    
    config = {
        "model": {
            "n_layers": model.config.num_hidden_layers,
            "n_heads": model.config.num_attention_heads,
            "d_model": model.config.hidden_size,
            "max_sequence_length": model.config.max_position_embeddings,
            "embedding_size": model.config.vocab_size,
            "pad_token_id": model.config.pad_token_id,
            "eos_token_id": model.config.eos_token_id,
            "weight_tying": model.config.tie_word_embeddings
        }
    }
    with open(os.path.join(output_path, "config.yaml"), "w") as f:
        yaml.dump(config, f)
    print(f"Config file saved to {output_path}/config.yaml")
    
    if tokenizer_path:
        tokenizer = AutoTokenizer.from_pretrained(hf_model_path)
        tokenizer.save_pretrained(output_path)
        print(f"Tokenizer saved to {output_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--hf_model_path", required=True, help="Path to the Hugging Face model.")
    parser.add_argument("--output_path", required=True, help="Directory to save the converted OLMo model.")
    parser.add_argument("--tokenizer_path", default=None, help="Optional: Path to Hugging Face tokenizer.")
    args = parser.parse_args()
    
    convert_hf_to_olmo(args.hf_model_path, args.output_path, args.tokenizer_path)

I wanted to know if the above code I wrote is safe for this conversion as I want to pretrain the allenai/OLMoE-1B-7B-0924-Instruct checkpoint from HF.

After converting this way, and checking in the model path in the pretraining code, it expects a file metadata.json which I am attaching:

metadata.json

I am not sure how to create this file.

[aditi184]

Basically I want to know how can I continually pretrain allenai/OLMoE-1B-7B-0924-Instruct checkpoint from HF using the available OLMoE pretraining code.

[Muennighoff]

Looks pretty good! I think the metadata file specifies some shapes for FSDP - to recreate it I think you can simply train a dummy model from scratch for 1 step with the FSDP distribution setup you want, save its checkpoint and take its metadata.json

[aditi184]

I actually realized that the script I pasted above is imperfect and leads to errors. I tried improving but couldn't suceed in loading the converted model using:

from olmo.model import OLMo
import torch
model = OLMo.from_checkpoint("/home/mila/k/khandela/scratch/ai2-llm/checkpoints/OLMoE/test")

Now I got more confused with a few things, for pretraining OLMoE, the model class still remains OLMo? Are there any new keys introduced for OLMoE?

Or am I loading it incorrectly?

Updated script I tried:

import argparse
import os
import torch
import yaml
import json
import gc
from pathlib import Path
from transformers import OlmoeForCausalLM, OlmoeConfig

def write_config_yaml(output_dir, config_dict):
    os.makedirs(output_dir, exist_ok=True)
    config_path = Path(output_dir) / "config.yaml"
    # Save the config under the "model" key as expected by OLMoE.
    with open(config_path, "w") as f:
        yaml.dump({"model": config_dict}, f, default_flow_style=False)

def convert_hf_to_olmoe(hf_dir, output_dir):
    print(f"Loading HF checkpoint from {hf_dir} …")
    model = OlmoeForCausalLM.from_pretrained(hf_dir, torch_dtype=torch.bfloat16)
    hf_state = model.state_dict()
    config: OlmoeConfig = model.config

    n_layers = config.num_hidden_layers
    n_heads = config.num_attention_heads
    hidden_size = config.hidden_size
    dims_per_head = hidden_size // n_heads
    base = config.rope_theta if getattr(config, "rope_theta", None) is not None else 10000.0
    max_position_embeddings = config.max_position_embeddings
    vocab_size = config.vocab_size
    pad_token_id = config.pad_token_id
    eos_token_id = config.eos_token_id

    # Determine key-value heads used (for multi-query or GQA)
    num_key_value_heads = config.num_key_value_heads if getattr(config, "num_key_value_heads", None) is not None else n_heads

    olmoe_state = {}

    # Process each transformer block/layer.
    for layer_i in range(n_layers):
        # --- Attention projection ---
        # HF model has separate q_proj, k_proj, v_proj weights.
        q_proj = hf_state[f"model.layers.{layer_i}.self_attn.q_proj.weight"]
        k_proj = hf_state[f"model.layers.{layer_i}.self_attn.k_proj.weight"]
        v_proj = hf_state[f"model.layers.{layer_i}.self_attn.v_proj.weight"]
        # Concatenate to form the original combined weight.
        att_proj_weight = torch.cat([q_proj, k_proj, v_proj], dim=0)
        key_att_proj = f"transformer.blocks.{layer_i}.att_proj.weight"
        olmoe_state[key_att_proj] = att_proj_weight
        # Create a bias for att_proj (if original expects it) as zeros.
        key_att_proj_bias = f"transformer.blocks.{layer_i}.att_proj.bias"
        olmoe_state[key_att_proj_bias] = torch.zeros(att_proj_weight.size(0), dtype=att_proj_weight.dtype)

        # --- Attention output projection ---
        o_proj_weight = hf_state[f"model.layers.{layer_i}.self_attn.o_proj.weight"]
        key_attn_out = f"transformer.blocks.{layer_i}.attn_out.weight"
        olmoe_state[key_attn_out] = o_proj_weight
        key_attn_out_bias = f"transformer.blocks.{layer_i}.attn_out.bias"
        olmoe_state[key_attn_out_bias] = torch.zeros(o_proj_weight.size(0), dtype=o_proj_weight.dtype)

        # --- Layer normalization ---
        # In the forward conversion, "input_layernorm.weight" was saved as attn_norm.weight.
        attn_norm_weight = hf_state[f"model.layers.{layer_i}.input_layernorm.weight"]
        key_attn_norm = f"transformer.blocks.{layer_i}.attn_norm.weight"
        olmoe_state[key_attn_norm] = attn_norm_weight
        key_attn_norm_bias = f"transformer.blocks.{layer_i}.attn_norm.bias"
        olmoe_state[key_attn_norm_bias] = torch.zeros(attn_norm_weight.size(0), dtype=attn_norm_weight.dtype)

        # Similarly, "post_attention_layernorm.weight" becomes ff_norm.weight.
        ff_norm_weight = hf_state[f"model.layers.{layer_i}.post_attention_layernorm.weight"]
        key_ff_norm = f"transformer.blocks.{layer_i}.ff_norm.weight"
        olmoe_state[key_ff_norm] = ff_norm_weight
        key_ff_norm_bias = f"transformer.blocks.{layer_i}.ff_norm.bias"
        olmoe_state[key_ff_norm_bias] = torch.zeros(ff_norm_weight.size(0), dtype=ff_norm_weight.dtype)

        # --- Feed-forward (FFN) router ---
        router_weight = hf_state[f"model.layers.{layer_i}.mlp.gate.weight"]
        key_ff_router = f"transformer.blocks.{layer_i}.ffn.router.layer.weight"
        olmoe_state[key_ff_router] = router_weight
        key_ff_router_bias = f"transformer.blocks.{layer_i}.ffn.router.layer.bias"
        # If the original did not have a bias, fill with zeros.
        olmoe_state[key_ff_router_bias] = torch.zeros(router_weight.size(0), dtype=router_weight.dtype)

        # --- Experts ---
        # Determine number of experts by checking available expert keys.
        num_experts = 0
        while f"model.layers.{layer_i}.mlp.experts.{num_experts}.gate_proj.weight" in hf_state:
            num_experts += 1

        expert_w1_list = []
        expert_v1_list = []
        expert_w2_list = []
        for expert_i in range(num_experts):
            gate_proj = hf_state[f"model.layers.{layer_i}.mlp.experts.{expert_i}.gate_proj.weight"]
            expert_w1_list.append(gate_proj)
            up_proj = hf_state[f"model.layers.{layer_i}.mlp.experts.{expert_i}.up_proj.weight"]
            expert_v1_list.append(up_proj)
            # down_proj was transposed during the forward conversion; reverse that here.
            down_proj = hf_state[f"model.layers.{layer_i}.mlp.experts.{expert_i}.down_proj.weight"]
            expert_w2_list.append(down_proj.t().contiguous())
        # Concatenate experts’ weights along the first dimension.
        key_experts_w1 = f"transformer.blocks.{layer_i}.ffn.experts.mlp.w1"
        olmoe_state[key_experts_w1] = torch.cat(expert_w1_list, dim=0)
        key_experts_v1 = f"transformer.blocks.{layer_i}.ffn.experts.mlp.v1"
        olmoe_state[key_experts_v1] = torch.cat(expert_v1_list, dim=0)
        key_experts_w2 = f"transformer.blocks.{layer_i}.ffn.experts.mlp.w2"
        olmoe_state[key_experts_w2] = torch.cat(expert_w2_list, dim=0)

        # --- Rotary embeddings ---
        # Recompute inv_freq using dims_per_head and base.
        inv_freq = 1.0 / (base ** (torch.arange(0, dims_per_head, 2, dtype=torch.float32) / dims_per_head))
        key_inv_freq = f"transformer.blocks.{layer_i}.self_attn.rotary_emb.inv_freq"
        olmoe_state[key_inv_freq] = inv_freq

        # --- (Optional) Feed-forward output projections ---
        # Some OLMoE implementations expect an additional FFN output projection.
        # If missing, we fill with zeros. Here we assume the expected shape is (hidden_size, hidden_size).
        key_ff_out_weight = f"transformer.blocks.{layer_i}.ff_out.weight"
        key_ff_out_bias = f"transformer.blocks.{layer_i}.ff_out.bias"
        olmoe_state[key_ff_out_weight] = torch.zeros(hidden_size, hidden_size, dtype=hf_state[f"model.layers.{layer_i}.mlp.gate.weight"].dtype)
        olmoe_state[key_ff_out_bias] = torch.zeros(hidden_size, dtype=hf_state[f"model.layers.{layer_i}.mlp.gate.weight"].dtype)

        # Similarly, if a feed-forward projection is expected.
        key_ff_proj_weight = f"transformer.blocks.{layer_i}.ff_proj.weight"
        key_ff_proj_bias = f"transformer.blocks.{layer_i}.ff_proj.bias"
        olmoe_state[key_ff_proj_weight] = torch.zeros(hidden_size, hidden_size, dtype=hf_state[f"model.layers.{layer_i}.mlp.gate.weight"].dtype)
        olmoe_state[key_ff_proj_bias] = torch.zeros(hidden_size, dtype=hf_state[f"model.layers.{layer_i}.mlp.gate.weight"].dtype)

    # --- Final (root) weights ---
    # Embedding tokens.
    olmoe_state["transformer.wte.weight"] = hf_state["model.embed_tokens.weight"]
    # Final LM head.
    olmoe_state["transformer.ff_out.weight"] = hf_state["lm_head.weight"]
    # Final layer norm.
    olmoe_state["transformer.ln_f.weight"] = hf_state["model.norm.weight"]
    olmoe_state["transformer.ln_f.bias"] = torch.zeros(hf_state["model.norm.weight"].size(0), dtype=hf_state["model.norm.weight"].dtype)

    # Save the OLMoE checkpoint as a single file "model.pt"
    os.makedirs(output_dir, exist_ok=True)
    output_model_path = os.path.join(output_dir, "model.pt")
    torch.save(olmoe_state, output_model_path)
    print(f"Saved OLMoE model weights to {output_model_path}")

    # --- Recreate the OLMoE config.yaml ---
    olmoe_config = {
        "n_layers": n_layers,
        "n_heads": n_heads,
        "d_model": hidden_size,
        "max_sequence_length": max_position_embeddings,
        "pad_token_id": pad_token_id,
        "eos_token_id": eos_token_id,
        "vocab_size": vocab_size,
        "weight_tying": config.tie_word_embeddings if hasattr(config, "tie_word_embeddings") else True,
        "multi_query_attention": True if num_key_value_heads == 1 else False,
        "clip_qkv": getattr(config, "clip_qkv", None),
    }
    write_config_yaml(output_dir, olmoe_config)
    print(f"Saved OLMoE config to {os.path.join(output_dir, 'config.yaml')}")

    # Cleanup.
    del olmoe_state, hf_state
    gc.collect()

def main():
    parser = argparse.ArgumentParser(
        description="Convert an HF (Transformers) OLMoE checkpoint to the original OLMoE format."
    )
    parser.add_argument("--input_dir", required=True, help="Path to the HF checkpoint directory (contains config.json, pytorch_model.bin, etc.).")
    parser.add_argument("--output_dir", required=True, help="Path where the OLMoE checkpoint (model.pt and config.yaml) will be saved.")
    args = parser.parse_args()
    convert_hf_to_olmoe(args.input_dir, args.output_dir)

if __name__ == "__main__":
    main()

[Muennighoff]

the model class still remains OLMo

yes! There are just some new keys related to the experts and gating that you should have from the ckpt

[aditi184]

Hey Niklas

Upon multiple trials to convert the HF checkpoint to OLMo class, I still am encountering issues like:

Missing key(s) in state_dict: "transformer.wpe.weight".
Unexpected key(s) in state_dict: "transformer.blocks.0.q_norm.weight", "transformer.blocks.0.k_norm.weight", "transformer.blocks.1.q_norm.weight", "transformer.blocks.1.k_norm.weight", "transformer.blocks.2.q_norm.weight", "transformer.blocks.2.k_norm.weight", "transformer.blocks.3.q_norm.weight", "transformer.blocks.3.k_norm.weight"

and shape mismatch errors.

Although I printed all shapes and they look fine, still the script isn't working.

Can you provide the checkpoint in the OLMo format so that I can load it directly? like there are checkpoints available on HuggingFace which are loaded without transformers library. Or if you can provide the script.

# Run script using python convert_hf_to_olmo2.py --hf_model_path allenai/OLMoE-1B-7B-0924-Instruct --output_dir /home/mila/k/khandela/scratch/ai2-llm/checkpoints/OLMoE/test

import argparse
import json
import os
import torch
import yaml
from pathlib import Path
from transformers import OlmoeForCausalLM, OlmoeConfig
import ipdb
def write_json(text, path):
    with open(path, "w") as f:
        json.dump(text, f, indent=4)

def save_olmo_checkpoint(model_path, hf_model_path):
    os.makedirs(model_path, exist_ok=True)
    tmp_model_path = os.path.join(model_path, "tmp")
    os.makedirs(tmp_model_path, exist_ok=True)
    
    # Load HF model
    hf_checkpoint = OlmoeForCausalLM.from_pretrained(hf_model_path, torch_dtype=torch.float32, low_cpu_mem_usage=True)
    hf_config = hf_checkpoint.config
    print(hf_config)
    n_layers = hf_config.num_hidden_layers
    n_heads = hf_config.num_attention_heads
    dim = hf_config.hidden_size
    dims_per_head = dim // n_heads
    num_key_value_heads = hf_config.num_key_value_heads
    num_experts = hf_config.intermediate_size // (dim // num_key_value_heads)
    print(num_experts)  # Estimate number of experts
    base = hf_config.rope_theta

    state_dict = hf_checkpoint.state_dict()
    
    olmoe_state_dict = {}
    
    # Convert transformer layers
    for layer_i in range(hf_config.num_hidden_layers):
        q_proj_weight = state_dict[f"model.layers.{layer_i}.self_attn.q_proj.weight"]
        k_proj_weight = state_dict[f"model.layers.{layer_i}.self_attn.k_proj.weight"]
        v_proj_weight = state_dict[f"model.layers.{layer_i}.self_attn.v_proj.weight"]
        
        # Fuse Q, K, V into a single tensor (OLMoE format)
        att_proj_weight = torch.cat([q_proj_weight, k_proj_weight, v_proj_weight], dim=0)

        # Map back to OLMoE format
        olmoe_state_dict[f"transformer.blocks.{layer_i}.att_proj.weight"] = att_proj_weight
        olmoe_state_dict[f"transformer.blocks.{layer_i}.attn_out.weight"] = state_dict[f"model.layers.{layer_i}.self_attn.o_proj.weight"]
        olmoe_state_dict[f"transformer.blocks.{layer_i}.q_norm.weight"] = state_dict[f"model.layers.{layer_i}.self_attn.q_norm.weight"]
        olmoe_state_dict[f"transformer.blocks.{layer_i}.k_norm.weight"] = state_dict[f"model.layers.{layer_i}.self_attn.k_norm.weight"]
        olmoe_state_dict[f"transformer.blocks.{layer_i}.ffn.router.layer.weight"] = state_dict[f"model.layers.{layer_i}.mlp.gate.weight"]
        olmoe_state_dict[f"transformer.blocks.{layer_i}.attn_norm.weight"] = state_dict[f"model.layers.{layer_i}.input_layernorm.weight"]
        olmoe_state_dict[f"transformer.blocks.{layer_i}.ff_norm.weight"] = state_dict[f"model.layers.{layer_i}.post_attention_layernorm.weight"]

        
        num_experts = state_dict[f"model.layers.{layer_i}.mlp.gate.weight"].shape[0]
        # dim_per_expert = state_dict[f"model.layers.{layer_i}.mlp.experts.0.gate_proj.weight"].shape[1]
        print(num_experts)
        # Reconstruct MoE expert weights
        w1_list = [state_dict[f"model.layers.{layer_i}.mlp.experts.{expert_i}.gate_proj.weight"] for expert_i in range(num_experts)]
        v1_list = [state_dict[f"model.layers.{layer_i}.mlp.experts.{expert_i}.up_proj.weight"] for expert_i in range(num_experts)]
        w2_list = [state_dict[f"model.layers.{layer_i}.mlp.experts.{expert_i}.down_proj.weight"].T for expert_i in range(num_experts)]

        olmoe_state_dict[f"transformer.blocks.{layer_i}.ffn.experts.mlp.w1"] = torch.cat(w1_list, dim=0)
        olmoe_state_dict[f"transformer.blocks.{layer_i}.ffn.experts.mlp.v1"] = torch.cat(v1_list, dim=0)
        olmoe_state_dict[f"transformer.blocks.{layer_i}.ffn.experts.mlp.w2"] = torch.cat(w2_list, dim=0)

        


    # ipdb.set_trace()
    dims_per_head = dim // n_heads
    print("here")
    print(n_heads)
    print(dim)
    print(dims_per_head)
    base = 10000.0
    # olmoe_state_dict["transformer.inv_freq"] = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head))
    olmoe_state_dict["transformer.wte.weight"] = state_dict["model.embed_tokens.weight"]
    olmoe_state_dict["transformer.ff_out.weight"] = state_dict["lm_head.weight"]
    olmoe_state_dict["transformer.ln_f.weight"] = state_dict["model.norm.weight"]
    
    # Save model.pt
    torch.save(olmoe_state_dict, os.path.join(model_path, "model.pt"))
    
    # Save config.yaml
    olmoe_config = {
        "model": {
            "n_layers": hf_config.num_hidden_layers,
            "n_heads": hf_config.num_attention_heads,
            "d_model": hf_config.hidden_size,
            "embedding_size": hf_config.vocab_size,
            "max_sequence_length": hf_config.max_position_embeddings,
            "pad_token_id": hf_config.pad_token_id,
            "eos_token_id": hf_config.eos_token_id,
            "weight_tying": hf_config.tie_word_embeddings,
            "block_type": "moe",
            "n_kv_heads" : None,
            "clip_qkv" : None,
            "include_bias" : False,
            "bias_for_layer_norm" : False       
             }
    }
    print(olmoe_config)
    with open(os.path.join(model_path, "config.yaml"), "w") as f:
        yaml.dump(olmoe_config, f)
    
    print(f"Converted model saved to {model_path}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--hf_model_path", required=True, help="Path to the HF model directory.")
    parser.add_argument("--output_dir", required=True, help="Directory to save OLMo checkpoint.")
    args = parser.parse_args()
    save_olmo_checkpoint(args.output_dir, args.hf_model_path)

if __name__ == "__main__":
    main()

[Muennighoff]

Tagging @aman-17 here 🙌

[aman-17]

@aditi184, did you solve this? The errors like missing keys such as transformer.wpe.weight and unexpected ones like q_norm.weight are expected if there’s a structural mismatch between the HF checkpoint and what the OLMo class expects. It’s likely due to architectural differences.

[aditi184]

Yes I have written the conversion script for it

I can create a PR if you would like to include it in the code-base.

[DominicSlw]

Hi, @aditi184 I am also working on the continual training process and want to convert weights from huggingface formats to olmoe checkpoint format. Could you share the scripts? Huge thanks!!

[aditi184]

Hey @DominicSlw

Here you go:

import argparse
import gc
import json
import os
import shutil
from pathlib import Path

import torch
import yaml
from tokenizers import Tokenizer

from transformers import OlmoeConfig, OlmoeForCausalLM
from transformers.models.gpt_neox.tokenization_gpt_neox_fast import GPTNeoXTokenizerFast


def compute_intermediate_size(n, ffn_dim_multiplier=1, multiple_of=256):
    return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3)) + multiple_of - 1) // multiple_of)


def read_json(path):
    with open(path, "r") as f:
        return json.load(f)


def write_json(text, path):
    with open(path, "w") as f:
        json.dump(text, f)


def write_model(model_path, input_base_path, tokenizer_path=None, safe_serialization=True, fix_eos_token_id=True):
    os.makedirs(model_path, exist_ok=True)
    tmp_model_path = os.path.join(model_path, "tmp")
    os.makedirs(tmp_model_path, exist_ok=True)

    config_path = Path(input_base_path) / "config.yaml"
    olmoe_config = yaml.safe_load(config_path.read_text())["model"]

    if fix_eos_token_id:
        olmoe_config["eos_token_id"] = 50279

    n_layers = olmoe_config["n_layers"]
    n_heads = olmoe_config["n_heads"]
    dim = olmoe_config["d_model"]
    dims_per_head = dim // n_heads
    base = 10000.0
    inv_freq = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head))
    max_position_embeddings = olmoe_config["max_sequence_length"]

    vocab_size = olmoe_config.get("embedding_size", olmoe_config["vocab_size"])

    if olmoe_config.get("n_kv_heads", None) is not None:
        num_key_value_heads = olmoe_config["n_kv_heads"]  # for GQA / MQA
    elif olmoe_config["multi_query_attention"]:  # compatibility with other checkpoints
        num_key_value_heads = 1
    else:
        num_key_value_heads = n_heads

    print(f"Fetching all parameters from the checkpoint at {input_base_path}.")

    # Not sharded
    loaded = torch.load(os.path.join(input_base_path, "model.pt"), map_location="cpu")

    param_count = 0
    index_dict = {"weight_map": {}}
    for layer_i in range(n_layers):
        filename = f"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin"
        fused_dims = [dim, dims_per_head * num_key_value_heads, dims_per_head * num_key_value_heads]
        q_proj_weight, k_proj_weight, v_proj_weight = torch.split(
            loaded[f"transformer.blocks.{layer_i}.att_proj.weight"], fused_dims, dim=0
        )
        state_dict = {
            f"model.layers.{layer_i}.self_attn.q_proj.weight": q_proj_weight,
            f"model.layers.{layer_i}.self_attn.k_proj.weight": k_proj_weight,
            f"model.layers.{layer_i}.self_attn.v_proj.weight": v_proj_weight,
            f"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[f"transformer.blocks.{layer_i}.attn_out.weight"],
            # f"model.layers.{layer_i}.self_attn.q_norm.weight": loaded[f"transformer.blocks.{layer_i}.q_norm.weight"],
            # f"model.layers.{layer_i}.self_attn.k_norm.weight": loaded[f"transformer.blocks.{layer_i}.k_norm.weight"],
            f"model.layers.{layer_i}.mlp.gate.weight": loaded[f"transformer.blocks.{layer_i}.ffn.router.layer.weight"],
            f"model.layers.{layer_i}.input_layernorm.weight": loaded[f"transformer.blocks.{layer_i}.attn_norm.weight"],
            f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[
                f"transformer.blocks.{layer_i}.ff_norm.weight"
            ],
        }

        num_experts = loaded[f"transformer.blocks.{layer_i}.ffn.router.layer.weight"].shape[0]
        dim_per_expert = loaded[f"transformer.blocks.{layer_i}.ffn.experts.mlp.w1"].shape[0] // num_experts
        for expert_i in range(num_experts):
            state_dict[f"model.layers.{layer_i}.mlp.experts.{expert_i}.gate_proj.weight"] = loaded[
                f"transformer.blocks.{layer_i}.ffn.experts.mlp.w1"
            ][dim_per_expert * expert_i : dim_per_expert * (expert_i + 1), :]
            state_dict[f"model.layers.{layer_i}.mlp.experts.{expert_i}.up_proj.weight"] = loaded[
                f"transformer.blocks.{layer_i}.ffn.experts.mlp.v1"
            ][dim_per_expert * expert_i : dim_per_expert * (expert_i + 1), :]
            state_dict[f"model.layers.{layer_i}.mlp.experts.{expert_i}.down_proj.weight"] = loaded[
                f"transformer.blocks.{layer_i}.ffn.experts.mlp.w2"
            ][dim_per_expert * expert_i : dim_per_expert * (expert_i + 1), :].T.contiguous()

        state_dict[f"model.layers.{layer_i}.self_attn.rotary_emb.inv_freq"] = inv_freq

        for k, v in state_dict.items():
            index_dict["weight_map"][k] = filename
            param_count += v.numel()
        torch.save(state_dict, os.path.join(tmp_model_path, filename))

    filename = f"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin"

    # Unsharded
    state_dict = {
        "model.embed_tokens.weight": loaded["transformer.wte.weight"],
        "lm_head.weight": loaded["transformer.ff_out.weight"],
        "model.norm.weight": loaded["transformer.ln_f.weight"],
    }

    for k, v in state_dict.items():
        index_dict["weight_map"][k] = filename
        param_count += v.numel()
    torch.save(state_dict, os.path.join(tmp_model_path, filename))

    # Write configs
    index_dict["metadata"] = {"total_size": param_count * 2}
    write_json(index_dict, os.path.join(tmp_model_path, "pytorch_model.bin.index.json"))

    config = OlmoeConfig(
        vocab_size=vocab_size,
        hidden_size=dim,
        intermediate_size=dim_per_expert,
        num_hidden_layers=n_layers,
        num_attention_heads=n_heads,
        num_key_value_heads=num_key_value_heads,
        max_position_embeddings=max_position_embeddings,
        pad_token_id=olmoe_config["pad_token_id"],
        bos_token_id=None,
        eos_token_id=olmoe_config["eos_token_id"],
        tie_word_embeddings=olmoe_config["weight_tying"],
        rope_theta=base,
        clip_qkv=olmoe_config.get("clip_qkv"),
    )
    config.save_pretrained(tmp_model_path)

    # Make space so we can load the model properly now.
    del state_dict
    del loaded
    gc.collect()

    if tokenizer_path is not None:
        _write_tokenizer(model_path, config, tokenizer_path, fix_eos_token_id)

    print("Loading the checkpoint in a OLMoE model.")
    model = OlmoeForCausalLM.from_pretrained(tmp_model_path, torch_dtype=torch.bfloat16)
    # Avoid saving this as part of the config.
    del model.config._name_or_path
    print("Saving in the Transformers format.")
    model.save_pretrained(model_path, safe_serialization=safe_serialization)
    shutil.rmtree(tmp_model_path)


def _write_tokenizer(
    output_path: Path, config: OlmoeConfig, input_tokenizer_path: Path, fix_eos_token_id: bool = True
) -> None:
    print(f"Saving a {GPTNeoXTokenizerFast.__name__} to {output_path}.")

    base_tokenizer = Tokenizer.from_file(str(input_tokenizer_path))

    eos_token_id = config.eos_token_id if config.eos_token_id is not None else base_tokenizer.get_vocab_size() - 1
    pad_token_id = config.pad_token_id if config.pad_token_id is not None else eos_token_id

    if fix_eos_token_id and eos_token_id == 0:
        # Fixing a bug in OLMo where eos token id was incorrectly set
        print("Changing eos_token_id from 0 to 50279.")
        eos_token_id = 50279

    tokenizer = GPTNeoXTokenizerFast(
        tokenizer_object=base_tokenizer,
        eos_token=base_tokenizer.decode([eos_token_id], skip_special_tokens=False),
        pad_token=base_tokenizer.decode([pad_token_id], skip_special_tokens=False),
        unk_token=None,
        bos_token=None,
    )

    tokenizer.save_pretrained(output_path)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_dir",
        required=True,
        help="Location of OLMoE weights, which contains config.yaml and model.pt.",
    )
    parser.add_argument(
        "--tokenizer_json_path",
        default=None,
        help="Location of OLMoE tokenizer json file.",
    )
    parser.add_argument(
        "--output_dir",
        required=True,
        help="Location to write HF model and tokenizer",
    )
    parser.add_argument(
        "--no_fix_eos_token_id",
        action="store_false",
        dest="fix_eos_token_id",
        help="If set, does not change eos token id from 0 to 50279 if it is 0. Changing 0 to 50279 is a bug fix, so use this option with care.",
    )
    parser.add_argument(
        "--safe_serialization", type=bool, default=True, help="Whether or not to save using `safetensors`."
    )
    args = parser.parse_args()
    write_model(
        model_path=args.output_dir,
        input_base_path=args.input_dir,
        safe_serialization=args.safe_serialization,
        tokenizer_path=args.tokenizer_json_path,
        fix_eos_token_id=args.fix_eos_token_id,
    )


if __name__ == "__main__":
    main()

[DominicSlw]

Hi, @aditi184. Thanks for your quick reply! This script seems to convert from olmo fomat to huggingface format instead of hf format to olmo format. Could you share the script on hf--> olmo conversion?

[aditi184]

Oh sorry, here it is:

# Run script using python convert_hf_to_olmo2.py --hf_model_path allenai/OLMoE-1B-7B-0924-Instruct --output_dir /home/mila/k/khandela/scratch/ai2-llm/checkpoints/OLMoE/base-0924

import argparse
import json
import os
import torch
import yaml
from pathlib import Path
from transformers import OlmoeForCausalLM, OlmoeConfig
import ipdb
def write_json(text, path):
    with open(path, "w") as f:
        json.dump(text, f, indent=4)

def save_olmo_checkpoint(model_path, hf_model_path):
    os.makedirs(model_path, exist_ok=True)
    tmp_model_path = os.path.join(model_path, "tmp")
    os.makedirs(tmp_model_path, exist_ok=True)
    
    # Load HF model
    hf_checkpoint = OlmoeForCausalLM.from_pretrained(hf_model_path, torch_dtype=torch.float32, low_cpu_mem_usage=True)
    hf_config = hf_checkpoint.config
    print(hf_config)
    n_layers = hf_config.num_hidden_layers
    n_heads = hf_config.num_attention_heads
    dim = hf_config.hidden_size
    dims_per_head = dim // n_heads
    num_key_value_heads = hf_config.num_key_value_heads
    num_experts = hf_config.intermediate_size // (dim // num_key_value_heads)
    print(num_experts)  # Estimate number of experts
    base = hf_config.rope_theta

    state_dict = hf_checkpoint.state_dict()
    
    olmoe_state_dict = {}
    
    # Convert transformer layers
    for layer_i in range(hf_config.num_hidden_layers):
        q_proj_weight = state_dict[f"model.layers.{layer_i}.self_attn.q_proj.weight"]
        k_proj_weight = state_dict[f"model.layers.{layer_i}.self_attn.k_proj.weight"]
        v_proj_weight = state_dict[f"model.layers.{layer_i}.self_attn.v_proj.weight"]
        
        # Fuse Q, K, V into a single tensor (OLMoE format)
        att_proj_weight = torch.cat([q_proj_weight, k_proj_weight, v_proj_weight], dim=0)

        # Map back to OLMoE format
        olmoe_state_dict[f"transformer.blocks.{layer_i}.att_proj.weight"] = att_proj_weight
        olmoe_state_dict[f"transformer.blocks.{layer_i}.attn_out.weight"] = state_dict[f"model.layers.{layer_i}.self_attn.o_proj.weight"]
        olmoe_state_dict[f"transformer.blocks.{layer_i}.q_norm.weight"] = state_dict[f"model.layers.{layer_i}.self_attn.q_norm.weight"]
        olmoe_state_dict[f"transformer.blocks.{layer_i}.k_norm.weight"] = state_dict[f"model.layers.{layer_i}.self_attn.k_norm.weight"]
        olmoe_state_dict[f"transformer.blocks.{layer_i}.ffn.router.layer.weight"] = state_dict[f"model.layers.{layer_i}.mlp.gate.weight"]
        olmoe_state_dict[f"transformer.blocks.{layer_i}.attn_norm.weight"] = state_dict[f"model.layers.{layer_i}.input_layernorm.weight"]
        olmoe_state_dict[f"transformer.blocks.{layer_i}.ff_norm.weight"] = state_dict[f"model.layers.{layer_i}.post_attention_layernorm.weight"]

        
        num_experts = state_dict[f"model.layers.{layer_i}.mlp.gate.weight"].shape[0]
        # dim_per_expert = state_dict[f"model.layers.{layer_i}.mlp.experts.0.gate_proj.weight"].shape[1]
        print(num_experts)
        # Reconstruct MoE expert weights
        w1_list = [state_dict[f"model.layers.{layer_i}.mlp.experts.{expert_i}.gate_proj.weight"] for expert_i in range(num_experts)]
        v1_list = [state_dict[f"model.layers.{layer_i}.mlp.experts.{expert_i}.up_proj.weight"] for expert_i in range(num_experts)]
        w2_list = [state_dict[f"model.layers.{layer_i}.mlp.experts.{expert_i}.down_proj.weight"].T for expert_i in range(num_experts)]

        olmoe_state_dict[f"transformer.blocks.{layer_i}.ffn.experts.mlp.w1"] = torch.cat(w1_list, dim=0)
        olmoe_state_dict[f"transformer.blocks.{layer_i}.ffn.experts.mlp.v1"] = torch.cat(v1_list, dim=0)
        olmoe_state_dict[f"transformer.blocks.{layer_i}.ffn.experts.mlp.w2"] = torch.cat(w2_list, dim=0)

        


    # ipdb.set_trace()
    dims_per_head = dim // n_heads
    print("here")
    print(n_heads)
    print(dim)
    print(dims_per_head)
    base = 10000.0
    # olmoe_state_dict["transformer.inv_freq"] = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head))
    olmoe_state_dict["transformer.wte.weight"] = state_dict["model.embed_tokens.weight"]
    olmoe_state_dict["transformer.ff_out.weight"] = state_dict["lm_head.weight"]
    olmoe_state_dict["transformer.ln_f.weight"] = state_dict["model.norm.weight"]
    
    # Save model.pt
    torch.save(olmoe_state_dict, os.path.join(model_path, "model.pt"))
    
    # Save config.yaml
    olmoe_config = {
        "model": {
            "n_layers": hf_config.num_hidden_layers,
            "n_heads": hf_config.num_attention_heads,
            "d_model": hf_config.hidden_size,
            "embedding_size": hf_config.vocab_size,
            "max_sequence_length": hf_config.max_position_embeddings,
            "pad_token_id": hf_config.pad_token_id,
            "eos_token_id": hf_config.eos_token_id,
            "weight_tying": hf_config.tie_word_embeddings,
            "block_type": "moe",
            "n_kv_heads" : None,
            "clip_qkv" : None,
            "include_bias" : False,
            "bias_for_layer_norm" : False       
             }
    }
    print(olmoe_config)
    with open(os.path.join(model_path, "config.yaml"), "w") as f:
        yaml.dump(olmoe_config, f)
    
    print(f"Converted model saved to {model_path}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--hf_model_path", required=True, help="Path to the HF model directory.")
    parser.add_argument("--output_dir", required=True, help="Directory to save OLMo checkpoint.")
    args = parser.parse_args()
    save_olmo_checkpoint(args.output_dir, args.hf_model_path)

if __name__ == "__main__":
    main()