bitnet.cpp/JiRackTernaryPyTorch_1b_bitnet.py · CMSManhattan/JiRackTernaryPro

JiRackTernaryPro_1b / bitnet.cpp /JiRackTernaryPyTorch_1b_bitnet.py

Update bitnet.cpp/JiRackTernaryPyTorch_1b_bitnet.py

6cac314 verified 2 months ago

17.7 kB

# =============================================================================
# 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.")

	# =============================================================================
	# 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.")