# ============================================================================= # COPYRIGHT © 2025-2026 Konstantin Vladimirovich Grabko. ALL RIGHTS RESERVED. # CMS Manhattan JiRack Technology — PATENT PENDING # # This code is proprietary. # Personal and non-commercial research use is allowed. # Any commercial use, derivative works for profit, or distribution # requires a paid license and 5% royalty. # # Unauthorized commercial use is strictly prohibited. # Contact: grabko@cmsmanhattan.com # ============================================================================= # # JiRack Ternary Transformer 1B - "bitnet.cpp-compatible" revision # # Derived from JiRackTernaryPyTorch_1b.py (CMSManhattan/JiRackTernaryPro_1b). # Three deliberate changes vs. that file, everything else preserved verbatim # so the existing trained checkpoint loads 1:1 by name: # # [C-1] SubLN added. Two new learnable RMSNorm per block - # attn_sub_norm (before out_proj) and ffn_sub_norm (before ffn_w2). # bitnet.cpp's BITNET architecture hard-requires these tensors; # without them the GGUF will not load. 16 layers x 2 = 32 new tensors. # # [C-2] lambda_ warmup knob added. The 1B original had NO lambda at all - # BitLinear was always fully quantized. A lambda_ buffer per BitLinear # (16 x 7 = 112) lets the network ease back into quantization while # the brand-new SubLN weights learn their scales. Registered as a # persistent buffer so it survives save/load. # # [C-3] Activation quantization no longer mean-centers. # The original did `x_norm = x - x.mean(dim=-1)` before absmax. # bitnet.cpp's I2_S/TL1/TL2 kernels implement plain absmax with NO # mean subtraction, so keeping the centering would mean training # against math the inference engine does not perform. Dropping it is # a real behavioral change - this is part of why the fine-tune below # is mandatory, not optional. # # RoPE IS UNCHANGED ON PURPOSE. The 1B original uses HF/Llama-style # half-split rotation (rotate_half + freqs repeated as [f0..f31, f0..f31]). # The 10B "_fixed" line used interleaved/NeoX instead. They are NOT # interchangeable - the trained weights encode which one they expect. # Consequence for export: the 1B converter MUST apply the standard # convert_hf_to_gguf permute to q_proj/k_proj. The 10B converter must not. # Do not copy one converter onto the other. # ============================================================================= import math import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.checkpoint import checkpoint # ========================= JIRACK 1B CONFIG CONSTANTS ========================= VOCAB_SIZE = 128256 HIDDEN_SIZE = 2048 INTERMEDIATE_SIZE = 8192 NUM_LAYERS = 16 NUM_HEADS = 32 NUM_KV_HEADS = 8 HEAD_DIM = HIDDEN_SIZE // NUM_HEADS # = 64 MAX_SEQ_LEN = 4096 ROPE_THETA = 500000.0 RMS_EPS = 1e-6 # NOTE: 1e-6 here, the 10B line uses 1e-5 INIT_STD = 0.02 # --- quantization stability, from the original file --- STABILITY_EPS = 1e-9 INT8_SCALE_TARGET = 127.0 QUANT_EPS = 1e-5 # clamp floor used inside BitLinear # ============================================================================= def _detect_sdpa_gqa() -> bool: try: q = torch.zeros(1, 2, 1, 8) kv = torch.zeros(1, 1, 1, 8) F.scaled_dot_product_attention(q, kv, kv, enable_gqa=True) return True except TypeError: return False except Exception: return False _SDPA_HAS_GQA = _detect_sdpa_gqa() class JiRackConfig1B: def __init__(self): self.vocab_size = VOCAB_SIZE self.hidden_size = HIDDEN_SIZE self.intermediate_size = INTERMEDIATE_SIZE self.num_hidden_layers = NUM_LAYERS self.num_attention_heads = NUM_HEADS self.num_key_value_heads = NUM_KV_HEADS self.head_dim = HEAD_DIM self.max_seq_len = MAX_SEQ_LEN self.max_position_embeddings = MAX_SEQ_LEN self.rope_theta = ROPE_THETA self.rms_norm_eps = RMS_EPS self.init_std = INIT_STD # ---------------------------------------------------------------------------- # RoPE - HF / Llama half-split style. IDENTICAL MATH to the 1B original. # freqs has head_dim/2 = 32 entries; it is tiled to 64 as [f0..f31, f0..f31], # and rotate_half splits the vector down the middle. This is what # convert_hf_to_gguf.py's permute() is designed to undo, hence the converter # for this model applies that permute. # ---------------------------------------------------------------------------- def precompute_freqs_cis(dim: int, seq_len: int, theta: float = ROPE_THETA): freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) t = torch.arange(seq_len).float() freqs = torch.outer(t, freqs) return torch.cos(freqs), torch.sin(freqs) def apply_rotary_emb(xq, xk, freqs_cos, freqs_sin): def rotate_half(x): x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) T = xq.shape[2] f_cos = freqs_cos[:T].to(device=xq.device, dtype=xq.dtype).view(1, 1, T, -1).repeat(1, 1, 1, 2) f_sin = freqs_sin[:T].to(device=xq.device, dtype=xq.dtype).view(1, 1, T, -1).repeat(1, 1, 1, 2) xq_out = (xq * f_cos) + (rotate_half(xq) * f_sin) xk_out = (xk * f_cos) + (rotate_half(xk) * f_sin) return xq_out, xk_out class BitLinear(nn.Linear): """Canonical BitNet b1.58 fake-quant linear. No internal normalization - that is the caller's job (norm1/norm2 for the pre-norms, attn_sub_norm/ffn_sub_norm for SubLN). Weight quant is absmean ternary; activation quant is plain absmax int8, NO mean centering, see [C-3] in the header. """ def __init__(self, in_features, out_features, bias=False): super().__init__(in_features, out_features, bias=bias) self.eps = QUANT_EPS # [C-2] persistent so it round-trips through save/load self.register_buffer("lambda_", torch.zeros(()), persistent=True) def forward(self, x: torch.Tensor) -> torch.Tensor: if not self.training and float(self.lambda_) < 1e-6: return F.linear(x, self.weight, self.bias) lam = self.lambda_.to(x.dtype) w = self.weight gamma = w.float().abs().mean().clamp(min=self.eps).to(w.dtype) w_quant = torch.clamp(torch.round(w / gamma), -1.0, 1.0) * gamma w_effective = w + lam * (w_quant - w).detach() x_scale = INT8_SCALE_TARGET / x.abs().max(dim=-1, keepdim=True).values.clamp(min=self.eps) x_quant = torch.clamp(torch.round(x * x_scale), -128.0, 127.0) / x_scale x_effective = x + lam * (x_quant - x).detach() return F.linear(x_effective, w_effective, self.bias) class RMSNorm(nn.Module): """Learnable RMSNorm. Used for the block pre-norms AND for SubLN - same math, different position in the graph.""" def __init__(self, dim, eps=RMS_EPS): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): dtype = x.dtype x = x.float() x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return (x * self.weight.float()).to(dtype) def repeat_kv(x, n_rep): if n_rep == 1: return x return x.repeat_interleave(n_rep, dim=1) class TransformerBlock(nn.Module): def __init__(self, config, use_checkpoint=False, bias=False): super().__init__() self.use_checkpoint = use_checkpoint self.n_heads = config.num_attention_heads self.n_kv_heads = config.num_key_value_heads self.head_dim = config.head_dim self.n_rep = self.n_heads // self.n_kv_heads self.q_proj = BitLinear(config.hidden_size, config.hidden_size, bias=bias) self.k_proj = BitLinear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=bias) self.v_proj = BitLinear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=bias) self.out_proj = BitLinear(config.hidden_size, config.hidden_size, bias=bias) self.ffn_w1 = BitLinear(config.hidden_size, config.intermediate_size, bias=bias) self.ffn_w3 = BitLinear(config.hidden_size, config.intermediate_size, bias=bias) self.ffn_w2 = BitLinear(config.intermediate_size, config.hidden_size, bias=bias) self.norm1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.norm2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # [C-1] SubLN, exactly where bitnet.cpp expects it self.attn_sub_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.ffn_sub_norm = RMSNorm(config.intermediate_size, eps=config.rms_norm_eps) def forward(self, x, freqs_cos, freqs_sin): if self.use_checkpoint and self.training: return checkpoint(self._forward_impl, x, freqs_cos, freqs_sin, use_reentrant=False) return self._forward_impl(x, freqs_cos, freqs_sin) def _forward_impl(self, x, freqs_cos, freqs_sin): h = self.norm1(x) B, T, D = x.shape q = self.q_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = self.k_proj(h).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(h).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) q, k = apply_rotary_emb(q, k, freqs_cos, freqs_sin) if self.n_rep > 1 and _SDPA_HAS_GQA: attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True) else: k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep) attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True) attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, D) # [C-1] SubLN before the ternary output projection attn_out = self.attn_sub_norm(attn_out) x = x + self.out_proj(attn_out) m = self.norm2(x) gate = F.silu(self.ffn_w1(m)) up = self.ffn_w3(m) ffn_hidden = gate * up # [C-1] SubLN before the ternary down projection ffn_hidden = self.ffn_sub_norm(ffn_hidden) x = x + self.ffn_w2(ffn_hidden) return x class JiRackTransformer1B(nn.Module): def __init__(self, config: JiRackConfig1B = None, use_checkpoint=False, bias=False): super().__init__() self.config = config if config is not None else JiRackConfig1B() self.use_checkpoint = use_checkpoint self.token_emb = nn.Embedding(self.config.vocab_size, self.config.hidden_size) self.blocks = nn.ModuleList([ TransformerBlock(self.config, self.use_checkpoint, bias=bias) for _ in range(self.config.num_hidden_layers) ]) self.ln_f = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False) cos, sin = precompute_freqs_cis( dim=self.config.head_dim, seq_len=self.config.max_seq_len, theta=self.config.rope_theta, ) # NOTE: the 1B original registered these as PERSISTENT buffers, so an # old checkpoint may or may not carry freqs_cos/freqs_sin depending on # how it was saved. Non-persistent here (they are recomputed at init); # the migration script tolerates them being present in the old file. self.register_buffer("freqs_cos", cos, persistent=False) self.register_buffer("freqs_sin", sin, persistent=False) self._init_weights() def _init_weights(self): std = self.config.init_std resid_std = std / math.sqrt(2 * self.config.num_hidden_layers) nn.init.normal_(self.token_emb.weight, mean=0.0, std=std) nn.init.normal_(self.lm_head.weight, mean=0.0, std=std) for block in self.blocks: for lin in (block.q_proj, block.k_proj, block.v_proj, block.ffn_w1, block.ffn_w3): nn.init.normal_(lin.weight, mean=0.0, std=std) if lin.bias is not None: nn.init.zeros_(lin.bias) for lin in (block.out_proj, block.ffn_w2): nn.init.normal_(lin.weight, mean=0.0, std=resid_std) if lin.bias is not None: nn.init.zeros_(lin.bias) # SubLN weights stay at their RMSNorm init of 1.0 -> near-identity, # which minimizes the shock of the architecture change. def load_from_old_checkpoint(self, old_state_dict: dict): """Load weights from the ORIGINAL 1B architecture (no SubLN, no lambda_). Tensor names are identical between the two files, so this is a plain strict=False load. Exactly three categories of difference are tolerated: 1. missing *.sub_norm.weight - the 32 new SubLN tensors [C-1]. Left at init 1.0, must be learned during the fine-tune. 2. missing *.lambda_ - the 112 new warmup buffers [C-2]. Left at 0.0, the correct start for a fresh warmup. 3. unexpected freqs_cos/freqs_sin - the original registered its RoPE tables as persistent buffers, so they may be baked into the old file. They are recomputed at init here, so they are ignored. Anything else is a real mismatch and raises. """ missing, unexpected = self.load_state_dict(old_state_dict, strict=False) new_subln = [m for m in missing if "sub_norm" in m] new_lambda = [m for m in missing if m.endswith("lambda_")] other_missing = [m for m in missing if "sub_norm" not in m and not m.endswith("lambda_")] rope_extra = [u for u in unexpected if u in ("freqs_cos", "freqs_sin")] other_unexpected = [u for u in unexpected if u not in ("freqs_cos", "freqs_sin")] if other_missing: raise RuntimeError( f"{len(other_missing)} unexpected missing keys beyond SubLN/lambda_: " f"{other_missing[:10]}" ) if other_unexpected: raise RuntimeError( f"{len(other_unexpected)} unexpected keys in old checkpoint: " f"{other_unexpected[:10]}" ) n_layers = self.config.num_hidden_layers print("Loaded old 1B checkpoint.") print(f" {len(new_subln)} SubLN tensors initialized at 1.0 " f"(expected {2 * n_layers}), e.g. {new_subln[:2]}") print(f" {len(new_lambda)} lambda_ buffers initialized at 0.0 " f"(expected {7 * n_layers}) - warmup starts fresh") if rope_extra: print(f" ignored baked-in RoPE tables: {rope_extra}") def set_lambda(self, lambda_value: float): for module in self.modules(): if isinstance(module, BitLinear): module.lambda_.fill_(lambda_value) def get_lambda(self) -> float: for module in self.modules(): if isinstance(module, BitLinear): return float(module.lambda_) return 0.0 def forward(self, input_ids): x = self.token_emb(input_ids) for block in self.blocks: x = block(x, self.freqs_cos, self.freqs_sin) return self.lm_head(self.ln_f(x)) @torch.no_grad() def export_ternary_state_dict(self): """Bake fake-quant weights into real ternary codes + per-tensor scale. SubLN weights are full precision and pass through untouched, same as embeddings / lm_head / norms.""" out = {} for name, module in self.named_modules(): if isinstance(module, BitLinear): w = module.weight.float() gamma = w.abs().mean().clamp(min=module.eps) codes = torch.clamp(torch.round(w / gamma), -1.0, 1.0).to(torch.int8) out[f"{name}.codes"] = codes out[f"{name}.gamma"] = gamma if module.bias is not None: out[f"{name}.bias"] = module.bias.detach().clone() out["token_emb.weight"] = self.token_emb.weight.detach().clone() out["lm_head.weight"] = self.lm_head.weight.detach().clone() out["ln_f.weight"] = self.ln_f.weight.detach().clone() for name, module in self.named_modules(): if isinstance(module, RMSNorm) and name != "ln_f": out[f"{name}.weight"] = module.weight.detach().clone() return out if __name__ == "__main__": class TinyConfig(JiRackConfig1B): def __init__(self): super().__init__() self.vocab_size = 256 self.hidden_size = 64 self.intermediate_size = 128 self.num_hidden_layers = 2 self.num_attention_heads = 4 self.num_key_value_heads = 2 self.head_dim = 16 self.max_seq_len = 64 torch.manual_seed(0) model = JiRackTransformer1B(TinyConfig()).eval() ids = torch.randint(0, 256, (2, 32)) with torch.no_grad(): model.set_lambda(0.0) y0 = model(ids) model.set_lambda(1.0) y1 = model(ids) assert torch.isfinite(y0).all() and torch.isfinite(y1).all() ratio = y1.std() / y0.std() print(f"std ratio lambda=1 vs lambda=0: {ratio.item():.3f} (should be O(1))") sd = model.state_dict() subln = [n for n in sd if "sub_norm" in n] lam = [n for n in sd if n.endswith("lambda_")] print(f"SubLN tensors: {len(subln)} (expect 2 x n_layers)") print(f"lambda_ buffers: {len(lam)} (expect 7 x n_layers)") print("Smoke test OK.")