| """Initialize a BitLMv18 binary student's latent weights from a trained FP32LM teacher. |
| |
| Copies: |
| - embedding codebook (FP32 embed → student.embed) |
| - output codebook (FP32 head → student.out_codebook) |
| - per-block attention projections (FP32 qkv split → student q/k/v, o → o) |
| - per-block FFN (FP32 g/u/d → student gate/up/down) |
| |
| Skips (not present in v18): positional embedding, RMSNorm. ALiBi replaces pos. |
| Teacher and student must have identical d_model, n_layers, n_heads, d_ff, vocab. |
| """ |
| import torch |
| from model_v18 import BitLMv18 |
| from model_fp32 import FP32LM |
|
|
|
|
| @torch.no_grad() |
| def init_binary_from_fp32(student: BitLMv18, teacher: FP32LM, mode='sign_rescale', |
| latent_scale=0.02): |
| """ |
| mode='raw' — copy teacher weights verbatim (tried in v36; failed because |
| teacher magnitudes are too large → weights frozen). |
| mode='sign_rescale' — copy sign(teacher) and rescale to `latent_scale`. Student |
| starts with teacher's sign pattern BUT has flippable |
| latent magnitudes. Preferred. |
| """ |
| assert student.d_model == teacher.d_model |
| assert len(student.blocks) == len(teacher.blocks) |
| assert student.vocab_size == teacher.vocab_size |
|
|
| def convert(w): |
| if mode == 'raw': |
| return w.clone() |
| elif mode == 'sign_rescale': |
| s = torch.sign(w) |
| s[s == 0] = 1 |
| return s * latent_scale |
| raise ValueError(mode) |
|
|
| student.embed.weight.copy_(convert(teacher.embed.weight)) |
| student.out_codebook.copy_(convert(teacher.embed.weight)) |
|
|
| for b_blk, t_blk in zip(student.blocks, teacher.blocks): |
| qkv = t_blk.a.qkv.weight |
| d = qkv.shape[1] |
| b_blk.attn.q_proj.raw.weight.copy_(convert(qkv[:d])) |
| b_blk.attn.k_proj.raw.weight.copy_(convert(qkv[d:2*d])) |
| b_blk.attn.v_proj.raw.weight.copy_(convert(qkv[2*d:])) |
| b_blk.attn.o_proj.raw.weight.copy_(convert(t_blk.a.o.weight)) |
|
|
| b_blk.ffn.gate.raw.weight.copy_(convert(t_blk.f.g.weight)) |
| b_blk.ffn.up.raw.weight.copy_(convert(t_blk.f.u.weight)) |
| b_blk.ffn.down.raw.weight.copy_(convert(t_blk.f.d.weight)) |
|
|
|
|
| def load_teacher(path, device='cuda'): |
| ck = torch.load(path, map_location=device, weights_only=False) |
| cfg = ck['args'] |
| t = FP32LM( |
| vocab_size=cfg['vocab_size'], d_model=cfg['d_model'], n_layers=cfg['n_layers'], |
| n_heads=cfg['n_heads'], d_ff=cfg['d_ff'], max_seq_len=cfg['seq_len'], |
| ).to(device) |
| t.load_state_dict(ck['model']) |
| t.eval() |
| return t, ck |
|
|
