Quillan-Ronin / scripts /fix_projections.py

Upload batch 15/16

7500ff7 verified 7 days ago

7.09 kB

#!/usr/bin/env python3
"""Phase 1: Fix expert projections by directly copying source weights.
The 3 source models were sliced into Quillan's 34 experts.
The problem: dimension mismatches were handled with random SVD projections
that corrupted the signal. Fix: pad/trim source weights directly to target dims.
"""
import os, sys, torch, gc
sys.path.insert(0, r'C:\Users\Admin\Quillan-Ronin')
from safetensors import safe_open
from pathlib import Path

BASE = Path(r'C:\Users\Admin\Quillan-Ronin')
MODEL_DIR = BASE / 'Quillan-v4.2-model'
CKPT_PATH = BASE / 'checkpoints' / 'quillan_finetuned.pt'
OUT_PATH = BASE / 'checkpoints' / 'quillan_fixed.pt'
H, F = 2048, 4096

def dequantize_bitnet(w_uint8, w_scale):
    """BitNet stores uint8 with per-tensor scale. 
    Formula: w_fp = (uint8 - 128.0) * scale"""
    return (w_uint8.float() - 128.0) * w_scale.float()

def pad_to_target(w, target_in, target_out):
    """Pad or trim a weight to target dimensions.
    w: [src_in, src_out] in Quillan [in, out] convention
    Returns: [target_in, target_out]
    """
    w = w.float()
    in_dim, out_dim = w.shape
    
    # Pad or trim in dimension
    if in_dim < target_in:
        pad_in = target_in - in_dim
        zeros = torch.zeros(pad_in, out_dim)
        w = torch.cat([w, zeros], dim=0)
    elif in_dim > target_in:
        w = w[:target_in, :]
    
    # Pad or trim out dimension
    if out_dim < target_out:
        pad_out = target_out - out_dim
        zeros = torch.zeros(target_in, pad_out)
        w = torch.cat([w, zeros], dim=1)
    elif out_dim > target_out:
        w = w[:, :target_out]
    
    return w

# ─── Load source models ─────────────────────────────────────────────────
print('=== PHASE 1: FIX EXPERT WEIGHTS ===')
print('\n[1] Loading source models...')

# Llama (C0-C7): direct copies, no changes needed
# Qwen (C8-C21): 14 experts, hidden=1024, ffn=3584
# BitNet (C22-C33): 12 experts, hidden=2560, ffn=1728

# Qwen FFN weights
print('  Qwen FFN weights...')
q_w1, q_wg, q_w2 = [], [], []
with safe_open(str(MODEL_DIR / 'qwen model.safetensors'), framework='pt') as f:
    for i in range(14):
        q_w1.append(f.get_tensor(f'model.language_model.layers.{i}.mlp.gate_proj.weight').float())
        q_wg.append(f.get_tensor(f'model.language_model.layers.{i}.mlp.up_proj.weight').float())
        q_w2.append(f.get_tensor(f'model.language_model.layers.{i}.mlp.down_proj.weight').float())
q_w1 = torch.stack(q_w1)  # [14, 3584, 1024] — [out, in]
q_wg = torch.stack(q_wg)
q_w2 = torch.stack(q_w2)  # [14, 1024, 3584] — [out, in]
print(f'  Qwen w1: {list(q_w1.shape)}, w2: {list(q_w2.shape)}')

# BitNet FFN weights (dequantize from uint8)
print('  BitNet FFN weights...')
b_w1, b_wg, b_w2 = [], [], []
with safe_open(str(MODEL_DIR / 'bitnet model.safetensors'), framework='pt') as f:
    for i in range(12):
        w1_u = f.get_tensor(f'model.layers.{i}.mlp.gate_proj.weight')
        w1_s = f.get_tensor(f'model.layers.{i}.mlp.gate_proj.weight_scale')
        b_w1.append(dequantize_bitnet(w1_u, w1_s))
        
        wg_u = f.get_tensor(f'model.layers.{i}.mlp.up_proj.weight')
        wg_s = f.get_tensor(f'model.layers.{i}.mlp.up_proj.weight_scale')
        b_wg.append(dequantize_bitnet(wg_u, wg_s))
        
        w2_u = f.get_tensor(f'model.layers.{i}.mlp.down_proj.weight')
        w2_s = f.get_tensor(f'model.layers.{i}.mlp.down_proj.weight_scale')
        b_w2.append(dequantize_bitnet(w2_u, w2_s))
