File size: 3,093 Bytes
7e1725c |
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 61 62 63 64 65 66 67 68 |
# prepare_donor.py
import torch
import os
import argparse
from tqdm import tqdm
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
def main(source_model_id, output_path):
"""
Loads a Qwen2.5 model, removes all '.bias' tensors, adds placeholder
'q_norm.weight' and 'k_norm.weight' tensors, and saves the result.
This creates an architecturally compatible donor for a Qwen3 merge.
"""
print(f"Loading source donor model: {source_model_id}")
# Load on CPU to save VRAM
model = AutoModelForCausalLM.from_pretrained(
source_model_id,
torch_dtype=torch.bfloat16,
device_map="cpu",
trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(source_model_id, trust_remote_code=True)
config = model.config
source_state_dict = model.state_dict()
new_state_dict = {}
# --- Part 1: Remove '.bias' tensors ---
print("Removing all '.bias' tensors...")
for name, tensor in tqdm(source_state_dict.items(), desc="Filtering Tensors"):
if not name.endswith(".bias"):
new_state_dict[name] = tensor
# --- Part 2: Add placeholder 'q_norm' and 'k_norm' tensors ---
print("Adding placeholder 'q_norm' and 'k_norm' tensors...")
# These norms are 1D vectors of size `head_dim` (128)
# A value of 1.0 is a standard, neutral initialization for a norm weight.
norm_dim = config.hidden_size // config.num_attention_heads # Should be 128 for this model
placeholder_norm = torch.ones(norm_dim, dtype=torch.bfloat16)
for i in tqdm(range(config.num_hidden_layers), desc="Adding Norm Tensors"):
q_norm_name = f"model.layers.{i}.self_attn.q_norm.weight"
k_norm_name = f"model.layers.{i}.self_attn.k_norm.weight"
new_state_dict[q_norm_name] = placeholder_norm.clone()
new_state_dict[k_norm_name] = placeholder_norm.clone()
# The original model is a fine container, we just need to load the modified state dict.
# strict=False is crucial because we have removed and added keys.
print("Loading the new state dict back into the model shell...")
model.load_state_dict(new_state_dict, strict=False, assign=True)
print(f"Saving the architecturally aligned model to: {output_path}")
os.makedirs(output_path, exist_ok=True)
model.save_pretrained(output_path)
tokenizer.save_pretrained(output_path)
print("\nDonor preparation complete!")
print(f"The aligned donor is ready at '{output_path}'.")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Prepare a Qwen2.5 donor model for merging with Qwen3.")
parser.add_argument("--source_model", type=str, default="Qwen/Qwen2.5-72B-Instruct", help="The Hugging Face model ID of the source model.")
parser.add_argument("--output_path", type=str, required=True, help="The local directory path to save the prepared donor model.")
args = parser.parse_args()
# Example: python prepare_donor.py --output_path ./Qwen2.5-72B-Instruct-Aligned
main(args.source_model, args.output_path)
|