Upload merge_model_weights.py

merge_model_weights.py

@@ -0,0 +1,59 @@
+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}")