b_w1 = torch.stack(b_w1)  # [12, 1728, 2560] — [out, in]
b_wg = torch.stack(b_wg)
b_w2 = torch.stack(b_w2)  # [12, 2560, 1728] — [out, in]
print(f'  BitNet w1: {list(b_w1.shape)}, w2: {list(b_w2.shape)}')

# ─── Convert to Quillan format [experts, in_dim, out_dim] ─────────────
print('\n[2] Converting to Quillan format...')

# Qwen: [14, 3584, 1024] [out, in] → transpose → [14, 1024, 3584] [in, out] → pad → [14, 2048, 4096]
q_w1 = torch.stack([pad_to_target(w, H, F) for w in q_w1.transpose(1, 2)])  # [14, H, F]
q_wg = torch.stack([pad_to_target(w, H, F) for w in q_wg.transpose(1, 2)])
q_w2 = torch.stack([pad_to_target(w, F, H) for w in q_w2.transpose(1, 2)])  # down: [in=F, out=H] → pad to [4096, 2048]

# BitNet: [12, 1728, 2560] [out, in] → transpose → [12, 2560, 1728] [in, out] → trim/pad → [12, 2048, 4096]
# In this case, in_dim=2560 > target=2048, so we TRIM the first 2048 dims
b_w1 = torch.stack([pad_to_target(w, H, F) for w in b_w1.transpose(1, 2)])
b_wg = torch.stack([pad_to_target(w, H, F) for w in b_wg.transpose(1, 2)])
b_w2 = torch.stack([pad_to_target(w, F, H) for w in b_w2.transpose(1, 2)])  # down: [in=F, out=H]

print(f'  Qwen converted: w1={list(q_w1.shape)}')
print(f'  BitNet converted: w1={list(b_w1.shape)}')

# ─── Load checkpoint and replace experts ───────────────────────────────
print('\n[3] Loading checkpoint...')
ckpt = torch.load(CKPT_PATH, map_location='cpu', weights_only=True)
sd = ckpt.get('state_dict', ckpt)
if 'model_state_dict' in ckpt: sd = ckpt['model_state_dict']

print('  Replacing C8-C21 (Qwen)...')
sd['moe.w1'][8:22] = q_w1.to(sd['moe.w1'].dtype)
sd['moe.wgate'][8:22] = q_wg.to(sd['moe.wgate'].dtype)
sd['moe.w2'][8:22] = q_w2.to(sd['moe.w2'].dtype)

# BitNet weights have ~1000x larger magnitude than Llama. Rescale to match.
llama_norm = sd['moe.w1'][:8].norm()
bitnet_norm = b_w1.norm()
scale_factor = llama_norm / bitnet_norm
print(f'  Rescaling BitNet by {scale_factor:.4f} ({llama_norm:.1f} / {bitnet_norm:.1f})')
b_w1 = b_w1 * scale_factor
b_wg = b_wg * scale_factor
b_w2 = b_w2 * scale_factor

print('  Replacing C22-C33 (BitNet)...')
sd['moe.w1'][22:34] = b_w1.to(sd['moe.w1'].dtype)
sd['moe.wgate'][22:34] = b_wg.to(sd['moe.wgate'].dtype)
sd['moe.w2'][22:34] = b_w2.to(sd['moe.w2'].dtype)

del q_w1, q_wg, q_w2, b_w1, b_wg, b_w2; gc.collect()

# ─── Verify ─────────────────────────────────────────────────────────────
s1 = str(tuple(sd['moe.w1'].shape))
s2 = str(tuple(sd['moe.w2'].shape))
sg = str(tuple(sd['moe.wgate'].shape))
dtype = sd['moe.w1'].dtype
print(f'  w1: {s1} {dtype}')
print(f'  w2: {s2} {dtype}')
print(f'  wgate: {sg} {dtype}')

# Quick sanity: check norm of Llama vs Qwen vs BitNet experts
for name, start, end in [('Llama C0', 0, 8), ('Qwen C8', 8, 22), ('BitNet C22', 22, 34)]:
    n = sd['moe.w1'][start:end].norm().item()
    print(f'  {name}: w1_norm={n:.2f}')

