File size: 3,002 Bytes
7500ff7 | 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 69 70 71 72 73 74 75 76 77 78 79 80 | import torch
import sys
from transplant_weights import QuillanConfig
def migrate(ckpt_path="quillan_transplanted.pt", out_path="quillan_transplanted_v8.pt"):
print(f"Loading {ckpt_path}...")
ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)
sd = ckpt['model_state_dict']
new_sd = {}
num_experts = 34
hidden_dim = 2048
ffn_dim = 4096
# Initialize the new parameter tensors for MoE
w1 = torch.zeros(num_experts, hidden_dim, ffn_dim)
wgate = torch.zeros(num_experts, hidden_dim, ffn_dim)
w2 = torch.zeros(num_experts, ffn_dim, hidden_dim)
# Initialize wgate normally since it didn't exist (using normal since Kaiming uses uniform or normal)
torch.nn.init.kaiming_normal_(wgate.view(-1, ffn_dim))
# Check old layers to map
for k, v in sd.items():
if k == 'txt_emb.weight':
new_sd['ingestion.txt_emb.weight'] = v
elif k == 'mod_emb.weight':
new_sd['ingestion.mod_emb.weight'] = v
elif k.startswith('experts.'):
parts = k.split('.')
idx = int(parts[1])
if parts[2] == 'w1':
# v shape: [ffn_dim, hidden_dim]
w1[idx] = v.T
elif parts[2] == 'w2':
# v shape: [hidden_dim, ffn_dim]
w2[idx] = v.T
elif parts[2] == 'swarm':
# e.g. experts.0.swarm.lora_A -> moe.expert_swarms.0.lora_A (if using lora_A/B in CouncilExpertSwarm)
# Let's map it directly. Wait, in old model it was swarm.A and swarm.B?
if parts[3] == 'A':
new_sd[f'moe.expert_swarms.{idx}.A'] = v
elif parts[3] == 'B':
new_sd[f'moe.expert_swarms.{idx}.B'] = v
else:
new_sd[f'moe.expert_swarms.{idx}.{parts[3]}'] = v
elif k == 'router.weight':
new_sd['moe.router.weight'] = v
elif k == 'router.bias':
new_sd['moe.router.bias'] = v
elif k.startswith('diffusion.0.'):
# diffusion.0.q_proj.weight -> diffusion_core.q_proj.weight
new_k = k.replace('diffusion.0.', 'diffusion_core.')
new_sd[new_k] = v
elif k.startswith('diffusion.'):
# Ignore diffusion layers > 0
pass
elif k == 'quillan_finalizer.bias':
# Ignore bias if the new layer doesn't use it, but if it's there we keep it.
new_sd[k] = v
elif k == 'quillan_finalizer.weight':
new_sd[k] = v
elif k == 'txt_dec.weight':
new_sd[k] = v
elif k.startswith('decomposition.'):
new_sd[k] = v
else:
new_sd[k] = v
new_sd['moe.w1'] = w1
new_sd['moe.w2'] = w2
new_sd['moe.wgate'] = wgate
print(f"Saving to {out_path}...")
torch.save({'model_state_dict': new_sd}, out_path)
print("Done!")
if __name__ == "__main__":
migrate()
|