# ============================================================================= # 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