Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,279 @@
+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)
+        # Optimized: removed redundant self.weight.float() cast for bfloat16/float16 throughput
+        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]
+    # x is [B, n_heads, T, head_dim]
+    # cos, sin are [max_len, head_dim//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  # flag for depth-scaled init
+
+        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  # flag for depth-scaled init
+
+    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)
+        
+        # NEFTune: Noisy Embedding Fine-Tuning for conversational quality boost
+        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)
+
+        # BUG 2 FIX: Additive Float Attention Mask (0.0 = attend, -inf = mask out)
+        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)     # [B, 1, 1, T]
+            bool_mask = causal.unsqueeze(0).unsqueeze(0) & key_mask   # [B, 1, T, T] True = allowed
+            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)
+        logits = self.head(x)
+
+        loss = None
+        if targets is not None:
+            # BUG 1 FIX: Explicit shift_labels flag support
+            if shift_labels:
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_targets = targets[..., 1:].contiguous()
+                loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_targets.view(-1), ignore_index=-100)
+            else:
+                # Targets are already pre-shifted (e.g. pretraining ShardPool y = tokens[i+1:])
+                loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100)
+
+            # z-loss: stabilizes logit magnitudes during pretraining; disable for SFT (z_loss_weight=0)
+            if self.cfg.z_loss_weight > 0:
+                z_loss = self.cfg.z_loss_weight * (torch.logsumexp(logits, dim=-1) ** 2).mean()
+                loss = loss + z_loss
+
+        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):
+            # Crop to context window
+            idx_cond = idx if idx.size(1) <= self.cfg.context_length else idx[:, -self.cfg.context_length:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :]  # only last position
+
+            # Repetition penalty (Standard Hugging Face formulation)
+            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
+
+            # Temperature scaling
+            logits = logits / max(temperature, 1e-8)
+            if not torch.isfinite(logits).any():
+                raise ValueError("Model logits became non-finite (NaN/Inf) during generation. Ensure stable precision or checkpoint validity.")
+
+            # Top-k filtering
+            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')
+
+            # Top-p (nucleus) filtering
+            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
+
+# Aliases for universal compatibility
+Transformer = ViuAI
+ModelArgs = ViuAIConfig