aurora_model.py

from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from .aurora_config import AuroraConfig

try:
    from cut_cross_entropy import linear_cross_entropy
except ImportError:
    linear_cross_entropy = None


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        scale = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return self.weight * x * scale


def precompute_rope_frequencies(
    seq_len: int, head_dim: int, theta: float, device: torch.device, dtype: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor]:
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    positions = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(positions, inv_freq)
    return freqs.cos().to(dtype=dtype), freqs.sin().to(dtype=dtype)


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    cos = cos[None, :, None, :]
    sin = sin[None, :, None, :]
    x_even = x[..., 0::2]
    x_odd = x[..., 1::2]
    out = torch.empty_like(x)
    out[..., 0::2] = x_even * cos - x_odd * sin
    out[..., 1::2] = x_even * sin + x_odd * cos
    return out


class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: AuroraConfig) -> None:
        super().__init__()
        self.cfg = cfg
        self.num_heads = cfg.num_attention_heads
        self.num_kv_heads = cfg.num_key_value_heads
        self.head_dim = cfg.head_dim
        self.kv_repeat = self.num_heads // self.num_kv_heads

        self.q_proj = nn.Linear(cfg.hidden_size, cfg.num_attention_heads * self.head_dim, bias=cfg.attention_bias)
        self.k_proj = nn.Linear(cfg.hidden_size, cfg.num_key_value_heads * self.head_dim, bias=cfg.attention_bias)
        self.v_proj = nn.Linear(cfg.hidden_size, cfg.num_key_value_heads * self.head_dim, bias=cfg.attention_bias)
        self.o_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=cfg.attention_bias)
        self.q_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps) if cfg.qk_norm else nn.Identity()
        self.k_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps) if cfg.qk_norm else nn.Identity()
        self.dropout_p = cfg.dropout

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        batch, seq_len, _ = x.shape
        q = self.q_proj(x).view(batch, seq_len, self.num_heads, self.head_dim)
        k = self.k_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim)
        v = self.v_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim)

        q = self.q_norm(q)
        k = self.k_norm(k)
        q = apply_rope(q, cos, sin).transpose(1, 2)
        k = apply_rope(k, cos, sin).transpose(1, 2)
        v = v.transpose(1, 2)

        # Explicit K/V expansion is mathematically equivalent to GQA and works
        # across CUDA, Apple MPS, and CPU PyTorch backends.
        if self.kv_repeat > 1:
            k = k.repeat_interleave(self.kv_repeat, dim=1)
            v = v.repeat_interleave(self.kv_repeat, dim=1)
        y = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=None,
            dropout_p=self.dropout_p if self.training else 0.0,
            is_causal=True,
        )
        y = y.transpose(1, 2).contiguous().view(batch, seq_len, self.cfg.hidden_size)
        return self.o_proj(y)


