| import math
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| from torch.utils.checkpoint import checkpoint
|
| from config import ViuAIConfig
|
|
|
| class RMSNorm(nn.Module):
|
| def __init__(self, dim, eps=1e-5):
|
| super().__init__()
|
| self.eps = eps
|
| self.weight = nn.Parameter(torch.ones(dim))
|
|
|
| def forward(self, x):
|
| x_fp32 = x.float()
|
| norm = x_fp32 * torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + self.eps)
|
|
|
| return (norm * self.weight).type_as(x)
|
|
|
| def precompute_rope(head_dim, max_len, theta=10000.0, device="cpu"):
|
| freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
|
| t = torch.arange(max_len, device=device).float()
|
| freqs = torch.outer(t, freqs)
|
| return torch.cos(freqs), torch.sin(freqs)
|
|
|
| def apply_rope(x, cos, sin):
|
| cos = cos.to(x.dtype)
|
| sin = sin.to(x.dtype)
|
| x1, x2 = x[..., ::2], x[..., 1::2]
|
|
|
|
|
| T = x.size(2)
|
| cos = cos[:T][None, None, :, :]
|
| sin = sin[:T][None, None, :, :]
|
| rotated = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
|
| return rotated.flatten(-2)
|
|
|
| def repeat_kv(x, n_rep):
|
| if n_rep == 1:
|
| return x
|
| return x.repeat_interleave(n_rep, dim=1)
|
|
|
| class Attention(nn.Module):
|
| def __init__(self, cfg: ViuAIConfig):
|
| super().__init__()
|
| self.n_heads = cfg.n_heads
|
| self.n_kv_heads = cfg.n_kv_heads
|
| self.head_dim = cfg.d_model // cfg.n_heads
|
| self.n_rep = self.n_heads // self.n_kv_heads
|
|
|
| self.q_proj = nn.Linear(cfg.d_model, cfg.n_heads * self.head_dim, bias=False)
|
| self.k_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
|
| self.v_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
|
| self.o_proj = nn.Linear(cfg.n_heads * self.head_dim, cfg.d_model, bias=False)
|
| self.o_proj._is_residual = True
|
|
|
| self.q_norm = RMSNorm(self.head_dim, cfg.norm_eps)
|
| self.k_norm = RMSNorm(self.head_dim, cfg.norm_eps)
|
| self.attn_dropout = cfg.attn_dropout
|
|
|
| def forward(self, x, cos, sin, attn_mask=None):
|
| B, T, C = x.shape
|
| q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
|
| k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
|
| v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
|
|
|
| q, k = self.q_norm(q), self.k_norm(k)
|
| q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
|
|
|
| dropout_p = self.attn_dropout if self.training else 0.0
|
| try:
|
| if attn_mask is not None:
|
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=dropout_p, enable_gqa=True)
|
| else:
|
| out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=dropout_p, enable_gqa=True)
|
| except TypeError:
|
| if self.n_rep > 1:
|
| k = repeat_kv(k, self.n_rep)
|
| v = repeat_kv(v, self.n_rep)
|
| if attn_mask is not None:
|
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=dropout_p)
|
| else:
|
| out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=dropout_p)
|
|
|
| out = out.transpose(1, 2).contiguous().view(B, T, -1)
|
| return self.o_proj(out)
|
|
|
| class SwiGLU(nn.Module):
|
| def __init__(self, cfg: ViuAIConfig):
|
| super().__init__()
|
| self.gate_proj = nn.Linear(cfg.d_model, cfg.ffn_hidden, bias=False)
|
| self.up_proj = nn.Linear(cfg.d_model, cfg.ffn_hidden, bias=False)
|
| self.down_proj = nn.Linear(cfg.ffn_hidden, cfg.d_model, bias=False)
|
| self.down_proj._is_residual = True
|
|
|
| def forward(self, x):
|
| return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|
|
|
| class Block(nn.Module):
|
| def __init__(self, cfg: ViuAIConfig):
|
| super().__init__()
|
| self.attn_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
|
| self.attn = Attention(cfg)
|
| self.ffn_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
|
| self.ffn = SwiGLU(cfg)
|
| self.resid_dropout = nn.Dropout(cfg.resid_dropout) if cfg.resid_dropout > 0 else nn.Identity()
|
|
|
| def forward(self, x, cos, sin, attn_mask=None):
|
| x = x + self.resid_dropout(self.attn(self.attn_norm(x), cos, sin, attn_mask))
|
| x = x + self.resid_dropout(self.ffn(self.ffn_norm(x)))
|
| return x
|
|
|
| class ViuAI(nn.Module):
|
| def __init__(self, cfg: ViuAIConfig):
|
| super().__init__()
|
| assert cfg.d_model % cfg.n_heads == 0, "d_model must be divisible by n_heads"
|
| assert cfg.n_heads % cfg.n_kv_heads == 0, "n_heads must be divisible by n_kv_heads"
|
| assert (cfg.d_model // cfg.n_heads) % 2 == 0, "head_dim must be even for RoPE"
|
|
|
| self.cfg = cfg
|
| self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
|
| self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
|
| self.final_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
|
| self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
|
| self.head.weight = self.tok_emb.weight
|
|
|
| head_dim = cfg.d_model // cfg.n_heads
|
| cos, sin = precompute_rope(head_dim, cfg.context_length, cfg.rope_theta)
|
| self.register_buffer("rope_cos", cos, persistent=False)
|
| self.register_buffer("rope_sin", sin, persistent=False)
|
|
|
| self.apply(self._init_weights)
|
|
|
| def _init_weights(self, module):
|
| if isinstance(module, nn.Linear):
|
| if hasattr(self, 'head') and module is self.head:
|
| return
|
| std = 0.02
|
| if getattr(module, '_is_residual', False):
|
| std *= (2 * self.cfg.n_layers) ** -0.5
|
| nn.init.normal_(module.weight, mean=0.0, std=std)
|
| elif isinstance(module, nn.Embedding):
|
| nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
|
|
| def forward(self, idx, targets=None, pad_id=None, shift_labels=False):
|
| B, T = idx.shape
|
| assert T <= self.cfg.context_length, f"Sequence length {T} exceeds context_length {self.cfg.context_length}"
|
| x = self.tok_emb(idx)
|
|
|
|
|
| if self.training and getattr(self.cfg, "neftune_alpha", 0.0) > 0.0:
|
| dims = x.size(1) * x.size(2)
|
| mag = self.cfg.neftune_alpha / math.sqrt(max(1, dims))
|
| noise = torch.zeros_like(x).uniform_(-1.0, 1.0) * mag
|
| x = x + noise
|
|
|
| cos = self.rope_cos.to(x.device)
|
| sin = self.rope_sin.to(x.device)
|
|
|
|
|
| attn_mask = None
|
| if pad_id is not None:
|
| causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device))
|
| key_mask = (idx != pad_id).unsqueeze(1).unsqueeze(2)
|
| bool_mask = causal.unsqueeze(0).unsqueeze(0) & key_mask
|
| attn_mask = torch.zeros(B, 1, T, T, dtype=x.dtype, device=x.device)
|
| attn_mask = attn_mask.masked_fill(~bool_mask, float("-inf"))
|
|
|
| for block in self.blocks:
|
| if self.training and self.cfg.use_checkpoint:
|
| if attn_mask is not None:
|
| x = checkpoint(block, x, cos, sin, attn_mask, use_reentrant=False)
|
| else:
|
| x = checkpoint(block, x, cos, sin, use_reentrant=False)
|
| else:
|
| x = block(x, cos, sin, attn_mask)
|
| x = self.final_norm(x)
|
|
|
|
|
|
|
|
|
|
|
|
|
| loss = None
|
| if targets is not None:
|
| V = self.cfg.vocab_size
|
| ce_chunk = 512
|
|
|
| if shift_labels:
|
|
|
| hidden = x[:, :-1, :].contiguous()
|
| tgt = targets[:, 1:].contiguous()
|
| else:
|
| hidden = x
|
| tgt = targets
|
|
|
| total_loss = torch.zeros(1, device=x.device, dtype=torch.float32)
|
| total_active = (tgt != -100).sum().clamp(min=1)
|
| seq_len = hidden.size(1)
|
|
|
| for i in range(0, seq_len, ce_chunk):
|
| end = min(i + ce_chunk, seq_len)
|
| chunk_logits = self.head(hidden[:, i:end, :])
|
| chunk_tgt = tgt[:, i:end]
|
| chunk_loss = F.cross_entropy(
|
| chunk_logits.reshape(-1, V), chunk_tgt.reshape(-1),
|
| ignore_index=-100, reduction="sum"
|
| )
|
| total_loss = total_loss + chunk_loss
|
|
|
| loss = total_loss / total_active
|
|
|
|
|
|
|
| if self.cfg.z_loss_weight > 0:
|
| z_sum = torch.zeros(1, device=x.device, dtype=torch.float32)
|
| z_count = 0
|
| x_flat_len = x.size(1)
|
| for i in range(0, x_flat_len, ce_chunk):
|
| end = min(i + ce_chunk, x_flat_len)
|
| chunk_logits = self.head(x[:, i:end, :])
|
| z_sum = z_sum + (torch.logsumexp(chunk_logits.float(), dim=-1) ** 2).sum()
|
| z_count += chunk_logits.size(0) * chunk_logits.size(1)
|
| z_loss = self.cfg.z_loss_weight * (z_sum / max(1, z_count))
|
| loss = loss + z_loss
|
|
|
|
|
|
|
| return None, loss
|
|
|
|
|
| logits = self.head(x)
|
| return logits, loss
|
|
|
| def num_params(self):
|
| """Count parameters, deduplicating tied weights."""
|
| seen = set()
|
| total = 0
|
| for p in self.parameters():
|
| if p.data_ptr() not in seen:
|
| seen.add(p.data_ptr())
|
| total += p.numel()
|
| return total
|
|
|
| @torch.no_grad()
|
| def generate(self, idx, max_new_tokens, temperature=1.0, top_k=50, top_p=0.9,
|
| eos_token_id=None, repetition_penalty=1.0):
|
| self.eval()
|
| idx = idx.to(next(self.parameters()).device)
|
| def _is_all_eos(next_tokens):
|
| if eos_token_id is None:
|
| return False
|
| if isinstance(eos_token_id, (list, tuple, set)):
|
| is_eos = torch.zeros(next_tokens.shape[0], dtype=torch.bool, device=next_tokens.device)
|
| for eid in eos_token_id:
|
| is_eos |= (next_tokens.squeeze(-1) == eid)
|
| return is_eos.all().item()
|
| return (next_tokens == eos_token_id).all().item()
|
|
|
| for _ in range(max_new_tokens):
|
|
|
| idx_cond = idx if idx.size(1) <= self.cfg.context_length else idx[:, -self.cfg.context_length:]
|
| logits, _ = self(idx_cond)
|
| logits = logits[:, -1, :]
|
|
|
|
|
| if repetition_penalty != 1.0:
|
| for b in range(idx.size(0)):
|
| prev_tokens = idx[b].unique()
|
| score = logits[b, prev_tokens]
|
| logits[b, prev_tokens] = torch.where(
|
| score < 0, score * repetition_penalty, score / repetition_penalty
|
| )
|
|
|
| if temperature <= 0:
|
| idx_next = logits.argmax(dim=-1, keepdim=True)
|
| idx = torch.cat([idx, idx_next], dim=1)
|
| if _is_all_eos(idx_next):
|
| break
|
| continue
|
|
|
|
|
| logits = logits / max(temperature, 1e-8)
|
| if not torch.isfinite(logits).all():
|
| raise ValueError("Model logits became non-finite (NaN/Inf) during generation. Ensure stable precision or checkpoint validity.")
|
|
|
|
|
| if top_k > 0:
|
| top_k_val = min(top_k, logits.size(-1))
|
| kth_vals, _ = torch.topk(logits, top_k_val)
|
| logits[logits < kth_vals[:, [-1]]] = float('-inf')
|
|
|
|
|
| if top_p < 1.0:
|
| sorted_logits, sorted_indices = torch.sort(logits, descending=True)
|
| cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
|
| sorted_indices_to_remove = cumulative_probs > top_p
|
| sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
|
| sorted_indices_to_remove[..., 0] = False
|
| indices_to_remove = torch.zeros_like(logits, dtype=torch.bool).scatter_(
|
| 1, sorted_indices, sorted_indices_to_remove
|
| )
|
| logits[indices_to_remove] = float('-inf')
|
|
|
| probs = F.softmax(logits, dim=-1)
|
| idx_next = torch.multinomial(probs, num_samples=1)
|
| idx = torch.cat([idx, idx_next], dim=1)
|
|
|
| if _is_all_eos(idx_next):
|
| break
|
|
|
| return idx
|
|
|
|
|
| Transformer = ViuAI
|
| ModelArgs = ViuAIConfig
|
|
|