# ─── Save ───────────────────────────────────────────────────────────────
print(f'\n[4] Saving to {OUT_PATH.name}...')
torch.save(dict(state_dict=sd, step=0, loss=0.0, version='v5.3.1-fixed'), OUT_PATH)
gb = os.path.getsize(OUT_PATH) / 1e9
print(f'  Saved: {gb:.2f} GB')
print('\nPHASE 1 COMPLETE')

	#!/usr/bin/env python3
	"""Phase 1: Fix expert projections by directly copying source weights.

	The 3 source models were sliced into Quillan's 34 experts.
	The problem: dimension mismatches were handled with random SVD projections
	that corrupted the signal. Fix: pad/trim source weights directly to target dims.
	"""
	import os, sys, torch, gc
	sys.path.insert(0, r'C:\Users\Admin\Quillan-Ronin')
	from safetensors import safe_open
	from pathlib import Path

	BASE = Path(r'C:\Users\Admin\Quillan-Ronin')
	MODEL_DIR = BASE / 'Quillan-v4.2-model'
	CKPT_PATH = BASE / 'checkpoints' / 'quillan_finetuned.pt'
	OUT_PATH = BASE / 'checkpoints' / 'quillan_fixed.pt'
	H, F = 2048, 4096

	def dequantize_bitnet(w_uint8, w_scale):
	"""BitNet stores uint8 with per-tensor scale.
	Formula: w_fp = (uint8 - 128.0) * scale"""
	return (w_uint8.float() - 128.0) * w_scale.float()

	def pad_to_target(w, target_in, target_out):
	"""Pad or trim a weight to target dimensions.
	w: [src_in, src_out] in Quillan [in, out] convention
	Returns: [target_in, target_out]
	"""
	w = w.float()
	in_dim, out_dim = w.shape

	# Pad or trim in dimension
	if in_dim < target_in:
	pad_in = target_in - in_dim
	zeros = torch.zeros(pad_in, out_dim)
	w = torch.cat([w, zeros], dim=0)
	elif in_dim > target_in:
	w = w[:target_in, :]

	# Pad or trim out dimension
	if out_dim < target_out:
	pad_out = target_out - out_dim
	zeros = torch.zeros(target_in, pad_out)
	w = torch.cat([w, zeros], dim=1)
	elif out_dim > target_out:
	w = w[:, :target_out]

	return w

	# ─── Load source models ─────────────────────────────────────────────────
	print('=== PHASE 1: FIX EXPERT WEIGHTS ===')
	print('\n[1] Loading source models...')

	# Llama (C0-C7): direct copies, no changes needed
	# Qwen (C8-C21): 14 experts, hidden=1024, ffn=3584
	# BitNet (C22-C33): 12 experts, hidden=2560, ffn=1728

	# Qwen FFN weights
	print(' Qwen FFN weights...')
	q_w1, q_wg, q_w2 = [], [], []
	with safe_open(str(MODEL_DIR / 'qwen model.safetensors'), framework='pt') as f:
	for i in range(14):
	q_w1.append(f.get_tensor(f'model.language_model.layers.{i}.mlp.gate_proj.weight').float())
	q_wg.append(f.get_tensor(f'model.language_model.layers.{i}.mlp.up_proj.weight').float())
	q_w2.append(f.get_tensor(f'model.language_model.layers.{i}.mlp.down_proj.weight').float())
	q_w1 = torch.stack(q_w1) # [14, 3584, 1024] — [out, in]
	q_wg = torch.stack(q_wg)
	q_w2 = torch.stack(q_w2) # [14, 1024, 3584] — [out, in]
	print(f' Qwen w1: {list(q_w1.shape)}, w2: {list(q_w2.shape)}')

	# BitNet FFN weights (dequantize from uint8)
	print(' BitNet FFN weights...')
	b_w1, b_wg, b_w2 = [], [], []
	with safe_open(str(MODEL_DIR / 'bitnet model.safetensors'), framework='pt') as f:
	for i in range(12):
	w1_u = f.get_tensor(f'model.layers.{i}.mlp.gate_proj.weight')
	w1_s = f.get_tensor(f'model.layers.{i}.mlp.gate_proj.weight_scale')
	b_w1.append(dequantize_bitnet(w1_u, w1_s))

	wg_u = f.get_tensor(f'model.layers.{i}.mlp.up_proj.weight')
	wg_s = f.get_tensor(f'model.layers.{i}.mlp.up_proj.weight_scale')
	b_wg.append(dequantize_bitnet(wg_u, wg_s))

	w2_u = f.get_tensor(f'model.layers.{i}.mlp.down_proj.weight')
	w2_s = f.get_tensor(f'model.layers.{i}.mlp.down_proj.weight_scale')
	b_w2.append(dequantize_bitnet(w2_u, w2_s))
	b_w1 = torch.stack(b_w1) # [12, 1728, 2560] — [out, in]
	b_wg = torch.stack(b_wg)
	b_w2 = torch.stack(b_w2) # [12, 2560, 1728] — [out, in]
	print(f' BitNet w1: {list(b_w1.shape)}, w2: {list(b_w2.shape)}')

	# ─── Convert to Quillan format [experts, in_dim, out_dim] ─────────────
	print('\n[2] Converting to Quillan format...')

	# Qwen: [14, 3584, 1024] [out, in] → transpose → [14, 1024, 3584] [in, out] → pad → [14, 2048, 4096]
	q_w1 = torch.stack([pad_to_target(w, H, F) for w in q_w1.transpose(1, 2)]) # [14, H, F]
	q_wg = torch.stack([pad_to_target(w, H, F) for w in q_wg.transpose(1, 2)])
	q_w2 = torch.stack([pad_to_target(w, F, H) for w in q_w2.transpose(1, 2)]) # down: [in=F, out=H] → pad to [4096, 2048]

	# BitNet: [12, 1728, 2560] [out, in] → transpose → [12, 2560, 1728] [in, out] → trim/pad → [12, 2048, 4096]
	# In this case, in_dim=2560 > target=2048, so we TRIM the first 2048 dims
	b_w1 = torch.stack([pad_to_target(w, H, F) for w in b_w1.transpose(1, 2)])
	b_wg = torch.stack([pad_to_target(w, H, F) for w in b_wg.transpose(1, 2)])
	b_w2 = torch.stack([pad_to_target(w, F, H) for w in b_w2.transpose(1, 2)]) # down: [in=F, out=H]

	print(f' Qwen converted: w1={list(q_w1.shape)}')
	print(f' BitNet converted: w1={list(b_w1.shape)}')

	# ─── Load checkpoint and replace experts ───────────────────────────────
	print('\n[3] Loading checkpoint...')
	ckpt = torch.load(CKPT_PATH, map_location='cpu', weights_only=True)
	sd = ckpt.get('state_dict', ckpt)
	if 'model_state_dict' in ckpt: sd = ckpt['model_state_dict']

	print(' Replacing C8-C21 (Qwen)...')
	sd['moe.w1'][8:22] = q_w1.to(sd['moe.w1'].dtype)
	sd['moe.wgate'][8:22] = q_wg.to(sd['moe.wgate'].dtype)
	sd['moe.w2'][8:22] = q_w2.to(sd['moe.w2'].dtype)

	# BitNet weights have ~1000x larger magnitude than Llama. Rescale to match.
	llama_norm = sd['moe.w1'][:8].norm()
	bitnet_norm = b_w1.norm()
	scale_factor = llama_norm / bitnet_norm
	print(f' Rescaling BitNet by {scale_factor:.4f} ({llama_norm:.1f} / {bitnet_norm:.1f})')
	b_w1 = b_w1 * scale_factor
	b_wg = b_wg * scale_factor
	b_w2 = b_w2 * scale_factor

	print(' Replacing C22-C33 (BitNet)...')
	sd['moe.w1'][22:34] = b_w1.to(sd['moe.w1'].dtype)
	sd['moe.wgate'][22:34] = b_wg.to(sd['moe.wgate'].dtype)
	sd['moe.w2'][22:34] = b_w2.to(sd['moe.w2'].dtype)

	del q_w1, q_wg, q_w2, b_w1, b_wg, b_w2; gc.collect()

	# ─── Verify ─────────────────────────────────────────────────────────────
	s1 = str(tuple(sd['moe.w1'].shape))
	s2 = str(tuple(sd['moe.w2'].shape))
	sg = str(tuple(sd['moe.wgate'].shape))
	dtype = sd['moe.w1'].dtype
	print(f' w1: {s1} {dtype}')
	print(f' w2: {s2} {dtype}')
	print(f' wgate: {sg} {dtype}')

	# Quick sanity: check norm of Llama vs Qwen vs BitNet experts
	for name, start, end in [('Llama C0', 0, 8), ('Qwen C8', 8, 22), ('BitNet C22', 22, 34)]:
	n = sd['moe.w1'][start:end].norm().item()
	print(f' {name}: w1_norm={n:.2f}')

	# ─── Save ───────────────────────────────────────────────────────────────
	print(f'\n[4] Saving to {OUT_PATH.name}...')
	torch.save(dict(state_dict=sd, step=0, loss=0.0, version='v5.3.1-fixed'), OUT_PATH)
	gb = os.path.getsize(OUT_PATH) / 1e9
	print(f' Saved: {gb:.2f} GB')
	print('\nPHASE 1 COMPLETE')