Checkpoint modification

import transformers
import torch
import tqdm
import json

####### INPUT SECTION ##################################################################

original_ckpt_location = "Qwen2.5-Coder-0.5B-Instruct/"             # Location of the original checkpoint
updated_ckpt_location = "Qwen2.5-Coder-0.5B-Instruct_surgery/"      # Location of the updated checkpoint

original_config_location = "0.5B_config.json"                       # Location of the original config.json
updated_config_location = "0.5b_config_surgery.json"                # Location of the updated config.json

####### END OF INPUT SECTION ###########################################################


# Load the original checkpoints
model = transformers.AutoModelForCausalLM.from_pretrained(original_ckpt_location)
tokenizer = transformers.AutoTokenizer.from_pretrained(original_ckpt_location)

# Load the original config
with open(original_config_location, "r") as f:
    original_config = json.load(f)

# Load the updated config
with open(updated_config_location, "r") as f:
    updated_config = json.load(f)

# Asserts to make sure that only num_attention_heads, num_key_value_heads, intermediate_size, tie_word_embeddings have changed
for param in original_config.keys():
    if param not in ["num_attention_heads", "intermediate_size", "tie_word_embeddings", "num_key_value_heads", "transformers_version", "eos_token_id"]:
        assert original_config[param] == updated_config[param], f"Updated config must have the same {param} as original config"

# Update possible changed parameters
if "head_dim" in original_config.keys():
    HEAD_SIZE = int(original_config["head_dim"])
else:
    HEAD_SIZE = int(original_config["hidden_size"] / original_config["num_attention_heads"])

OLD_NUM_KV_HEADS = int(original_config["num_key_value_heads"])
NEW_NUM_KV_HEADS = int(updated_config["num_key_value_heads"])

OLD_NUM_HEADS = int(original_config["num_attention_heads"])
NEW_NUM_HEADS = int(updated_config["num_attention_heads"])

# To change the intermediate size, update these parameters (or comment out to keep original intermediate size)
OLD_INTERMEDIATE_SIZE = int(original_config["intermediate_size"])
NEW_INTERMEDIATE_SIZE = int(updated_config["intermediate_size"])

