JiRackNative_3b.py

# =============================================================================
# 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
# =============================================================================
#
# CHANGE LOG (this revision) — two correctness fixes, no behaviour change to
# the rest of the architecture:
#   FIX 1: RMSNorm now reduces the mean-of-squares in float32 and casts back.
#          Prevents bf16 precision loss that can cause loss/perplexity spikes.
#   FIX 2: BitLinear activation quantization no longer subtracts the mean
#          (per-token absmax, matching BitNet b1.58) and uses a symmetric
#          [-127, 127] clamp. Removes the double-centering vs. RMSNorm and the
#          forward/STE mismatch.
# =============================================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

# --- JIRACK 3B CONSTANTS ---
VOCAB_SIZE = 128256
HIDDEN_SIZE = 3072
NUM_LAYERS = 20
NUM_HEADS = 24
NUM_KV_HEADS = 8
# NOTE: with INTERMEDIATE_SIZE = 4096 the model is ~2.05B params, not 3B.
# Restore 8192 for a true ~2.8-3B model (wider SwiGLU FFN). Your call.
#INTERMEDIATE_SIZE = 8192
INTERMEDIATE_SIZE = 4096
MAX_SEQ_LEN = 4096
RMS_EPS = 1e-6
STABILITY_EPS = 1e-9
INT8_SCALE_TARGET = 127.0
TERNARY = False

class TernaryConfig:
    def __init__(self):
        self.vocab_size = VOCAB_SIZE
        self.hidden_size = HIDDEN_SIZE
        self.num_hidden_layers = NUM_LAYERS
        self.num_attention_heads = NUM_HEADS
        self.num_key_value_heads = NUM_KV_HEADS
        self.intermediate_size = INTERMEDIATE_SIZE
        self.max_position_embeddings = MAX_SEQ_LEN
        self.rms_norm_eps = RMS_EPS
        self.tie_word_embeddings = False
        self.model_type = "jirack_ternary"
        self.ternary = TERNARY  # Флаг теперь внутри конфига

    def get(self, key, default=None):
        return getattr(self, key, default)

    def __getitem__(self, key):
        return getattr(self, key)

class BitLinear(nn.Linear):
    def __init__(self, in_features, out_features, bias=False, ternary=False):
        super().__init__(in_features, out_features, bias)
        self.ternary = ternary

    def forward(self, x):
        if not self.ternary:
            return F.linear(x, self.weight, self.bias)
        # Weight Quantization (ternary {-1,0,+1}, absmean scale) — unchanged
        w = self.weight
        gamma = w.abs().mean().clamp(min=STABILITY_EPS)
        w_quant = torch.clamp(torch.round(w / gamma), -1, 1)
        w_final = w + (w_quant * gamma - w).detach()

        # Activation Quantization (per-token absmax, BitNet b1.58 style)
        # FIX 2: no mean-centering (there is already an RMSNorm before this
        # projection), and symmetric [-127, 127] clamp to match INT8_SCALE_TARGET.
        x_max = x.abs().amax(dim=-1, keepdim=True).clamp(min=STABILITY_EPS)
        scale = INT8_SCALE_TARGET / x_max
        x_quant = (x * scale).round().clamp(-INT8_SCALE_TARGET, INT8_SCALE_TARGET) / scale
        x_final = x + (x_quant - x).detach()

        return F.linear(x_final, w_final, self.bias)

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=RMS_EPS):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        # FIX 1: reduce in float32 then cast back (bf16-safe, prevents spikes)
        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 precompute_freqs_cis(dim, seq_len, theta=500000.0):
    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)
    return (xq * f_cos) + (rotate_half(xq) * f_sin), (xk * f_cos) + (rotate_half(xk) * f_sin)

class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads = config.num_attention_heads
        self.n_kv_heads = config.num_key_value_heads
        self.n_rep = self.n_heads // self.n_kv_heads
        self.head_dim = config.hidden_size // self.n_heads
        # Передаем параметр ternary из конфигурации
        self.q_proj = BitLinear(config.hidden_size, config.hidden_size, ternary=config.ternary)
        self.k_proj = BitLinear(config.hidden_size, self.n_kv_heads * self.head_dim, ternary=config.ternary)
        self.v_proj = BitLinear(config.hidden_size, self.n_kv_heads * self.head_dim, ternary=config.ternary)
        self.out_proj = BitLinear(config.hidden_size, config.hidden_size, ternary=config.ternary)

        self.ffn_w1 = BitLinear(config.hidden_size, config.intermediate_size, ternary=config.ternary)
        self.ffn_w3 = BitLinear(config.hidden_size, config.intermediate_size, ternary=config.ternary)
        self.ffn_w2 = BitLinear(config.intermediate_size, config.hidden_size, ternary=config.ternary)

        self.norm1, self.norm2 = RMSNorm(config.hidden_size), RMSNorm(config.hidden_size)

    def forward(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:
            k = k[:, :, None, :, :].expand(B, self.n_kv_heads, self.n_rep, T, self.head_dim).reshape(B, self.n_heads, T, self.head_dim)
            v = v[:, :, None, :, :].expand(B, self.n_kv_heads, self.n_rep, T, self.head_dim).reshape(B, self.n_heads, T, self.head_dim)

        # Полностью автоматический выбор кернела силами PyTorch
        attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

        x = x + self.out_proj(attn_out.transpose(1, 2).reshape(B, T, D))
        m = self.norm2(x)
        x = x + self.ffn_w2(F.silu(self.ffn_w1(m)) * self.ffn_w3(m))
        return x

class TernaryTransformer3B(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.token_emb = nn.Embedding(config.vocab_size, config.hidden_size)
        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_hidden_layers)])
        self.ln_f = RMSNorm(config.hidden_size)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.head_dim = config.hidden_size // config.num_attention_heads
        self.gradient_checkpointing = False
        self._set_rope_cache(config.max_position_embeddings)
        print(f"Ternary={config.ternary} | Native Auto-SDPA Activated")

    def gradient_checkpointing_enable(self, **kwargs):
        self.gradient_checkpointing = True

    def _set_rope_cache(self, seq_len):
        cos, sin = precompute_freqs_cis(self.head_dim, seq_len)
        self.register_buffer("freqs_cos", cos, persistent=False)
        self.register_buffer("freqs_sin", sin, persistent=False)

    def forward(self, input_ids):
        input_ids = input_ids.to(torch.long)
        T = input_ids.shape[1]
        if T > self.freqs_cos.shape[0]:
            self._set_rope_cache(T)

        x = self.token_emb(input_ids)

        for block in self.blocks:
            if self.gradient_checkpointing and self.training:
                x = checkpoint(block, x, self.freqs_cos, self.freqs_sin, use_reentrant=False)
            else:
                x = block(x, self.freqs_cos, self.freqs_sin)

        logits = self.lm_head(self.ln_f(x))
        return logits, None