""" AnuLM -- a small, trainable model with Sarvam 30B / 105B's architecture. Distilled from the annotated reference in sarvam/. Every non-obvious choice here traces to a specific line of Sarvam's released modelling code; the mapping is in README.md and in the comments below. What is faithful: * sparse MoE with a shared expert and a dense first layer * sigmoid router, aux-loss-free bias balancing, group-limited top-k * BOTH attention variants -- GQA (30B) and MLA (105B) -- switchable * QK-norm, pre-norm blocks, RMSNorm, SwiGLU, untied embeddings * YaRN context extension with the band blend and the logit-temperature term What is deliberately different: * it can actually train (the released 105B file cannot -- see README) * RoPE uses the split-halves convention throughout, not the 105B checkpoint's interleaved layout, since we train from scratch * no expert parallelism, no vLLM branch, no FlashAttention subclass: SDPA everywhere, which is what a single-GPU run wants anyway """ from __future__ import annotations import math import dataclasses from dataclasses import dataclass, field from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F # ============================================================================= # CONFIG # ============================================================================= @dataclass class AnuLMConfig: # --- tokens / context ----------------------------------------------- vocab_size: int = 259 # 256 bytes + BOS + EOS + PAD (byte mode) block_size: int = 256 tokenizer_path: Optional[str] = None # set when trained on BPE-encoded # data; sample.py reads it from the # checkpoint to encode/decode # --- trunk ------------------------------------------------------------ n_layer: int = 8 hidden_size: int = 384 n_head: int = 6 # --- attention --------------------------------------------------------- attn: str = "gqa" # "gqa" (30B-style) | "mla" (105B-style) head_dim: Optional[int] = None # gqa: defaults to hidden_size // n_head n_kv_head: int = 2 # gqa only kv_lora_rank: int = 48 # mla only qk_nope_head_dim: int = 48 # mla only qk_rope_head_dim: int = 24 # mla only v_head_dim: int = 48 # mla only use_qk_norm: bool = True # Sliding-window attention on layers [0, max_window_layers) -- Qwen2's # semantics, which the real 30B config inherits: it ships # `max_window_layers: 19` with no window set, i.e. off. None disables. sliding_window: Optional[int] = None max_window_layers: int = 0 # --- MoE --------------------------------------------------------------- num_experts: int = 16 num_experts_per_tok: int = 2 num_shared_experts: int = 1 moe_intermediate_size: int = 96 intermediate_size: int = 768 # the dense layer-0 MLP first_k_dense_replace: int = 1 n_group: int = 1 # group-limited routing; 1 disables it topk_group: int = 1 routed_scaling_factor: float = 2.5 bias_update_rate: float = 1e-3 # aux-loss-free balancing step size # "sparse": gather each expert's tokens and run one GEMM per expert. What a # real MoE does, and 2.1-2.35x faster here in eager -- but it uses # `nonzero`, whose output shape is data-dependent, so torch.compile breaks # the graph 8 times per forward and buys you nothing. # "dense": run every expert on every token and combine with a one-hot # weight matrix. Bit-identical output, 1 graph and 0 breaks under Dynamo, # at the cost of num_experts/top_k times the expert FLOPs. # "grouped": sort the (token, slot) pairs by expert once, then run ALL # experts as one grouped GEMM per projection (F.grouped_mm). No `nonzero`, # so shapes are static and Dynamo captures it in one graph; on a GPU it # turns num_experts small launches into one. bf16 only -- meant for # autocast, where the sparse path's Linears run in bf16 anyway. moe_impl: str = "sparse" # DeepSeek-V3's complementary sequence-wise balance loss, and a router # z-loss. Both default OFF: the bias rule does the balancing and these are # the belt to its braces as data gets more diverse. They enter the training # loss only; train.py reports them separately from the LM loss. seq_balance_alpha: float = 0.0 router_z_alpha: float = 0.0 # --- RoPE -------------------------------------------------------------- rope_theta: float = 10000.0 yarn: bool = False yarn_factor: float = 8.0 yarn_original_context: int = 256 yarn_beta_fast: float = 32.0 yarn_beta_slow: float = 1.0 yarn_mscale: float = 1.0 yarn_mscale_all_dim: float = 1.0 # --- misc -------------------------------------------------------------- rms_norm_eps: float = 1e-6 dropout: float = 0.0 tie_word_embeddings: bool = False # Sarvam unties both models initializer_range: float = 0.02 def __post_init__(self): assert self.attn in ("gqa", "mla") assert self.moe_impl in ("sparse", "dense", "grouped") if self.moe_impl == "grouped": assert hasattr(F, "grouped_mm"), ( "moe_impl='grouped' needs torch.nn.functional.grouped_mm (present in 2.11)") if self.sliding_window is not None: assert self.sliding_window >= 1 and self.max_window_layers >= 1, ( "sliding_window needs max_window_layers >= 1 to apply to any layer") assert self.num_experts % self.n_group == 0, "n_group must divide num_experts" # The footgun that group-limited routing hides: you cannot select more # experts than the surviving groups actually contain. per_group = self.num_experts // self.n_group assert self.num_experts_per_tok <= per_group * self.topk_group, ( f"top_k={self.num_experts_per_tok} exceeds the {per_group * self.topk_group} " f"experts reachable through {self.topk_group} of {self.n_group} groups" ) assert self.topk_group <= self.n_group if self.attn == "gqa": if self.head_dim is None: self.head_dim = self.hidden_size // self.n_head assert self.n_head % self.n_kv_head == 0 else: self.head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim # --- presets: the two real models, scaled down ------------------------ @staticmethod def nano_30b(**kw) -> "AnuLMConfig": """GQA + high rope_theta, no YaRN. Sarvam 30B's shape.""" cfg = dict( attn="gqa", n_layer=8, hidden_size=384, n_head=6, head_dim=64, n_kv_head=2, num_experts=16, num_experts_per_tok=2, moe_intermediate_size=96, # hidden/4, as in the real model intermediate_size=768, # hidden*2 n_group=1, topk_group=1, # grouping is a no-op, like the 30B rope_theta=1e6, yarn=False, ) cfg.update(kw) return AnuLMConfig(**cfg) @staticmethod def nano_350m(**kw) -> "AnuLMConfig": """353M total / 152M active. The largest that trains comfortably on 8 GB. Sized by measurement, not arithmetic. Memory is the binding constraint: optimizer state covers ALL parameters even though only 43% are active per token, so an MoE is memory-hungry to train however sparse its forward pass is. Measured peaks at batch 8 x 512, bf16 autocast, gradient checkpointing on, AdamW: 443M -> 9.74 GB spills to system RAM on an 8 GB card (12x slower) 353M -> 6.34 GB fits, 4955 tok/s <- this preset 281M -> 5.15 GB fits, 5887 tok/s Requires `--grad-ckpt`. Without it the same model peaks at 15.75 GB and thrashes. Muon fits a larger model in the same memory but costs 2.6x per step here -- see muon.py and docs/RESULTS.md. """ cfg = dict( attn="gqa", n_layer=20, hidden_size=1024, n_head=16, head_dim=64, n_kv_head=4, # 4:1, closer to the real 30B's 16:1 num_experts=12, num_experts_per_tok=3, moe_intermediate_size=384, intermediate_size=2816, n_group=1, topk_group=1, rope_theta=1e6, yarn=False, block_size=512, ) cfg.update(kw) return AnuLMConfig(**cfg) @staticmethod def nano_105b(**kw) -> "AnuLMConfig": """MLA + YaRN-ready. Sarvam 105B's shape.""" cfg = dict( attn="mla", n_layer=12, hidden_size=384, n_head=6, kv_lora_rank=48, # hidden/8, as in the real model qk_nope_head_dim=48, qk_rope_head_dim=24, v_head_dim=48, num_experts=16, num_experts_per_tok=2, moe_intermediate_size=192, # hidden/2 intermediate_size=1536, # hidden*4 n_group=4, topk_group=2, # group-limited, like the 105B rope_theta=10000.0, yarn=False, ) cfg.update(kw) return AnuLMConfig(**cfg) # ============================================================================= # NORM / MLP # ============================================================================= class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x): dtype = x.dtype x = x.float() x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return (x.to(dtype)) * self.weight class SwiGLU(nn.Module): """The expert body and the dense layer-0 MLP. 3 * D * I params.""" def __init__(self, hidden_size: int, intermediate_size: int): super().__init__() self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) # ============================================================================= # ROPE (see sarvam/yarn_rotary_annotated.py) # ============================================================================= def _yarn_correction_dim(num_rotations, dim, base, orig_ctx): return (dim * math.log(orig_ctx / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) def _yarn_correction_range(beta_fast, beta_slow, dim, base, orig_ctx): low = math.floor(_yarn_correction_dim(beta_fast, dim, base, orig_ctx)) high = math.ceil(_yarn_correction_dim(beta_slow, dim, base, orig_ctx)) return max(low, 0), min(high, dim - 1) def yarn_mscale(scale: float, mscale: float) -> float: if scale <= 1: return 1.0 return 0.1 * mscale * math.log(scale) + 1.0 class Rotary(nn.Module): """One instance at model level, shared by every layer (the 30B's pattern -- the 105B rebuilds this per layer and pays 2.15 GB for it at load).""" def __init__(self, cfg: AnuLMConfig, dim: int): super().__init__() self.dim = dim self.cfg = cfg inv_freq = self._build_inv_freq(cfg, dim) self.register_buffer("inv_freq", inv_freq, persistent=False) # YaRN's temperature, split the way DeepSeek (and so the 105B) splits # it: attention multiplies its softmax scale by attn_mscale**2, and the # cos/sin tables are scaled by mscale(f, mscale) / mscale(f, mscale_all_dim). # With the shipped 1.0 / 1.0 the table factor is exactly 1.0 and the # whole correction lands on the attention scale. if cfg.yarn: self.attn_mscale = yarn_mscale(cfg.yarn_factor, cfg.yarn_mscale_all_dim) self.table_mscale = (yarn_mscale(cfg.yarn_factor, cfg.yarn_mscale) / yarn_mscale(cfg.yarn_factor, cfg.yarn_mscale_all_dim)) else: self.attn_mscale = self.table_mscale = 1.0 self._cached_len = 0 self._cached_key = None # (device, dtype) the tables were built for self.register_buffer("cos_cached", torch.empty(0), persistent=False) self.register_buffer("sin_cached", torch.empty(0), persistent=False) @staticmethod def _build_inv_freq(cfg: AnuLMConfig, dim: int) -> torch.Tensor: idx = torch.arange(0, dim, 2, dtype=torch.float32) / dim freq_extra = 1.0 / (cfg.rope_theta ** idx) if not cfg.yarn: return freq_extra # NTK-by-parts: keep the fast dims, divide the slow dims by `factor`, # ramp linearly between the two correction dims. freq_inter = 1.0 / (cfg.yarn_factor * (cfg.rope_theta ** idx)) low, high = _yarn_correction_range( cfg.yarn_beta_fast, cfg.yarn_beta_slow, dim, cfg.rope_theta, cfg.yarn_original_context, ) if low == high: high += 0.001 ramp = (torch.arange(dim // 2, dtype=torch.float32) - low) / (high - low) ramp = ramp.clamp(0, 1) keep = 1.0 - ramp # 1 -> extrapolate, 0 -> interpolate return freq_inter * (1 - keep) + freq_extra * keep def _grow(self, seq_len: int, device, dtype): # Keyed on device and dtype as well as length: a model moved with .to() # or run under a different autocast dtype must not be handed stale tables. key = (device, dtype) if seq_len <= self._cached_len and self._cached_key == key: return n = max(seq_len, self._cached_len) t = torch.arange(n, device=device, dtype=torch.float32) freqs = torch.outer(t, self.inv_freq.to(device)) # (T, dim/2) emb = torch.cat((freqs, freqs), dim=-1) # (T, dim) self.cos_cached = (emb.cos() * self.table_mscale).to(dtype) self.sin_cached = (emb.sin() * self.table_mscale).to(dtype) self._cached_len = n self._cached_key = key def forward(self, seq_len: int, device, dtype, start: int = 0): """cos/sin for absolute positions [start, start + seq_len). `start` is the number of tokens already in the KV cache when decoding; 0 otherwise.""" self._grow(start + seq_len, device, dtype) return self.cos_cached[start:start + seq_len], self.sin_cached[start:start + seq_len] def rotate_half(x): x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) def apply_rope(q, k, cos, sin): """q, k: (B, H, S, Dr). cos/sin: (S, Dr) -> broadcast over batch and heads. Split-halves (GPT-NeoX) convention. Sarvam's 105B checkpoint stores its rope dims interleaved and permutes before rotating; we train from scratch, so we pick one convention and stay in it. """ cos = cos[None, None, :, :] sin = sin[None, None, :, :] return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin) def _attn_mask(S: int, T: int, start: int, window: Optional[int], device): """Boolean (S, T) mask, True = may attend. The S queries sit at absolute positions start..start+S-1; the T keys at 0..T-1 (T == start + S). Returns None whenever SDPA's built-in `is_causal` (or no mask at all, for a single decode step) is enough, so the fused kernels keep their fast path. An explicit mask is needed for a multi-token step against a non-empty cache -- `is_causal` aligns top-left, which is wrong there -- and for windows. """ if window is None and (start == 0 or S == 1): return None q_pos = torch.arange(start, start + S, device=device)[:, None] k_pos = torch.arange(T, device=device)[None, :] mask = k_pos <= q_pos if window is not None: mask &= k_pos > q_pos - window return mask def _sdpa(q, k, v, mask, S, scale, dropout_p): return F.scaled_dot_product_attention( q, k, v, attn_mask=mask, is_causal=(mask is None and S > 1), scale=scale, dropout_p=dropout_p, ) # ============================================================================= # ATTENTION -- GQA (see sarvam/gqa_forward_annotated.py) # ============================================================================= class GQAttention(nn.Module): def __init__(self, cfg: AnuLMConfig, rotary: Rotary): super().__init__() self.n_head = cfg.n_head self.n_kv_head = cfg.n_kv_head self.n_rep = cfg.n_head // cfg.n_kv_head self.head_dim = cfg.head_dim self.dropout = cfg.dropout # One fused projection: (H + 2*Hkv) * Dh, exactly as the 30B does. self.query_key_value = nn.Linear( cfg.hidden_size, (cfg.n_head + 2 * cfg.n_kv_head) * cfg.head_dim, bias=False ) self.dense = nn.Linear(cfg.n_head * cfg.head_dim, cfg.hidden_size, bias=False) self.use_qk_norm = cfg.use_qk_norm if cfg.use_qk_norm: self.query_layernorm = RMSNorm(cfg.head_dim, cfg.rms_norm_eps) self.key_layernorm = RMSNorm(cfg.head_dim, cfg.rms_norm_eps) # YaRN's temperature term lives here, not in the rotary tables -- # the same place MLA puts it. Without this, `--yarn` on a GQA model # would apply the band blend and silently drop half of YaRN. self.scale = cfg.head_dim ** -0.5 * (rotary.attn_mscale ** 2) def forward(self, x, cos, sin, cache: Optional[dict] = None, start: int = 0, window: Optional[int] = None): B, S, _ = x.shape qkv = self.query_key_value(x) qkv = qkv.view(B, S, self.n_head + 2 * self.n_kv_head, self.head_dim) q, k, v = qkv.split([self.n_head, self.n_kv_head, self.n_kv_head], dim=-2) q = q.transpose(1, 2) # (B, H, S, Dh) k = k.transpose(1, 2) # (B, Hkv, S, Dh) v = v.transpose(1, 2) # (B, Hkv, S, Dh) if self.use_qk_norm: # before RoPE, as in the 30B q = self.query_layernorm(q) k = self.key_layernorm(k) q, k = apply_rope(q, k, cos, sin) if cache is not None: # Cache the Hkv heads, not their H broadcast copies -- that ratio # (16:1 in the real 30B) is GQA's entire inference saving. if "k" in cache: k = torch.cat([cache["k"], k], dim=2) v = torch.cat([cache["v"], v], dim=2) cache["k"], cache["v"] = k, v T = k.shape[2] # Broadcast each KV head across its query group. (torch>=2.5 could do # this inside SDPA with enable_gqa=True; done explicitly for version safety.) k = k.repeat_interleave(self.n_rep, dim=1) # (B, H, T, Dh) v = v.repeat_interleave(self.n_rep, dim=1) y = _sdpa(q, k, v, _attn_mask(S, T, start, window, x.device), S, self.scale, self.dropout if self.training else 0.0) # (B, H, S, Dh) y = y.transpose(1, 2).contiguous().view(B, S, -1) return self.dense(y) # ============================================================================= # ATTENTION -- MLA (see sarvam/mla_forward_annotated.py) # ============================================================================= class MLAttention(nn.Module): def __init__(self, cfg: AnuLMConfig, rotary: Rotary): super().__init__() self.n_head = cfg.n_head self.qk_nope = cfg.qk_nope_head_dim self.qk_rope = cfg.qk_rope_head_dim self.q_head_dim = cfg.qk_nope_head_dim + cfg.qk_rope_head_dim self.v_head_dim = cfg.v_head_dim self.kv_lora_rank = cfg.kv_lora_rank self.dropout = cfg.dropout self.q_proj = nn.Linear(cfg.hidden_size, cfg.n_head * self.q_head_dim, bias=False) # The one projection that produces both the shared latent and the single # MQA-style positional key. Width = kv_lora_rank + qk_rope_head_dim. self.kv_a_proj_with_mqa = nn.Linear( cfg.hidden_size, cfg.kv_lora_rank + cfg.qk_rope_head_dim, bias=False ) self.kv_a_layernorm = RMSNorm(cfg.kv_lora_rank, cfg.rms_norm_eps) self.kv_b_proj = nn.Linear( cfg.kv_lora_rank, cfg.n_head * (cfg.qk_nope_head_dim + cfg.v_head_dim), bias=False ) self.o_proj = nn.Linear(cfg.n_head * cfg.v_head_dim, cfg.hidden_size, bias=False) self.use_qk_norm = cfg.use_qk_norm if cfg.use_qk_norm: self.q_layernorm = RMSNorm(self.q_head_dim, cfg.rms_norm_eps) # YaRN's temperature term lives here, not in the rotary tables. self.scale = self.q_head_dim ** -0.5 * (rotary.attn_mscale ** 2) def forward(self, x, cos, sin, cache: Optional[dict] = None, start: int = 0, window: Optional[int] = None): B, S, _ = x.shape q = self.q_proj(x).view(B, S, self.n_head, self.q_head_dim).transpose(1, 2) if self.use_qk_norm: q = self.q_layernorm(q) q_nope, q_pe = q.split([self.qk_nope, self.qk_rope], dim=-1) compressed = self.kv_a_proj_with_mqa(x) # (B, S, Lkv+Rk) compressed, k_pe = compressed.split([self.kv_lora_rank, self.qk_rope], dim=-1) k_pe = k_pe.view(B, S, 1, self.qk_rope).transpose(1, 2) # (B, 1, S, Rk) # RoPE touches only the rope slice; the latent stays position-free -- # which is exactly what lets the latent be cached below. q_pe, k_pe = apply_rope(q_pe, k_pe, cos, sin) compressed = self.kv_a_layernorm(compressed) # (B, S, Lkv) if cache is not None: # MLA's cache IS the point of MLA: kv_lora_rank + qk_rope_head_dim # numbers per token (48 + 24 here; 512 + 64 in the 105B) regardless # of head count, against GQA's 2 * Hkv * Dh. K and V are re-expanded # from it through kv_b_proj every step. if "c" in cache: compressed = torch.cat([cache["c"], compressed], dim=1) k_pe = torch.cat([cache["k_pe"], k_pe], dim=2) cache["c"], cache["k_pe"] = compressed, k_pe T = compressed.shape[1] kv = self.kv_b_proj(compressed) kv = kv.view(B, T, self.n_head, self.qk_nope + self.v_head_dim).transpose(1, 2) k_nope, v = kv.split([self.qk_nope, self.v_head_dim], dim=-1) k_pe = k_pe.expand(-1, self.n_head, -1, -1) # (B, H, T, Rk) q = torch.cat([q_nope, q_pe], dim=-1) # (B, H, S, Qd) k = torch.cat([k_nope, k_pe], dim=-1) # (B, H, T, Qd) # SDPA allows V's head dim to differ from Q/K's -- which MLA needs. y = _sdpa(q, k, v, _attn_mask(S, T, start, window, x.device), S, self.scale, self.dropout if self.training else 0.0) # (B, H, S, Vd) y = y.transpose(1, 2).contiguous().view(B, S, self.n_head * self.v_head_dim) return self.o_proj(y) # ============================================================================= # MoE (see sarvam/moe_routing_annotated.py) # ============================================================================= class Router(nn.Module): """Sigmoid scoring + aux-loss-free bias balancing + group-limited top-k. The bias is NOT learned by gradient descent. It is nudged after each optimizer step by `Model.update_expert_biases()`, using the load counts this module accumulates. That is the whole aux-loss-free trick: balancing steers selection without ever entering the loss. """ def __init__(self, cfg: AnuLMConfig): super().__init__() self.top_k = cfg.num_experts_per_tok self.num_experts = cfg.num_experts self.n_group = cfg.n_group self.topk_group = cfg.topk_group self.routed_scaling_factor = cfg.routed_scaling_factor self.weight = nn.Parameter(torch.empty(cfg.num_experts, cfg.hidden_size)) nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) # Buffers, not parameters: no grad, but they must ride along in the # state dict and follow .to(device). self.register_buffer("expert_bias", torch.zeros(cfg.num_experts)) self.register_buffer("load_counts", torch.zeros(cfg.num_experts)) self.seq_balance_alpha = cfg.seq_balance_alpha self.router_z_alpha = cfg.router_z_alpha def forward(self, x_flat, seq_len: Optional[int] = None): # (N, D) """Returns (topk_idx, combine weights, aux) -- aux is None unless one of the regularisers is on and the module is training. `seq_len` lets the sequence-wise balance loss see the (B, S) layout of a flat batch.""" # fp32 router: top-k is discrete, and bf16 rounding would flip choices. # Casting the inputs is not enough -- autocast re-casts F.linear to # bf16 whatever dtype it is handed, and every GPU run before this fix # routed in bf16 (3.2% of expert sets differed from fp32 on a real # checkpoint). Autocast has to be off for the whole gate. with torch.autocast(device_type=x_flat.device.type, enabled=False): topk_idx, w, aux = self._gate(x_flat.float(), seq_len) return topk_idx, w.to(x_flat.dtype), aux def _gate(self, x_flat, seq_len: Optional[int]): logits = F.linear(x_flat, self.weight.float()) # (N, E) scores = torch.sigmoid(logits) # independent, not softmax scores_for_choice = scores + self.expert_bias[None, :] if self.n_group > 1: N = scores.shape[0] per_group = self.num_experts // self.n_group grouped = scores_for_choice.view(N, self.n_group, per_group) # A group is worth its best two experts. g_score = grouped.topk(min(2, per_group), dim=-1)[0].sum(-1) # (N, G) g_idx = g_score.topk(self.topk_group, dim=-1)[1] # (N, Gk) g_mask = torch.zeros_like(g_score).scatter_(1, g_idx, 1.0) mask = g_mask.unsqueeze(-1).expand(N, self.n_group, per_group).reshape(N, -1) scores_for_choice = scores_for_choice.masked_fill(~mask.bool(), float("-inf")) topk_idx = scores_for_choice.topk(self.top_k, dim=-1)[1] # (N, k) # Gather from the UNBIASED scores: the bias picks, it never weights. w = scores.gather(1, topk_idx) # (N, k) if self.top_k > 1: w = w / (w.sum(-1, keepdim=True) + 1e-20) w = w * self.routed_scaling_factor aux = None if self.training: with torch.no_grad(): self.load_counts += torch.bincount( topk_idx.flatten(), minlength=self.num_experts ).to(self.load_counts.dtype) if self.seq_balance_alpha > 0 or self.router_z_alpha > 0: aux = logits.new_zeros(()) if self.seq_balance_alpha > 0: # DeepSeek-V3's complementary loss (their eqs. 17-20), per # sequence: f_i is expert i's share of the sequence's k*T slots # scaled so a uniform router scores 1.0; P_i is its mean # normalised sigmoid score. Only P carries gradient. Tiny alpha # (1e-4 in V3): it is there to stop one *sequence* collapsing # onto an expert, which the batch-level bias rule cannot see. T = seq_len or scores.shape[0] onehot = F.one_hot(topk_idx, self.num_experts).sum(1).float() # (N, E) f = onehot.view(-1, T, self.num_experts).mean(1) * (self.num_experts / self.top_k) P = (scores / scores.sum(-1, keepdim=True)).view(-1, T, self.num_experts).mean(1) aux = aux + self.seq_balance_alpha * (f * P).sum(-1).mean() if self.router_z_alpha > 0: # Router z-loss (ST-MoE): keeps the logits from drifting to # magnitudes where sigmoid saturates and top-k stops moving. aux = aux + self.router_z_alpha * torch.logsumexp(logits, dim=-1).pow(2).mean() return topk_idx, w, aux class MoE(nn.Module): def __init__(self, cfg: AnuLMConfig): super().__init__() self.top_k = cfg.num_experts_per_tok self.num_experts = cfg.num_experts self.impl = cfg.moe_impl self._stacked = None # eval-mode cache for _route_grouped self.router = Router(cfg) self.experts = nn.ModuleList( [SwiGLU(cfg.hidden_size, cfg.moe_intermediate_size) for _ in range(cfg.num_experts)] ) self.shared_expert = ( SwiGLU(cfg.hidden_size, cfg.moe_intermediate_size * cfg.num_shared_experts) if cfg.num_shared_experts > 0 else None ) def forward(self, x): # (B, S, D) """Returns (y, aux): the routed + shared output, and the router's regulariser (None unless enabled and training).""" B, S, D = x.shape x_flat = x.view(-1, D) # (N, D) topk_idx, topk_w, aux = self.router(x_flat, seq_len=S) route = {"dense": self._route_dense, "grouped": self._route_grouped}.get(self.impl, self._route_sparse) y = route(x_flat, topk_idx, topk_w).view(B, S, D) if self.shared_expert is not None: y = y + self.shared_expert(x) # dense, every token, outside routing return y, aux def _route_sparse(self, x_flat, topk_idx, topk_w): """One GEMM per expert over its own tokens. Default.""" N, D = x_flat.shape flat_expert = topk_idx.reshape(-1) # (N*k,) # Which token each (token, slot) assignment belongs to. tok_of_slot = torch.arange(N, device=x_flat.device).repeat_interleave(self.top_k) out = x_flat.new_zeros(N * self.top_k, D) for e in range(self.num_experts): sel = (flat_expert == e).nonzero(as_tuple=True)[0] if sel.numel() == 0: continue # index_select + index_put with distinct indices: differentiable, # so unlike the released 105B this path trains. No repeat_interleave # of the hidden states, so peak memory stays at O(N*k) rows touched # rather than O(N*k) rows materialised up front. out[sel] = self.experts[e](x_flat[tok_of_slot[sel]]).to(out.dtype) return (out.view(N, self.top_k, D) * topk_w.unsqueeze(-1)).sum(1) def _route_dense(self, x_flat, topk_idx, topk_w): """Every expert on every token, combined through a one-hot weight matrix. Static shapes throughout, so Dynamo captures the whole block in one graph -- at num_experts/top_k times the expert FLOPs.""" N, D = x_flat.shape weights = x_flat.new_zeros(N, self.num_experts).scatter_(1, topk_idx, topk_w) y = x_flat.new_zeros(N, D) for e in range(self.num_experts): y = y + weights[:, e: e + 1] * self.experts[e](x_flat) return y def _stacked_weights(self): """The ModuleList's expert weights as three (E, in, out) bf16 tensors. Stacking `.t()` views gives contiguous (E, D, I) directly -- the layout grouped_mm wants for x @ W -- and keeps the checkpoint layout identical to the sparse path. In training that is one copy of the expert weights per step (~1% of the step's bytes). In eval the weights do not change, so the stack is built once and reused: for decoding one token at a time the copy would otherwise cost more than the GEMMs it feeds. """ # Keyed on the parameters' in-place version counters, so a # load_state_dict() after the first forward invalidates the copy too. key = sum(p._version for e in self.experts for p in e.parameters()) if not self.training and self._stacked is not None and self._stacked[0] == key: return self._stacked[1] stacked = tuple( torch.stack([getattr(e, name).weight.t() for e in self.experts]).to(torch.bfloat16) for name in ("gate_proj", "up_proj", "down_proj")) self._stacked = (key, stacked) if not self.training else None return stacked def train(self, mode: bool = True): self._stacked = None # weights are about to move; drop the copy return super().train(mode) def _route_grouped(self, x_flat, topk_idx, topk_w): """Every expert in ONE grouped GEMM per projection (F.grouped_mm). Sort the N*k (token, slot) assignments by expert; the per-expert row ranges are then contiguous and `offs` (a cumsum of the bincount) tells the kernel where each group ends. argsort, bincount and cumsum all have static output shapes, so unlike `nonzero` this traces in one graph. bf16 only, by the kernel's contract -- run it under autocast. """ N, D = x_flat.shape E, k = self.num_experts, self.top_k flat_expert = topk_idx.reshape(-1) # (N*k,) order = torch.argsort(flat_expert, stable=True) # slots, grouped by expert # Per-expert counts via scatter_add rather than bincount: bincount's # output length depends on the data, which is itself a graph break. counts = torch.zeros(E, dtype=torch.int64, device=x_flat.device).scatter_add_( 0, flat_expert, torch.ones_like(flat_expert)) offs = counts.cumsum(0).to(torch.int32) xs = x_flat[order // k].to(torch.bfloat16) # (N*k, D) expert-sorted w_gate, w_up, w_down = self._stacked_weights() h = F.silu(F.grouped_mm(xs, w_gate, offs=offs)) * F.grouped_mm(xs, w_up, offs=offs) ys = F.grouped_mm(h, w_down, offs=offs) # (N*k, D) expert-sorted out = x_flat.new_zeros(N * k, D) out[order] = ys.to(out.dtype) # back to slot order return (out.view(N, k, D) * topk_w.unsqueeze(-1)).sum(1) # ============================================================================= # BLOCK / MODEL # ============================================================================= class Block(nn.Module): def __init__(self, cfg: AnuLMConfig, layer_idx: int, rotary: Rotary): super().__init__() self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) self.attn = ( GQAttention(cfg, rotary) if cfg.attn == "gqa" else MLAttention(cfg, rotary) ) self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) # Layer 0 dense, the rest MoE -- first_k_dense_replace = 1 in both models. self.is_moe = layer_idx >= cfg.first_k_dense_replace self.mlp = MoE(cfg) if self.is_moe else SwiGLU(cfg.hidden_size, cfg.intermediate_size) # Sliding window on the first max_window_layers layers, full attention # on the rest (Qwen2 semantics). None = full attention. self.window = cfg.sliding_window if layer_idx < cfg.max_window_layers else None def forward(self, x, cos, sin, cache: Optional[dict] = None, start: int = 0): """Returns (x, aux). aux is the MoE regulariser or None.""" x = x + self.attn(self.input_layernorm(x), cos, sin, cache=cache, start=start, window=self.window) h = self.mlp(self.post_attention_layernorm(x)) h, aux = h if self.is_moe else (h, None) return x + h, aux class AnuLM(nn.Module): def __init__(self, cfg: AnuLMConfig): super().__init__() self.cfg = cfg self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size) rope_dim = cfg.qk_rope_head_dim if cfg.attn == "mla" else cfg.head_dim self.rotary = Rotary(cfg, rope_dim) self.layers = nn.ModuleList([Block(cfg, i, self.rotary) for i in range(cfg.n_layer)]) self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False) if cfg.tie_word_embeddings: self.lm_head.weight = self.embed_tokens.weight self.grad_ckpt = False # set via enable_gradient_checkpointing() self.apply(self._init_weights) # nanoGPT's scaled init for the residual output projections. for name, p in self.named_parameters(): if name.endswith(("o_proj.weight", "dense.weight", "down_proj.weight")): nn.init.normal_(p, mean=0.0, std=cfg.initializer_range / math.sqrt(2 * cfg.n_layer)) def _init_weights(self, m): if isinstance(m, nn.Linear): nn.init.normal_(m.weight, mean=0.0, std=self.cfg.initializer_range) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, mean=0.0, std=self.cfg.initializer_range) def new_cache(self) -> list[dict]: """One dict per layer; each attention module owns its own entries.""" return [dict() for _ in self.layers] def forward(self, idx, targets=None, cache: Optional[list] = None, start: int = 0): """`cache` (from `new_cache()`) holds K/V for the `start` tokens already seen; this call attends to them and appends the new ones. Positions are absolute -- start..start+S-1 -- so RoPE agrees with the cached keys. With targets, the returned loss is LM cross-entropy plus any router regulariser; the regulariser alone is kept in `self.aux_loss` (detached) so the training log can report the two separately. """ B, S = idx.shape assert start + S <= self.cfg.block_size, \ f"positions up to {start + S} exceed block_size {self.cfg.block_size}" x = self.embed_tokens(idx) # (B, S, D) cos, sin = self.rotary(S, x.device, x.dtype, start=start) aux_total = None for i, layer in enumerate(self.layers): if self.grad_ckpt and self.training: # Recompute activations in the backward pass instead of storing # them: roughly 30% slower per step, but frees the activation # memory that otherwise caps batch size on a small card. x, aux = torch.utils.checkpoint.checkpoint( layer, x, cos, sin, use_reentrant=False) else: x, aux = layer(x, cos, sin, cache[i] if cache is not None else None, start) if aux is not None: aux_total = aux if aux_total is None else aux_total + aux x = self.norm(x) self.aux_loss = aux_total.detach() if aux_total is not None else None if targets is None: # Inference shortcut: only the last position matters. logits = self.lm_head(x[:, -1:, :]) return logits, None logits = self.lm_head(x) loss = F.cross_entropy( logits.view(-1, logits.size(-1)).float(), targets.reshape(-1), ignore_index=-1 ) if aux_total is not None: loss = loss + aux_total return logits, loss def enable_gradient_checkpointing(self, enable: bool = True): """Unlike the released Sarvam 105B file -- which advertises support and then never calls the checkpoint function -- this actually takes effect. See sarvam/decoder_and_model_annotated.py note 3.""" self.grad_ckpt = enable # --- aux-loss-free load balancing ------------------------------------- @torch.no_grad() def update_expert_biases(self) -> float: """Call once per optimizer step. Returns the load-imbalance ratio (max expert load / mean load; 1.0 is perfect, top_k*... is collapse). DeepSeek-V3's rule: nudge each expert's bias by a fixed step against the sign of its load error. Overloaded experts get harder to pick, starved ones easier -- and because the bias never touches the combine weights, none of this perturbs the language-modelling gradient. """ ratios = [] for m in self.modules(): if not isinstance(m, Router): continue load = m.load_counts total = load.sum() if total <= 0: continue mean = load.mean() m.expert_bias += self.cfg.bias_update_rate * torch.sign(mean - load) ratios.append((load.max() / mean).item()) load.zero_() return max(ratios) if ratios else 1.0 # --- bookkeeping -------------------------------------------------------- def num_params(self) -> tuple[int, int]: """(total, active-per-token). Mirrors the count in the annotated files.""" total = sum(p.numel() for p in self.parameters()) cfg = self.cfg inactive = 0 for layer in self.layers: if isinstance(layer.mlp, MoE): per_expert = sum(p.numel() for p in layer.mlp.experts[0].parameters()) inactive += per_expert * (cfg.num_experts - cfg.num_experts_per_tok) return total, total - inactive @torch.no_grad() def generate(self, idx, max_new_tokens: int, temperature: float = 1.0, top_k: int | None = None, eos_id: int | None = None, use_cache: bool = True, repetition_penalty: float = 1.0): """Incremental decoding. The prompt is prefilled once; every later step runs a single token against the KV cache. Once the context is full (block_size tokens) the cache is rebuilt from the last block_size tokens each step, which is exactly what the uncached path computes -- so the two paths agree token-for-token, and the tests hold them to it. Stops early when every row has produced `eos_id`; finished rows keep receiving eos so the batch stays rectangular. `repetition_penalty` > 1 (CTRL's rule: divide a positive logit, multiply a negative one, for every token already in the sequence) discourages the loops a small model falls into at low temperature. 1.0 is off, and the default, so nothing above changes. """ B = idx.shape[0] bs = self.cfg.block_size done = torch.zeros(B, dtype=torch.bool, device=idx.device) cache, start, pending = None, 0, idx for _ in range(max_new_tokens): if not use_cache: logits, _ = self(idx[:, -bs:]) else: if cache is None or start + pending.shape[1] > bs: cache, start, pending = self.new_cache(), 0, idx[:, -bs:] logits, _ = self(pending, cache=cache, start=start) start += pending.shape[1] logits = logits[:, -1, :].float() if repetition_penalty != 1.0: seen = logits.gather(1, idx) seen = torch.where(seen > 0, seen / repetition_penalty, seen * repetition_penalty) logits = logits.scatter(1, idx, seen) logits = logits / max(temperature, 1e-5) if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float("inf") probs = F.softmax(logits, dim=-1) nxt = torch.multinomial(probs, num_samples=1) # (B, 1) if eos_id is not None: nxt = torch.where(done[:, None], torch.full_like(nxt, eos_id), nxt) idx = torch.cat((idx, nxt), dim=1) pending = nxt if eos_id is not None: done |= nxt[:, 0] == eos_id if bool(done.all()): break return idx if __name__ == "__main__": for name, cfg in [("nano_30b (GQA)", AnuLMConfig.nano_30b()), ("nano_105b (MLA)", AnuLMConfig.nano_105b())]: # Seed per preset, not once up front: every number below is then # reproducible and quotable in the README, and stays that way if the # list is reordered or another preset is added. It matters most for the # "params without grad" count, which is a property of which experts this # particular 128-token batch happened to miss. torch.manual_seed(0) model = AnuLM(cfg) total, active = model.num_params() print(f"{name:18s} total {total/1e6:6.2f}M active/token {active/1e6:6.2f}M " f"({100*active/total:.1f}%) layers {cfg.n_layer}") x = torch.randint(0, cfg.vocab_size, (2, 64)) model.train() logits, loss = model(x, x) loss.backward() # Smoke test: an untrained model should sit at ln(V), every parameter # should have received a gradient, and the bias update should run. no_grad = [n for n, p in model.named_parameters() if p.grad is None] imbalance = model.update_expert_biases() print(f"{'':18s} logits {tuple(logits.shape)} loss {loss.item():.4f} " f"(ln(V) = {math.log(cfg.vocab_size):.4f})") print(f"{'':18s} backward OK params without grad: {len(no_grad)}" f"{' -> ' + ', '.join(no_grad[:3]) if no_grad else ''}") print(f"{'':18s} expert load imbalance after 1 step: {imbalance:.2f}x\n") # The project was called nanosarvam until 2026-09-18. Checkpoints saved before # then pickle their config as ``model.NanoSarvamConfig``; keep both names # resolvable so torch.load keeps working on them. NanoSarvamConfig = AnuLMConfig NanoSarvam = AnuLM def load_checkpoint(path, map_location="cpu") -> dict: """Load a training checkpoint (``.pt``) or a folder written by ``export_hf.py`` (``config.json`` + ``model.safetensors``) into the same dict shape: ``model`` (state dict), ``cfg`` (AnuLMConfig), and whatever ``step``, ``val_loss``, ``qa_templates`` and ``base_ckpt`` were saved.""" import json from pathlib import Path as _P p = _P(path) if not p.is_dir(): return torch.load(str(p), map_location=map_location, weights_only=False) from safetensors.torch import load_file meta = json.loads((p / "config.json").read_text(encoding="utf-8")) cfg_d = dict(meta["config"]) if cfg_d.get("tokenizer_path"): cfg_d["tokenizer_path"] = str(p / cfg_d["tokenizer_path"]) known = {f.name for f in dataclasses.fields(AnuLMConfig)} cfg = AnuLMConfig(**{k: v for k, v in cfg_d.items() if k in known}) dev = map_location if isinstance(map_location, str) else str(map_location) sd = load_file(str(p / "model.safetensors"), device=dev) ck = {"model": sd, "cfg": cfg} for k in ("step", "val_loss", "qa_templates", "base_ckpt"): if meta.get(k) is not None: ck[k] = meta[k] return ck