Create JiRackNative_3b.py

JiRackNative_3b.py

@@ -0,0 +1,193 @@
+# =============================================================================
+# 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