Create JiRackTernaryPyTorch_1b_bitnet.py

bitnet.cpp/JiRackTernaryPyTorch_1b_bitnet.py

@@ -0,0 +1,412 @@
+# =============================================================================
+# COPYRIGHT (c) 2025-2026 Konstantin Vladimirovich Grabko. ALL RIGHTS RESERVED.
+# CMS Manhattan JiRack Technology - PATENT PENDING
+# 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.")