File size: 4,427 Bytes
3df47d2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
from tqdm import tqdm
# Set environment variables and configure the number of threads
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
torch.set_num_threads(1)
# Define model names and output directory
base_model_name1 = "huihui-ai/QwQ-32B-Preview-abliterated"
base_model_name2 = "huihui-ai/Qwen2.5-Coder-32B-Instruct-abliterated"
output_model_dir = "huihui-ai/QwQ-32B-Coder-Fusion-9010" # Directory to save the merged model
# Specify the device for computation
device = torch.device("cpu")
# Load models and tokenizer onto the CPU
base_model1 = AutoModelForCausalLM.from_pretrained(base_model_name1, torch_dtype=torch.bfloat16).to(device)
base_model2 = AutoModelForCausalLM.from_pretrained(base_model_name2, torch_dtype=torch.bfloat16).to(device)
tokenizer = AutoTokenizer.from_pretrained(base_model_name1)
def merge_model_weights(base_model1, base_model2, alpha=0.9):
"""
Merge the weights of two models based on a specified ratio and return the updated model.
The parameter alpha determines the blending ratio, with a default of 0.9
(90% from base_model1 and 10% from base_model2).
"""
# Merge weights of the output layers
base_model1.lm_head.weight.data = (alpha * base_model1.lm_head.weight.data + (1 - alpha) * base_model2.lm_head.weight.data).to(device)
base_model1.model.embed_tokens.weight.data = (alpha * base_model1.model.embed_tokens.weight.data + (1 - alpha) * base_model2.model.embed_tokens.weight.data).to(device)
# Merge weights for each transformer layer
with tqdm(total=len(base_model1.model.layers), desc="Merging weights for layers") as pbar:
for i in range(len(base_model1.model.layers)):
base_model1.model.layers[i].input_layernorm.weight.data = (alpha * base_model1.model.layers[i].input_layernorm.weight.data + (1 - alpha) * base_model2.model.layers[i].input_layernorm.weight.data).to(device)
base_model1.model.layers[i].mlp.down_proj.weight.data = (alpha * base_model1.model.layers[i].mlp.down_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].mlp.down_proj.weight.data).to(device)
base_model1.model.layers[i].mlp.gate_proj.weight.data = (alpha * base_model1.model.layers[i].mlp.gate_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].mlp.gate_proj.weight.data).to(device)
base_model1.model.layers[i].mlp.up_proj.weight.data = (alpha * base_model1.model.layers[i].mlp.up_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].mlp.up_proj.weight.data).to(device)
base_model1.model.layers[i].post_attention_layernorm.weight.data = (alpha * base_model1.model.layers[i].post_attention_layernorm.weight.data + (1 - alpha) * base_model2.model.layers[i].post_attention_layernorm.weight.data).to(device)
base_model1.model.layers[i].self_attn.q_proj.weight.data = (alpha * base_model1.model.layers[i].self_attn.q_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].self_attn.q_proj.weight.data).to(device)
base_model1.model.layers[i].self_attn.k_proj.weight.data = (alpha * base_model1.model.layers[i].self_attn.k_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].self_attn.k_proj.weight.data).to(device)
base_model1.model.layers[i].self_attn.v_proj.weight.data = (alpha * base_model1.model.layers[i].self_attn.v_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].self_attn.v_proj.weight.data).to(device)
base_model1.model.layers[i].self_attn.o_proj.weight.data = (alpha * base_model1.model.layers[i].self_attn.o_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].self_attn.o_proj.weight.data).to(device)
pbar.update(1) # Update the progress bar
# Merge weights for the final normalization layer
base_model1.model.norm.weight.data = (alpha * base_model1.model.norm.weight.data + (1 - alpha) * base_model2.model.norm.weight.data).to(device)
return base_model1
# Merge the weights of the two models with a blending ratio of 0.9
merged_model = merge_model_weights(base_model1, base_model2, alpha=0.9)
# Save the merged model and tokenizer
merged_model.save_pretrained(output_model_dir)
tokenizer.save_pretrained(output_model_dir)
print(f"Merged model and tokenizer saved to {output_model_dir}")
|