| """ |
| ViuMini-Dense-360M — Sovereign 32-Layer Dense Transformer Architecture |
| Canonical implementation for 366.6M parameter Dense SLM. |
| - 32 layers: 100% Dense SwiGLU FFN (No MoE routing overhead) |
| - GQA (15 Q heads, 5 KV heads, head_dim 64, 3:1 ratio) |
| - RMSNorm + QK-Norm + RoPE (theta 500k) + Hybrid SWA/Full Attention (ratio 3) |
| - Tied Embeddings (vocab 48,000, dim 960) -> 366,617,536 params |
| - Native PyTorch SDPA (FlashAttention / memory-efficient kernel) |
| Run: |
| python viu1_dense.py --config ../configs/viu1_dense_config.yaml |
| """ |
| import argparse |
| import math |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| try: |
| import yaml |
| except ImportError: |
| yaml = None |
|
|
|
|
| @dataclass |
| class ViuDenseArgs: |
| dim: int = 960 |
| n_layers: int = 32 |
| n_heads: int = 15 |
| n_kv_heads: int = 5 |
| vocab_size: int = 48000 |
| multiple_of: int = 64 |
| ffn_dim_multiplier: float = 2.7 |
| norm_eps: float = 1e-5 |
| rope_theta: float = 500000.0 |
| max_seq_len: int = 2048 |
| sliding_window: int = 512 |
| hybrid_ratio: int = 3 |
| dropout: float = 0.0 |
| bias: bool = False |
| tie_embeddings: bool = True |
| qk_norm: bool = True |
| multi_token: bool = False |
| |
| moe_every: int = 999 |
| moe_num_experts: int = 0 |
| moe_top_k: int = 0 |
| moe_expert_div: int = 1 |
| moe_shared: bool = False |
| moe_balance_weight: float = 0.0 |
| moe_capacity_factor: float = 1.0 |
| switch_gate: bool = False |
| moe_num_routed_experts: int = None |
| moe_shared_experts: int = None |
|
|
|
|
| |
| ViuArgs = ViuDenseArgs |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-5): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| v = x.pow(2).mean(-1, keepdim=True) |
| return self.weight * (x * torch.rsqrt(v + self.eps)) |
|
|
|
|
| def rope_cos_sin(head_dim: int, seq_len: int, theta: float, device: torch.device, dtype: torch.dtype = torch.float32): |
| assert head_dim % 2 == 0 |
| inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)) |
| t = torch.arange(seq_len, device=device, dtype=torch.float32) |
| f = torch.outer(t, inv) |
| return torch.cos(f).to(dtype), torch.sin(f).to(dtype) |
|
|
|
|
| def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| d = x.shape[-1] |
| x1, x2 = x[..., : d // 2], x[..., d // 2 :] |
| cos = cos[: x.shape[1]].unsqueeze(0).unsqueeze(2).to(dtype=x.dtype) |
| sin = sin[: x.shape[1]].unsqueeze(0).unsqueeze(2).to(dtype=x.dtype) |
| return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, a: ViuDenseArgs, sliding: bool): |
| super().__init__() |
| self.n_heads = a.n_heads |
| self.n_kv = a.n_kv_heads |
| self.head_dim = a.dim // a.n_heads |
| assert a.dim % a.n_heads == 0, f"dim {a.dim} must be divisible by n_heads {a.n_heads}" |
| self.n_rep = a.n_heads // a.n_kv_heads |
| self.scale = 1.0 / math.sqrt(self.head_dim) |
| self.sliding = sliding |
| self.window = a.sliding_window |
| self.qk_norm = a.qk_norm |
| self._mask_cache = None |
|
|
| self.wq = nn.Linear(a.dim, a.n_heads * self.head_dim, bias=a.bias) |
| self.wk = nn.Linear(a.dim, self.n_kv * self.head_dim, bias=a.bias) |
| self.wv = nn.Linear(a.dim, self.n_kv * self.head_dim, bias=a.bias) |
| self.wo = nn.Linear(a.n_heads * self.head_dim, a.dim, bias=a.bias) |
| if a.qk_norm: |
| self.q_norm = RMSNorm(self.head_dim, a.norm_eps) |
| self.k_norm = RMSNorm(self.head_dim, a.norm_eps) |
| self.drop = nn.Dropout(a.dropout) |
|
|
| def _get_mask(self, T: int, device: torch.device): |
| if self._mask_cache is not None and self._mask_cache.shape == (T, T) and self._mask_cache.device == device: |
| return self._mask_cache |
| causal = torch.full((T, T), float("-inf"), device=device) |
| causal = torch.triu(causal, diagonal=1) |
| if self.sliding: |
| win = torch.full((T, T), float("-inf"), device=device) |
| idx = torch.arange(T, device=device) |
| ok = (idx.unsqueeze(1) - idx.unsqueeze(0) < self.window) |
| win = torch.where(ok, torch.zeros((), device=device), win) |
| causal = torch.minimum(causal, win) |
| self._mask_cache = causal |
| return causal |
|
|
| def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| B, T, _ = x.shape |
| q = self.wq(x).view(B, T, self.n_heads, self.head_dim) |
| k = self.wk(x).view(B, T, self.n_kv, self.head_dim) |
| v = self.wv(x).view(B, T, self.n_kv, self.head_dim) |
| if self.qk_norm: |
| q = self.q_norm(q) |
| k = self.k_norm(k) |
| q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) |
| if self.n_rep > 1: |
| k = k.repeat_interleave(self.n_rep, dim=2) |
| v = v.repeat_interleave(self.n_rep, dim=2) |
| q = q.transpose(1, 2) |
| k = k.transpose(1, 2) |
| v = v.transpose(1, 2) |
| drop_p = self.drop.p if self.training else 0.0 |
| if self.sliding: |
| causal = self._get_mask(T, x.device).to(dtype=q.dtype) |
| attn_out = F.scaled_dot_product_attention(q, k, v, attn_mask=causal, dropout_p=drop_p) |
| else: |
| attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=drop_p) |
| attn_out = attn_out.transpose(1, 2).contiguous() |
| out = attn_out.view(B, T, -1) |
| return self.wo(out) |
|
|
|
|
| def ffn_hidden(dim: int, mult: float, multiple_of: int) -> int: |
| h = int(2 * 4 * dim / 3) |
| if mult: |
| h = int(mult * dim) |
| return multiple_of * ((h + multiple_of - 1) // multiple_of) |
|
|
|
|
| class DenseFFN(nn.Module): |
| def __init__(self, a: ViuDenseArgs): |
| super().__init__() |
| h = ffn_hidden(a.dim, a.ffn_dim_multiplier, a.multiple_of) |
| self.w1 = nn.Linear(a.dim, h, bias=a.bias) |
| self.w2 = nn.Linear(h, a.dim, bias=a.bias) |
| self.w3 = nn.Linear(a.dim, h, bias=a.bias) |
| self.drop = nn.Dropout(a.dropout) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.drop(self.w2(F.silu(self.w1(x)) * self.w3(x))) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, a: ViuDenseArgs, idx: int): |
| super().__init__() |
| sliding = ((idx + 1) % a.hybrid_ratio != 0) |
| self.sliding = sliding |
| self.attn_norm = RMSNorm(a.dim, a.norm_eps) |
| self.attn = Attention(a, sliding=sliding) |
| self.ffn_norm = RMSNorm(a.dim, a.norm_eps) |
| self.ffn = DenseFFN(a) |
| self.is_moe = False |
|
|
| def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): |
| x = x + self.attn(self.attn_norm(x), cos, sin) |
| x = x + self.ffn(self.ffn_norm(x)) |
| return x, None, None |
|
|
|
|
| class Viu1Dense(nn.Module): |
| def __init__(self, a: ViuDenseArgs): |
| super().__init__() |
| self.args = a |
| self.tok_emb = nn.Embedding(a.vocab_size, a.dim) |
| self.drop = nn.Dropout(a.dropout) |
| self.layers = nn.ModuleList([Block(a, i) for i in range(a.n_layers)]) |
| self.norm = RMSNorm(a.dim, a.norm_eps) |
| self.head = nn.Linear(a.dim, a.vocab_size, bias=False) |
| self.mt_head = nn.Linear(a.dim, a.vocab_size, bias=False) if a.multi_token else None |
| self.use_checkpointing = False |
| self.apply(self._init) |
| |
| if a.tie_embeddings: |
| self.head.weight = self.tok_emb.weight |
|
|
| def gradient_checkpointing_enable(self): |
| self.use_checkpointing = True |
|
|
| def gradient_checkpointing_disable(self): |
| self.use_checkpointing = False |
|
|
| @staticmethod |
| def _init(m): |
| if isinstance(m, nn.Linear): |
| nn.init.xavier_uniform_(m.weight, gain=0.7) |
| if m.bias is not None: |
| nn.init.zeros_(m.bias) |
| elif isinstance(m, nn.Embedding): |
| nn.init.normal_(m.weight, std=0.02) |
|
|
| def forward(self, ids: torch.Tensor, targets: torch.Tensor = None): |
| B, T = ids.shape |
| assert T <= self.args.max_seq_len, f"Sequence length {T} exceeds max_seq_len {self.args.max_seq_len}" |
| x = self.drop(self.tok_emb(ids)) |
| cos, sin = rope_cos_sin( |
| self.args.dim // self.args.n_heads, T, |
| self.args.rope_theta, ids.device, |
| ) |
| for blk in self.layers: |
| if self.use_checkpointing and self.training: |
| def _create_custom_forward(module): |
| def _custom_forward(*inputs): |
| return module(*inputs) |
| return _custom_forward |
| nx, _, _ = torch.utils.checkpoint.checkpoint( |
| _create_custom_forward(blk), x, cos, sin, use_reentrant=False |
| ) |
| else: |
| nx, _, _ = blk(x, cos, sin) |
| x = nx |
|
|
| h = self.norm(x) |
| loss, mt_loss, logits = None, None, None |
| if targets is not None: |
| flat_h = h.reshape(-1, self.args.dim) |
| flat_tgt = targets.reshape(-1) |
| chunk_size = 1024 if (h.is_cuda and torch.cuda.get_device_properties(h.device).total_memory >= 14 * 1024**3) else 512 |
| n_tokens = flat_tgt.numel() |
| total_loss = 0.0 |
| for ci in range(0, n_tokens, chunk_size): |
| ch_h = flat_h[ci : ci + chunk_size] |
| ch_tgt = flat_tgt[ci : ci + chunk_size] |
| ch_logits = self.head(ch_h) |
| total_loss = total_loss + F.cross_entropy(ch_logits, ch_tgt, reduction="sum") |
| loss = total_loss / max(n_tokens, 1) |
|
|
| if self.mt_head is not None and T > 1: |
| flat_h_mt = h[:, :-1, :].reshape(-1, self.args.dim) |
| flat_tgt_mt = targets[:, 1:].reshape(-1) |
| n_tokens_mt = flat_tgt_mt.numel() |
| total_mt_loss = 0.0 |
| for ci in range(0, n_tokens_mt, chunk_size): |
| ch_h_mt = flat_h_mt[ci : ci + chunk_size] |
| ch_tgt_mt = flat_tgt_mt[ci : ci + chunk_size] |
| ch_logits_mt = self.mt_head(ch_h_mt) |
| total_mt_loss = total_mt_loss + F.cross_entropy(ch_logits_mt, ch_tgt_mt, reduction="sum") |
| mt_loss = total_mt_loss / max(n_tokens_mt, 1) |
| loss = loss + 0.1 * mt_loss |
| else: |
| logits = self.head(h) |
|
|
| aux_info = { |
| "moe_aux": 0.0, |
| "mt_loss": float(mt_loss.detach()) if isinstance(mt_loss, torch.Tensor) else 0.0, |
| "usages": [], |
| } |
| return logits, loss, aux_info |
|
|
| def param_counts(self): |
| total = sum(p.numel() for p in self.parameters()) |
| |
| active = total |
| return total, active |
|
|
|
|
| |
| Viu1MoE = Viu1Dense |
|
|
|
|
| def load_args(path: str) -> ViuDenseArgs: |
| if path and yaml and Path(path).exists(): |
| d = yaml.safe_load(open(path, encoding="utf-8")) or {} |
| return ViuDenseArgs(**{k: v for k, v in d.items() if k in ViuDenseArgs.__dataclass_fields__}) |
| return ViuDenseArgs() |
|
|
|
|
| __all__ = [ |
| "ViuDenseArgs", |
| "Viu1Dense", |
| "ViuArgs", |
| "Viu1MoE", |
| "RMSNorm", |
| "Attention", |
| "DenseFFN", |
| "Block", |
| "rope_cos_sin", |
| "apply_rope", |
| "ffn_hidden", |
| "load_args", |
| ] |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--config", default="../configs/viu1_dense_config.yaml") |
| args = ap.parse_args() |
| cfg = load_args(args.config) |
| m = Viu1Dense(cfg) |
| total, active = m.param_counts() |
| print(f"[config] dim={cfg.dim} L={cfg.n_layers} H={cfg.n_heads} KV={cfg.n_kv_heads} vocab={cfg.vocab_size}") |
| print(f"[params] total={total:,} ({total/1e6:.1f}M) | active={active:,} ({active/1e6:.1f}M)") |
| print(f"[target] ViuMini-Dense-360M (32L Dense, 366.6M params, 100% active)") |
| m.eval() |
| with torch.no_grad(): |
| ids = torch.randint(0, cfg.vocab_size, (1, 16)) |
| logits, loss, aux = m(ids, ids) |
| logits_str = tuple(logits.shape) if logits is not None else "(loss mode, logits purged)" |
| print(f"[smoke] logits {logits_str} loss {loss.item():.3f} aux {aux} OK") |
| n_dense = len(m.layers) |
| n_sliding = sum(1 for b in m.layers if b.sliding) |
| n_full = sum(1 for b in m.layers if not b.sliding) |
| print(f"[layers] dense={n_dense} sliding={n_sliding} full={n_full}") |
|
|
