JiRackTernaryUltra_14b.py
19.6 KB · 466 lines · python Raw
1 # =============================================================================
2 # COPYRIGHT © 2026 Konstantin Vladimirovich Grabko. ALL RIGHTS RESERVED.
3 # JiRack Ultra Ternary Transformer
4 #
5 # CMS Manhattan JiRack Technology — PATENT PENDING
6 #
7 # This code is proprietary.
8 # Personal and non-commercial research use is allowed.
9 # Any commercial use, derivative works for profit, or distribution
10 # requires a paid license and 5% royalty.
11 #
12 # Unauthorized commercial use is strictly prohibited.
13 # Contact: grabko@cmsmanhattan.com
14 # =============================================================================
15 import math
16
17 import torch
18 import torch.nn as nn
19 import torch.nn.functional as F
20 from torch.utils.checkpoint import checkpoint
21
22 # ==================== CONFIG CONSTANTS [DS14-1..3] ====================
23 VOCAB_SIZE = 152064
24 HIDDEN_SIZE = 5120
25 INTERMEDIATE_SIZE = 13824
26 NUM_LAYERS = 48
27 NUM_HEADS = 40
28 NUM_KV_HEADS = 8
29 HEAD_DIM = 128 # 40 * 128 = 5120 = hidden (q); kv dim = 8*128 = 1024
30 MAX_SEQ_LEN = 4096 # raise for long-context (ckpt supports 131072)
31 ROPE_THETA = 1000000.0 # [DS14-2] Qwen2.5-14B value
32 RMS_EPS = 1e-5 # [DS14-3] Qwen2.5-14B value
33 ROPE_SCALE_FACTOR = 1.0
34 INIT_STD = 0.02
35 ATTN_QKV_BIAS = True # [DS14-5] Qwen2: bias on q/k/v only
36 # =================================================================
37
38
39 # [FIX-6] Feature-detect native GQA support in SDPA (PyTorch >= 2.5).
40 def _detect_sdpa_gqa() -> bool:
41 try:
42 q = torch.zeros(1, 2, 1, 8)
43 kv = torch.zeros(1, 1, 1, 8)
44 F.scaled_dot_product_attention(q, kv, kv, enable_gqa=True)
45 return True
46 except TypeError:
47 return False
48 except Exception:
49 return False
50
51 _SDPA_HAS_GQA = _detect_sdpa_gqa()
52
53
54 class JiRackConfig:
55 def __init__(self):
56 self.vocab_size = VOCAB_SIZE
57 self.hidden_size = HIDDEN_SIZE
58 self.intermediate_size = INTERMEDIATE_SIZE
59 self.num_hidden_layers = NUM_LAYERS
60 self.num_attention_heads = NUM_HEADS
61 self.num_key_value_heads = NUM_KV_HEADS
62 self.head_dim = HEAD_DIM
63 self.max_seq_len = MAX_SEQ_LEN
64 self.rope_theta = ROPE_THETA
65 self.rms_norm_eps = RMS_EPS
66 self.rope_scale_factor = ROPE_SCALE_FACTOR
67 self.init_std = INIT_STD
68 self.attn_qkv_bias = ATTN_QKV_BIAS
69
70
71 # ==================== RoPE — HALF-SPLIT (HF convention) [DS-3] ====================
72 def precompute_freqs_cis(
73 dim: int,
74 end: int,
75 theta: float = ROPE_THETA,
76 scale_factor: float = ROPE_SCALE_FACTOR,
77 ):
78 """cos/sin of shape (end, dim), HF half-split layout: the (dim/2)
79 frequency vector is CONCATENATED with itself (torch.cat), not
80 interleaved. Matches transformers' LlamaRotaryEmbedding/Qwen2."""
81 freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
82 if scale_factor > 1.0:
83 freqs = freqs / scale_factor
84 t = torch.arange(end, dtype=torch.float32)
85 freqs = torch.outer(t, freqs) # (end, dim/2)
86 emb = torch.cat((freqs, freqs), dim=-1) # (end, dim) — half-split
87 return torch.cos(emb), torch.sin(emb)
88
89
90 def rotate_half(x):
91 """HF convention: (-x2, x1) where x1/x2 are the two HALVES of head_dim."""
92 x1 = x[..., : x.shape[-1] // 2]
93 x2 = x[..., x.shape[-1] // 2:]
94 return torch.cat((-x2, x1), dim=-1)
95
96
97 def apply_rotary_emb(xq, xk, cos, sin):
98 """Half-split RoPE, identical math to transformers.apply_rotary_pos_emb.
99 cos/sin: (T, head_dim); q/k: (B, H, T, head_dim)."""
100 cos = cos[None, None, :, :]
101 sin = sin[None, None, :, :]
102 xq_out = (xq * cos) + (rotate_half(xq) * sin)
103 xk_out = (xk * cos) + (rotate_half(xk) * sin)
104 return xq_out, xk_out
105
106
107 class BitLinear(nn.Linear):
108 """BitNet b1.58-style fake-quant linear with lambda warmup.
109 Identical to the 7B version ([FIX-1..5] preserved); bias — when
110 present ([DS-2]) — stays full precision ([DS-5])."""
111
112 def __init__(self, in_features, out_features, bias=False):
113 super().__init__(in_features, out_features, bias=bias)
114 self.eps = 1e-5
115 # [FIX-4] Buffer -> saved in state_dict, survives checkpoint resume.
116 self.register_buffer("lambda_", torch.zeros(()), persistent=True)
117
118 def forward(self, x: torch.Tensor) -> torch.Tensor:
119 # Fast path (exact at lambda=0 by continuity).
120 if not self.training and float(self.lambda_) < 1e-6:
121 return F.linear(x, self.weight, self.bias)
122
123 lam = self.lambda_.to(x.dtype)
124
125 # === Weights: per-tensor absmean ternary (b1.58) ===
126 w = self.weight
127 gamma = w.float().abs().mean().clamp(min=self.eps).to(w.dtype) # [FIX-5]
128 w_quant = torch.clamp(torch.round(w / gamma), -1.0, 1.0) * gamma
129 w_effective = w + lam * (w_quant - w).detach()
130
131 # === Activations: per-token absmax int8 ([FIX-2],[FIX-3]) ===
132 x_scale = 127.0 / x.abs().max(dim=-1, keepdim=True).values.clamp(min=self.eps)
133 x_quant = torch.clamp(torch.round(x * x_scale), -128.0, 127.0) / x_scale
134 x_effective = x + lam * (x_quant - x).detach()
135
136 # [FIX-1] Dequantized operands -> no post-matmul rescale.
137 # [DS-5] bias added in full precision by F.linear.
138 return F.linear(x_effective, w_effective, self.bias)
139
140
141 class RMSNorm(nn.Module):
142 def __init__(self, dim, eps=RMS_EPS):
143 super().__init__()
144 self.eps = eps
145 self.weight = nn.Parameter(torch.ones(dim))
146
147 def forward(self, x):
148 # [FIX-5] Compute statistics in fp32, cast back to input dtype.
149 dtype = x.dtype
150 x = x.float()
151 x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
152 return (x * self.weight.float()).to(dtype)
153
154
155 class TransformerBlock(nn.Module):
156 def __init__(self, config, use_checkpoint=False):
157 super().__init__()
158 self.use_checkpoint = use_checkpoint
159 self.n_heads = config.num_attention_heads
160 self.n_kv_heads = config.num_key_value_heads
161 self.head_dim = config.head_dim
162 self.n_rep = self.n_heads // self.n_kv_heads
163
164 self.norm1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
165 self.norm2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
166
167 qkv_bias = config.attn_qkv_bias # [DS-2]
168 self.q_proj = BitLinear(config.hidden_size,
169 self.n_heads * self.head_dim, bias=qkv_bias)
170 self.k_proj = BitLinear(config.hidden_size,
171 self.n_kv_heads * self.head_dim, bias=qkv_bias)
172 self.v_proj = BitLinear(config.hidden_size,
173 self.n_kv_heads * self.head_dim, bias=qkv_bias)
174 self.out_proj = BitLinear(self.n_heads * self.head_dim,
175 config.hidden_size, bias=False)
176
177 self.ffn_w1 = BitLinear(config.hidden_size, config.intermediate_size, bias=False) # gate
178 self.ffn_w3 = BitLinear(config.hidden_size, config.intermediate_size, bias=False) # up
179 self.ffn_w2 = BitLinear(config.intermediate_size, config.hidden_size, bias=False) # down
180
181 def forward(self, x, freqs_cos, freqs_sin):
182 if self.use_checkpoint and self.training:
183 return checkpoint(
184 self._forward_impl, x, freqs_cos, freqs_sin, use_reentrant=False
185 )
186 return self._forward_impl(x, freqs_cos, freqs_sin)
187
188 def _forward_impl(self, x, freqs_cos, freqs_sin):
189 h = self.norm1(x)
190 B, T, _ = h.shape
191
192 q = self.q_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
193 k = self.k_proj(h).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
194 v = self.v_proj(h).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
195
196 q, k = apply_rotary_emb(q, k, freqs_cos, freqs_sin) # [DS-3] half-split
197
198 if self.n_rep > 1 and _SDPA_HAS_GQA: # [FIX-6]
199 attn_out = F.scaled_dot_product_attention(
200 q, k, v, is_causal=True, enable_gqa=True
201 )
202 else:
203 if self.n_rep > 1:
204 k = k.repeat_interleave(self.n_rep, dim=1)
205 v = v.repeat_interleave(self.n_rep, dim=1)
206 attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
207
208 attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, -1)
209
210 x = x + self.out_proj(attn_out)
211
212 m = self.norm2(x)
213 gate = F.silu(self.ffn_w1(m))
214 up = self.ffn_w3(m)
215 x = x + self.ffn_w2(gate * up)
216
217 return x
218
219
220 class JiRackTransformer(nn.Module):
221 def __init__(self, config: JiRackConfig = None, use_checkpoint=False):
222 super().__init__()
223 self.config = config if config is not None else JiRackConfig()
224 self.use_checkpoint = use_checkpoint
225
226 self.token_emb = nn.Embedding(self.config.vocab_size, self.config.hidden_size)
227 self.blocks = nn.ModuleList([
228 TransformerBlock(self.config, self.use_checkpoint)
229 for _ in range(self.config.num_hidden_layers)
230 ])
231 self.ln_f = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
232 # tie_word_embeddings = False in the 14B distill -> separate lm_head.
233 self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
234
235 cos, sin = precompute_freqs_cis(
236 dim=self.config.head_dim,
237 end=self.config.max_seq_len,
238 theta=self.config.rope_theta,
239 scale_factor=self.config.rope_scale_factor,
240 )
241 self.register_buffer("freqs_cos", cos, persistent=False)
242 self.register_buffer("freqs_sin", sin, persistent=False)
243
244 # [FIX-8] Only relevant when training from scratch; harmless before
245 # load_hf_state_dict() overwrites everything.
246 self._init_weights()
247
248 def _init_weights(self):
249 std = self.config.init_std
250 resid_std = std / math.sqrt(2 * self.config.num_hidden_layers)
251
252 nn.init.normal_(self.token_emb.weight, mean=0.0, std=std)
253 nn.init.normal_(self.lm_head.weight, mean=0.0, std=std)
254
255 for block in self.blocks:
256 for lin in (block.q_proj, block.k_proj, block.v_proj,
257 block.ffn_w1, block.ffn_w3):
258 nn.init.normal_(lin.weight, mean=0.0, std=std)
259 if lin.bias is not None:
260 nn.init.zeros_(lin.bias)
261 for lin in (block.out_proj, block.ffn_w2):
262 nn.init.normal_(lin.weight, mean=0.0, std=resid_std)
263 if lin.bias is not None:
264 nn.init.zeros_(lin.bias)
265
266 # ---------------- lambda warmup hooks (unchanged) ----------------
267 def set_lambda(self, lambda_value: float):
268 for module in self.modules():
269 if isinstance(module, BitLinear):
270 module.lambda_.fill_(lambda_value)
271
272 def get_lambda(self) -> float:
273 for module in self.modules():
274 if isinstance(module, BitLinear):
275 return float(module.lambda_)
276 return 0.0
277
278 def forward(self, input_ids):
279 seq_len = input_ids.shape[1]
280 x = self.token_emb(input_ids)
281
282 cos = self.freqs_cos[:seq_len].to(device=x.device, dtype=x.dtype)
283 sin = self.freqs_sin[:seq_len].to(device=x.device, dtype=x.dtype)
284
285 for block in self.blocks:
286 x = block(x, cos, sin)
287
288 return self.lm_head(self.ln_f(x))
289
290 # ------------------------------------------------------------------
291 # [DS-4] HF Qwen2 -> JiRack weight mapping.
292 # Usage:
293 # from transformers import AutoModelForCausalLM
294 # hf = AutoModelForCausalLM.from_pretrained(
295 # "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B", torch_dtype=torch.float32)
296 # model.load_hf_state_dict(hf.state_dict())
297 # or load safetensors shards directly and merge them into one dict.
298 # No RoPE permutation is needed: this model uses the same half-split
299 # rotation as HF ([DS-3]).
300 # ------------------------------------------------------------------
301 @torch.no_grad()
302 def load_hf_state_dict(self, hf_sd: dict, strict: bool = True):
303 mapped = {}
304 mapped["token_emb.weight"] = hf_sd["model.embed_tokens.weight"]
305 mapped["ln_f.weight"] = hf_sd["model.norm.weight"]
306 if "lm_head.weight" in hf_sd:
307 mapped["lm_head.weight"] = hf_sd["lm_head.weight"]
308 else:
309 # tied-embedding checkpoints (e.g. the 1.5B distill)
310 mapped["lm_head.weight"] = hf_sd["model.embed_tokens.weight"]
311
312 for i in range(self.config.num_hidden_layers):
313 hf = f"model.layers.{i}"
314 jr = f"blocks.{i}"
315 mapped[f"{jr}.norm1.weight"] = hf_sd[f"{hf}.input_layernorm.weight"]
316 mapped[f"{jr}.norm2.weight"] = hf_sd[f"{hf}.post_attention_layernorm.weight"]
317
318 mapped[f"{jr}.q_proj.weight"] = hf_sd[f"{hf}.self_attn.q_proj.weight"]
319 mapped[f"{jr}.k_proj.weight"] = hf_sd[f"{hf}.self_attn.k_proj.weight"]
320 mapped[f"{jr}.v_proj.weight"] = hf_sd[f"{hf}.self_attn.v_proj.weight"]
321 mapped[f"{jr}.q_proj.bias"] = hf_sd[f"{hf}.self_attn.q_proj.bias"]
322 mapped[f"{jr}.k_proj.bias"] = hf_sd[f"{hf}.self_attn.k_proj.bias"]
323 mapped[f"{jr}.v_proj.bias"] = hf_sd[f"{hf}.self_attn.v_proj.bias"]
324 mapped[f"{jr}.out_proj.weight"] = hf_sd[f"{hf}.self_attn.o_proj.weight"]
325
326 mapped[f"{jr}.ffn_w1.weight"] = hf_sd[f"{hf}.mlp.gate_proj.weight"]
327 mapped[f"{jr}.ffn_w3.weight"] = hf_sd[f"{hf}.mlp.up_proj.weight"]
328 mapped[f"{jr}.ffn_w2.weight"] = hf_sd[f"{hf}.mlp.down_proj.weight"]
329
330 missing, unexpected = self.load_state_dict(mapped, strict=False)
331 # lambda_ buffers are OURS (not in HF) — they legitimately stay missing.
332 real_missing = [k for k in missing if not k.endswith("lambda_")]
333 if strict:
334 assert not real_missing, f"missing from HF checkpoint: {real_missing}"
335 assert not unexpected, f"unexpected keys: {unexpected}"
336 print(f"✅ HF weights loaded: {len(mapped)} tensors "
337 f"({len(real_missing)} missing, {len(unexpected)} unexpected)")
338 return real_missing, unexpected
339
340 # ------------------------------------------------------------------
341 # [FIX-9] Ternary export (biases included, full precision, [DS-5]).
342 # ------------------------------------------------------------------
343 @torch.no_grad()
344 def export_ternary_state_dict(self):
345 out = {}
346 for name, module in self.named_modules():
347 if isinstance(module, BitLinear):
348 w = module.weight.float()
349 gamma = w.abs().mean().clamp(min=module.eps)
350 codes = torch.clamp(torch.round(w / gamma), -1.0, 1.0).to(torch.int8)
351 out[f"{name}.codes"] = codes
352 out[f"{name}.gamma"] = gamma
353 if module.bias is not None:
354 out[f"{name}.bias"] = module.bias.detach().clone()
355 out["token_emb.weight"] = self.token_emb.weight.detach().clone()
356 out["lm_head.weight"] = self.lm_head.weight.detach().clone()
357 out["ln_f.weight"] = self.ln_f.weight.detach().clone()
358 for name, module in self.named_modules():
359 if isinstance(module, RMSNorm) and name != "ln_f":
360 out[f"{name}.weight"] = module.weight.detach().clone()
361 return out
362
363
364 # Convenience aliases so the training script barely changes:
365 JiRackConfig = JiRackConfig # drop-in name compat (optional)
366 JiRackTransformer = JiRackTransformer
367
368
369 # =============================================================================
370 # Smoke test: python JiRackTernaryUltra_14b.py (tiny config, CPU, seconds)
371 # =============================================================================
372 if __name__ == "__main__":
373 class TinyConfig(JiRackConfig):
374 def __init__(self):
375 super().__init__()
376 self.vocab_size = 256
377 self.hidden_size = 64
378 self.intermediate_size = 128
379 self.num_hidden_layers = 2
380 self.num_attention_heads = 4
381 self.num_key_value_heads = 2
382 self.head_dim = 16
383 self.max_seq_len = 64
384
385 torch.manual_seed(0)
386 model = JiRackTransformer(TinyConfig()).eval()
387 ids = torch.randint(0, 256, (2, 32))
388
389 with torch.no_grad():
390 model.set_lambda(0.0)
391 y0 = model(ids)
392 model.set_lambda(1e-4)
393 y_eps = model(ids)
394 model.set_lambda(1.0)
395 y1 = model(ids)
396
397 # 1) Continuity in lambda.
398 rel_jump = (y_eps - y0).norm() / y0.norm()
399 print(f"relative change at lambda=1e-4: {rel_jump.item():.2e} (must be ~1e-4)")
400 assert rel_jump < 1e-2, "lambda warmup is not continuous!"
401
402 # 2) Output scale sanity at full quantization.
403 ratio = y1.std() / y0.std()
404 print(f"std ratio lambda=1 vs lambda=0: {ratio.item():.3f} (must be O(1))")
405 assert 0.1 < ratio.item() < 10.0, "output scale collapsed or exploded!"
406
407 # 3) Gradients flow through STE at lambda=1 (weights AND qkv biases).
408 model.train()
409 model.set_lambda(1.0)
410 loss = model(ids).float().pow(2).mean()
411 loss.backward()
412 g = model.blocks[0].q_proj.weight.grad
413 gb = model.blocks[0].q_proj.bias.grad
414 assert g is not None and torch.isfinite(g).all() and g.abs().sum() > 0
415 assert gb is not None and torch.isfinite(gb).all(), "qkv bias got no grad!"
416 print(f"grad norms q_proj: weight={g.norm().item():.4f}, bias={gb.norm().item():.4f}")
417
418 # 4) lambda survives a state_dict round-trip.
419 sd = model.state_dict()
420 model2 = JiRackTransformer(TinyConfig())
421 model2.load_state_dict(sd)
422 assert abs(model2.get_lambda() - 1.0) < 1e-9, "lambda_ not serialized!"
423 print("lambda serialization: OK")
424
425 # 5) HF name mapping round-trip on the tiny config: build a fake HF
426 # dict from our own weights, load it back, outputs must match.
427 fake_hf = {
428 "model.embed_tokens.weight": model.token_emb.weight.detach().clone(),
429 "model.norm.weight": model.ln_f.weight.detach().clone(),
430 "lm_head.weight": model.lm_head.weight.detach().clone(),
431 }
432 for i, blk in enumerate(model.blocks):
433 p = f"model.layers.{i}"
434 fake_hf[f"{p}.input_layernorm.weight"] = blk.norm1.weight.detach().clone()
435 fake_hf[f"{p}.post_attention_layernorm.weight"] = blk.norm2.weight.detach().clone()
436 fake_hf[f"{p}.self_attn.q_proj.weight"] = blk.q_proj.weight.detach().clone()
437 fake_hf[f"{p}.self_attn.k_proj.weight"] = blk.k_proj.weight.detach().clone()
438 fake_hf[f"{p}.self_attn.v_proj.weight"] = blk.v_proj.weight.detach().clone()
439 fake_hf[f"{p}.self_attn.q_proj.bias"] = blk.q_proj.bias.detach().clone()
440 fake_hf[f"{p}.self_attn.k_proj.bias"] = blk.k_proj.bias.detach().clone()
441 fake_hf[f"{p}.self_attn.v_proj.bias"] = blk.v_proj.bias.detach().clone()
442 fake_hf[f"{p}.self_attn.o_proj.weight"] = blk.out_proj.weight.detach().clone()
443 fake_hf[f"{p}.mlp.gate_proj.weight"] = blk.ffn_w1.weight.detach().clone()
444 fake_hf[f"{p}.mlp.up_proj.weight"] = blk.ffn_w3.weight.detach().clone()
445 fake_hf[f"{p}.mlp.down_proj.weight"] = blk.ffn_w2.weight.detach().clone()
446
447 model3 = JiRackTransformer(TinyConfig()).eval()
448 model3.load_hf_state_dict(fake_hf)
449 model3.set_lambda(0.0)
450 model.eval(); model.set_lambda(0.0)
451 with torch.no_grad():
452 y_ref = model(ids)
453 y_map = model3(ids)
454 assert torch.allclose(y_ref, y_map, atol=1e-5), "HF mapping mismatch!"
455 print("HF name-mapping round-trip: OK")
456
457 # 6) Export produces genuinely ternary codes + preserved biases.
458 exported = model.export_ternary_state_dict()
459 codes = exported["blocks.0.q_proj.codes"]
460 assert set(codes.unique().tolist()) <= {-1, 0, 1}
461 assert "blocks.0.q_proj.bias" in exported, "qkv bias lost in export!"
462 print(f"export: {sum('codes' in k for k in exported)} ternary tensors, "
463 f"biases preserved, OK")
464
465 print("\nAll smoke tests passed.")
466