for layer in tqdm.tqdm(model.model.layers):

    # Update Q/O_Proj sizes if num_attention_heads changes
    if OLD_NUM_HEADS != NEW_NUM_HEADS:

        old_group_size = int(OLD_NUM_HEADS / OLD_NUM_KV_HEADS)
        new_group_size = int(NEW_NUM_HEADS / OLD_NUM_KV_HEADS)

        # Update Q Weight Matrix
        new_q_weight_matrix = torch.zeros((NEW_NUM_HEADS * HEAD_SIZE, OLD_NUM_HEADS * HEAD_SIZE))
        for i in range(OLD_NUM_KV_HEADS):
            new_q_weight_matrix[i * new_group_size * HEAD_SIZE: ((i * new_group_size) + old_group_size) * HEAD_SIZE, :] = layer.self_attn.q_proj.weight[i * old_group_size * HEAD_SIZE: (i+1) * old_group_size * HEAD_SIZE,:]
        layer.self_attn.q_proj.weight.data = new_q_weight_matrix

        # Update Q Bias (in case of Qwen2)
        if model.config.model_type == "qwen2":
            new_q_bias = torch.zeros((NEW_NUM_HEADS * HEAD_SIZE))
            for i in range(OLD_NUM_KV_HEADS):
                new_q_bias[i * new_group_size * HEAD_SIZE:((i * new_group_size) + old_group_size) * HEAD_SIZE] = layer.self_attn.q_proj.bias[i * old_group_size * HEAD_SIZE: (i+1) * old_group_size * HEAD_SIZE]
            layer.self_attn.q_proj.bias.data = new_q_bias

        # Update O Weight Matrix
        new_o_weight_matrix = torch.zeros((OLD_NUM_HEADS * HEAD_SIZE, NEW_NUM_HEADS * HEAD_SIZE))
        for i in range(OLD_NUM_KV_HEADS):
            new_o_weight_matrix[:,i * new_group_size * HEAD_SIZE:((i * new_group_size) + old_group_size) * HEAD_SIZE] = layer.self_attn.o_proj.weight[:, i * old_group_size * HEAD_SIZE: (i+1) * old_group_size * HEAD_SIZE]
        layer.self_attn.o_proj.weight.data = new_o_weight_matrix

        # Update KV Gemm Weight Matrices if needed
        if OLD_NUM_KV_HEADS != NEW_NUM_KV_HEADS:
            assert OLD_NUM_KV_HEADS == 2, "Updates to num_key_value_heads are only allowed for Qwen2.5 0.5B"

            duplication_factor = int(NEW_NUM_KV_HEADS / OLD_NUM_KV_HEADS)

            layer.self_attn.k_proj.weight.data = torch.cat((layer.self_attn.k_proj.weight[:HEAD_SIZE, :].repeat(duplication_factor, 1),
                                                        layer.self_attn.k_proj.weight[HEAD_SIZE:, :].repeat(duplication_factor, 1)), dim=0)
            layer.self_attn.v_proj.weight.data = torch.cat((layer.self_attn.v_proj.weight[:HEAD_SIZE, :].repeat(duplication_factor, 1),
                                                            layer.self_attn.v_proj.weight[HEAD_SIZE:, :].repeat(duplication_factor, 1)), dim=0)

            # Update KV Bias (in case of Qwen2)
            if model.config.model_type == "qwen2":
                layer.self_attn.k_proj.bias.data = torch.cat((layer.self_attn.k_proj.bias[:HEAD_SIZE].repeat(duplication_factor),
                                                                layer.self_attn.k_proj.bias[HEAD_SIZE:].repeat(duplication_factor)), dim=0)
                layer.self_attn.v_proj.bias.data = torch.cat((layer.self_attn.v_proj.bias[:HEAD_SIZE].repeat(duplication_factor),
                                                            layer.self_attn.v_proj.bias[HEAD_SIZE:].repeat(duplication_factor)), dim=0)

    # Update FFN Gemm sizes if intermediate dim changes
    if OLD_INTERMEDIATE_SIZE != NEW_INTERMEDIATE_SIZE:

        padding = NEW_INTERMEDIATE_SIZE - OLD_INTERMEDIATE_SIZE

        # Update Down Gemm Weight Matrix
        layer.mlp.down_proj.weight.data = torch.cat((layer.mlp.down_proj.weight.data,
                                                torch.zeros((layer.mlp.down_proj.weight.data.shape[0], padding)).to(layer.mlp.down_proj.weight.device)),dim=1)

        # Update Gate/Up Gemm Weight Matrix
        layer.mlp.gate_proj.weight.data = torch.cat((layer.mlp.gate_proj.weight.data,
                                                torch.zeros((padding, layer.mlp.gate_proj.weight.data.shape[1],)).to(layer.mlp.gate_proj.weight.device)),dim=0)
        layer.mlp.up_proj.weight.data = torch.cat((layer.mlp.up_proj.weight.data,
                                                torch.zeros((padding, layer.mlp.up_proj.weight.data.shape[1],)).to(layer.mlp.up_proj.weight.device)),dim=0)

# Update new model config
model.config.head_dim = HEAD_SIZE
model.config.num_attention_heads = NEW_NUM_HEADS
model.config.intermediate_size = NEW_INTERMEDIATE_SIZE
model.config.num_key_value_heads = NEW_NUM_KV_HEADS

# Check if weight needs to be untied offline
if original_config["tie_word_embeddings"] != updated_config["tie_word_embeddings"]:
    new_weight = torch.nn.Parameter(model.get_input_embeddings().weight.data.clone())
    assert model.get_input_embeddings().weight.data.data_ptr() != new_weight.data.data_ptr()
    model.get_output_embeddings().weight = new_weight
    model.config.tie_word_embeddings = False

# Save updated model weights in specified location
model.save_pretrained(updated_ckpt_location)
tokenizer.save_pretrained(updated_ckpt_location)