class SwiGLU(nn.Module):
    def __init__(self, cfg: AuroraConfig, intermediate_size: int | None = None) -> None:
        super().__init__()
        intermediate_size = intermediate_size or cfg.intermediate_size
        self.gate_proj = nn.Linear(cfg.hidden_size, intermediate_size, bias=cfg.mlp_bias)
        self.up_proj = nn.Linear(cfg.hidden_size, intermediate_size, bias=cfg.mlp_bias)
        self.down_proj = nn.Linear(intermediate_size, cfg.hidden_size, bias=cfg.mlp_bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class Top1MoE(nn.Module):
    """Top-1 routed SwiGLU experts with Switch-style router regularization."""

    def __init__(self, cfg: AuroraConfig) -> None:
        super().__init__()
        self.num_experts = cfg.num_experts
        self.router_aux_loss_coef = cfg.router_aux_loss_coef
        self.router_z_loss_coef = cfg.router_z_loss_coef
        self.router_noise_scale = cfg.router_noise_scale
        self.capacity_factor = cfg.moe_capacity_factor
        self.use_gate_weight = cfg.router_use_gate_weight
        # Keep routing in BF16/FP32 rather than quantizing its logits to FP8.
        self.router = nn.Linear(cfg.hidden_size, cfg.num_experts, bias=False)
        self.experts = nn.ModuleList([SwiGLU(cfg) for _ in range(cfg.num_experts)])
        # Detached summaries from the latest batch, for collapse detection in
        # the trainer.  They are intentionally not persistent model state.
        self.last_expert_fraction: torch.Tensor | None = None
        self.last_preferred_expert_fraction: torch.Tensor | None = None
        self.last_forced_fraction: torch.Tensor | None = None
        self.last_selected_gate_probability: torch.Tensor | None = None

    def _capacity_constrained_route(
        self, scores: torch.Tensor, preferred_index: torch.Tensor
    ) -> torch.Tensor:
        """Assign exactly one expert/token while bounding every expert load.

        Experts keep their highest-scoring first-choice tokens.  Overflow is
        deterministically retried against each token's next preference.  With
        five experts this small eager-only matching pass is far cheaper than
        an expert MLP and prevents a collapsed router from starving experts.
        """
        token_count = scores.size(0)
        capacity = max(
            math.ceil(token_count / self.num_experts),
            math.ceil(token_count * self.capacity_factor / self.num_experts),
        )
        rankings = torch.argsort(scores, dim=-1, descending=True)
        assigned = torch.full_like(preferred_index, -1)
        remaining = [capacity for _ in range(self.num_experts)]

        for rank in range(self.num_experts):
            for expert_index in range(self.num_experts):
                slots = remaining[expert_index]
                if slots <= 0:
                    continue
                candidates = torch.nonzero(
                    (assigned < 0) & (rankings[:, rank] == expert_index), as_tuple=False
                ).flatten()
                candidate_count = candidates.numel()
                if candidate_count == 0:
                    continue
                if candidate_count > slots:
                    candidate_scores = scores.index_select(0, candidates)[:, expert_index]
                    best_positions = torch.topk(candidate_scores, k=slots, sorted=False).indices
                    candidates = candidates.index_select(0, best_positions)
                    candidate_count = slots
                assigned.index_fill_(0, candidates, expert_index)
                remaining[expert_index] -= candidate_count

        # The combined capacity is at least the token count and every token
        # ranks every expert, so this is a logic invariant rather than an
        # expected fallback.
        if bool(torch.any(assigned < 0)):
            raise RuntimeError("capacity-constrained MoE routing left tokens unassigned")
        return assigned

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        original_shape = x.shape
        flat_x = x.reshape(-1, original_shape[-1])
        router_logits = self.router(flat_x).float()
        router_probs = torch.softmax(router_logits, dim=-1)
        # Routing indices are discrete.  Keep this bookkeeping out of the
        # autograd graph; language gradients still reach the selected gate
        # probability when gate weighting is enabled, while router losses use
        # the clean differentiable probabilities below.
        with torch.no_grad():
            routing_logits = router_logits
            if self.training and self.router_noise_scale > 0:
                # Noisy top-1 routing keeps early training exploratory.  Only
                # the discrete expert choice is noisy; probability weights and
                # regularization remain based on clean BF16/FP32 logits.
                gumbel_noise = -torch.empty_like(router_logits).exponential_().log()
                routing_logits = router_logits + self.router_noise_scale * gumbel_noise
            preferred_index = torch.argmax(routing_logits, dim=-1)
            route_index = (
                self._capacity_constrained_route(routing_logits, preferred_index)
                if self.capacity_factor
                else preferred_index
            )
        selected_router_probability = router_probs.gather(1, route_index.unsqueeze(1)).squeeze(1)
        route_weight = (
            selected_router_probability
            if self.use_gate_weight
            else torch.ones_like(selected_router_probability)
        )

        output = torch.zeros_like(flat_x)
        for expert_index, expert in enumerate(self.experts):
            token_indices = torch.nonzero(route_index == expert_index, as_tuple=False).flatten()
            if token_indices.numel() == 0:
                continue
            expert_input = flat_x.index_select(0, token_indices)
            real_token_count = expert_input.size(0)
            # TorchAO's FP8 GEMMs require their M dimension to be divisible
            # by 16.  Sparse routing gives every expert a variable number of
            # tokens, so pad only this temporary dispatch buffer and discard
            # the corresponding outputs.  This changes no real-token math.
            fp8_padding = (-real_token_count) % 16
            if fp8_padding:
                expert_input = torch.cat(
                    (expert_input, expert_input.new_zeros((fp8_padding, expert_input.size(-1)))), dim=0
                )
            expert_output = expert(expert_input)[:real_token_count]
            routed_output = expert_output * route_weight.index_select(0, token_indices).to(expert_output.dtype).unsqueeze(-1)
            # RMSNorm can promote the residual stream to FP32, while the FP8
            # expert projections return BF16 under autocast.  Restore the
            # residual dtype before scattering selected expert outputs.
            output = output.index_copy(0, token_indices, routed_output.to(output.dtype))

        expert_fraction = F.one_hot(route_index, num_classes=self.num_experts).to(router_probs.dtype).mean(dim=0)
        preferred_fraction = F.one_hot(preferred_index, num_classes=self.num_experts).to(router_probs.dtype).mean(dim=0)
        mean_router_prob = router_probs.mean(dim=0)
        self.last_expert_fraction = expert_fraction.detach()
        self.last_preferred_expert_fraction = preferred_fraction.detach()
        self.last_forced_fraction = (route_index != preferred_index).float().mean().detach()
        self.last_selected_gate_probability = selected_router_probability.mean().detach()
        # Balance the router's *clean preference* rather than the capacity-
        # constrained dispatch, which is intentionally already near-uniform.
        aux_loss = self.router_aux_loss_coef * self.num_experts * torch.sum(
            preferred_fraction * mean_router_prob
        )
        z_loss = self.router_z_loss_coef * torch.logsumexp(router_logits, dim=-1).square().mean()
        return output.reshape(original_shape), aux_loss + z_loss


class DecoderBlock(nn.Module):
    def __init__(self, cfg: AuroraConfig) -> None:
        super().__init__()
        self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.self_attn = CausalSelfAttention(cfg)
        self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.mlp: nn.Module = Top1MoE(cfg) if cfg.num_experts > 1 else SwiGLU(cfg)

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        x = x + self.self_attn(self.input_layernorm(x), cos, sin)
        mlp_input = self.post_attention_layernorm(x)
        if isinstance(self.mlp, Top1MoE):
            mlp_output, router_loss = self.mlp(mlp_input)
        else:
            mlp_output = self.mlp(mlp_input)
            router_loss = x.new_zeros((), dtype=torch.float32)
        return x + mlp_output, router_loss


class AuroraForCausalLM(nn.Module):
    def __init__(self, cfg: AuroraConfig) -> None:
        super().__init__()
        cfg.validate()
        self.cfg = cfg
        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
        self.layers = nn.ModuleList([DecoderBlock(cfg) for _ in range(cfg.num_layers)])
        self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
        self.last_router_loss: torch.Tensor | None = None
        self.register_buffer("rope_cos_cached", torch.empty(0), persistent=False)
        self.register_buffer("rope_sin_cached", torch.empty(0), persistent=False)
        if cfg.tie_word_embeddings:
            self.lm_head.weight = self.embed_tokens.weight
        self.apply(self._init_weights)

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def _rope_cache(
        self, seq_len: int, device: torch.device, dtype: torch.dtype
    ) -> tuple[torch.Tensor, torch.Tensor]:
        cache_miss = (
            self.rope_cos_cached.numel() == 0
            or self.rope_cos_cached.size(0) < seq_len
            or self.rope_cos_cached.device != device
            or self.rope_cos_cached.dtype != dtype
        )
        if cache_miss:
            cos, sin = precompute_rope_frequencies(
                self.cfg.context_length,
                self.cfg.head_dim,
                self.cfg.rope_theta,
                device,
                dtype,
            )
            self.rope_cos_cached = cos
            self.rope_sin_cached = sin
        return self.rope_cos_cached[:seq_len], self.rope_sin_cached[:seq_len]

    def _rope_dtype(self, x: torch.Tensor) -> torch.dtype:
        if x.device.type == "cuda" and torch.is_autocast_enabled("cuda"):
            return torch.get_autocast_dtype("cuda")
        if x.device.type == "cpu" and torch.is_autocast_enabled("cpu"):
            return torch.get_autocast_dtype("cpu")
        return x.dtype

    def forward(
        self, input_ids: torch.Tensor, labels: torch.Tensor | None = None
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        x = self.embed_tokens(input_ids)
        cos, sin = self._rope_cache(x.size(1), x.device, self._rope_dtype(x))
        router_loss = torch.zeros((), device=x.device, dtype=torch.float32)
        for layer in self.layers:
            x, layer_router_loss = layer(x, cos, sin)
            router_loss = router_loss + layer_router_loss
        # Each layer produces a regularizer of the same scale.  Average them
        # so the configured coefficient has the same meaning regardless of
        # depth (instead of becoming 16x stronger in this MoE model).
        router_loss = router_loss / max(1, len(self.layers))
        self.last_router_loss = router_loss.detach()
        x = self.norm(x)
        if labels is not None and linear_cross_entropy is not None:
            # RMSNorm may promote activations to FP32, but Cut Cross Entropy's
            # backward kernel requires BF16/FP16 hidden states.
            loss = linear_cross_entropy(x.to(self.lm_head.weight.dtype), self.lm_head.weight, labels, shift=True)
            logits = x.new_empty(0)
        else:
            logits = self.lm_head(x)
            loss = None
            if labels is not None:
                loss = F.cross_entropy(
                    logits[:, :-1].contiguous().view(-1, logits.size(-1)),
                    labels[:, 1:].contiguous().view(-1),
                )
        # Keep evaluation perplexity comparable to dense models: router regularization
        # shapes gradients only during training and is not language-model loss.
        if loss is not None and self.training:
            loss = loss + router_loss
        return logits, loss


def count_parameters(model: nn.Module) -> int:
    seen: set[int] = set()
    total = 0
    for param in model.parameters():
        ident = id(param)
        if ident not in seen:
            seen.add(ident)
            total += param.numel()
    return total


def count_active_parameters(model: nn.Module) -> int:
    """Count parameters used by one top-1 path, without double-counting ties."""
    seen: set[int] = set()
    total = 0
    for name, param in model.named_parameters():
        if ".mlp.experts." in name:
            expert_index = name.split(".mlp.experts.", 1)[1].split(".", 1)[0]
            if expert_index != "0":
                continue
        ident = id(param)
        if ident not in seen:
            seen.add(ident)
            total += param.numel()
    return total


def estimate_parameter_count(cfg: AuroraConfig) -> int:
    model = AuroraForCausalLM(cfg)
    return count_parameters(model)