Fix normal Transformers inference loading

README.md

@@ -49,7 +49,7 @@ tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
 messages = [{"role": "user", "content": "What is Python?"}]
 inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
-outputs = model.generate(inputs, max_new_tokens=96, do_sample=False)
+outputs = model.generate(**inputs, max_new_tokens=96, do_sample=False, use_cache=False)
 print(tokenizer.decode(outputs[0], skip_special_tokens=True))
 ```
 

aurora_config.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, fields
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+@dataclass
+class AuroraConfig:
+    model_name: str = "Ember Proelia"
+    vocab_size: int = 16000
+    hidden_size: int = 896
+    num_layers: int = 23
+    num_attention_heads: int = 14
+    num_key_value_heads: int = 2
+    intermediate_size: int = 2432
+    context_length: int = 2048
+    rope_theta: float = 500000.0
+    rms_norm_eps: float = 1.0e-5
+    qk_norm: bool = True
+    tie_word_embeddings: bool = True
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    dropout: float = 0.0
+    num_experts: int = 1
+    router_aux_loss_coef: float = 0.0
+    router_z_loss_coef: float = 0.0
+    router_noise_scale: float = 0.0
+    moe_capacity_factor: float = 0.0
+    router_use_gate_weight: bool = False
+
+    @property
+    def head_dim(self) -> int:
+        return self.hidden_size // self.num_attention_heads
+
+    def validate(self) -> None:
+        if self.vocab_size <= 0 or self.hidden_size <= 0 or self.num_layers <= 0:
+            raise ValueError("vocab_size, hidden_size, and num_layers must be positive")
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError("hidden_size must divide evenly by num_attention_heads")
+        if self.num_attention_heads % self.num_key_value_heads != 0:
+            raise ValueError("num_attention_heads must divide evenly by num_key_value_heads")
+        if self.head_dim % 2 != 0:
+            raise ValueError("head_dim must be even for RoPE")
+        if self.context_length <= 0:
+            raise ValueError("context_length must be positive")
+        if self.num_experts <= 0:
+            raise ValueError("num_experts must be positive")
+
+
+def load_model_config(path: str | Path) -> AuroraConfig:
+    path = Path(path)
+    with path.open("r", encoding="utf-8") as handle:
+        raw: dict[str, Any] = yaml.safe_load(handle) or {}
+    allowed = {field.name for field in fields(AuroraConfig)}
+    unknown = sorted(set(raw) - allowed)
+    if unknown:
+        raise ValueError(f"Unknown model config keys: {unknown}")
+    cfg = AuroraConfig(**raw)
+    cfg.validate()
+    return cfg

aurora_model.py

@@ -0,0 +1,376 @@
+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)

modeling_aurora.py

@@ -8,8 +8,8 @@ from transformers import PreTrainedModel
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
 from .configuration_aurora import AuroraHFConfig
-from .aurora.config import AuroraConfig as NativeAuroraConfig
-from .aurora.model import AuroraForCausalLM as NativeAuroraForCausalLM
+from .aurora_config import AuroraConfig as NativeAuroraConfig
+from .aurora_model import AuroraForCausalLM as NativeAuroraForCausalLM
 
 
 class AuroraForCausalLM(PreTrainedModel):
@@ -17,6 +17,8 @@ class AuroraForCausalLM(PreTrainedModel):
     base_model_prefix = "aurora"
     main_input_name = "input_ids"
     _supports_cache_class = False
+    _tied_weights_keys = {}
+    all_tied_weights_keys = {}
 
     def __init__(self, config: AuroraHFConfig):
         super().__init__(config)

tokenizer_config.json

@@ -5,5 +5,5 @@
   "pad_token": "<pad>",
   "model_max_length": 2048,
   "clean_up_tokenization_spaces": false,
-  "chat_template": "{% for message in messages %}<|im_start|>{{ message['role'] }}\\n{{ message['content'] }}<|im_end|>\\n{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\\n{% endif %}"
+  "chat_template": "{% for message in messages %}<|im_start|>{{ message['role'] }}\n{{ message['content'] }}<|im_end|>\n{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
 }