modeling_eve.py

# modeling_eve.py
# Self-contained Eve MoE model definition with training-safe loss, PEFT compatibility,
# and Hugging Face generation support.
#
# Key fixes vs. earlier versions:
# - Correct *shifted* causal LM loss (predict token t+1 from position t).
# - Returns a proper Transformers ModelOutput (CausalLMOutputWithPast).
# - Implements get_input_embeddings / get_output_embeddings for PEFT checkpointing.
# - Supports prompt-masked SFT via ignore_index=-100.
#
# Notes:
# - This model does NOT implement kv-cache; generate() will work but be slower.
# - Attention masking for padding is not applied (is_causal=True); use right-padding.

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Tuple, Any, Dict

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

from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast


from configuration_eve import EveConfig


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:
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight


def precompute_rope_freqs(
    head_dim: int,
    max_seq_len: int,
    theta: float = 10000.0,
    device: Optional[torch.device] = None,
) -> torch.Tensor:
    """Precompute complex RoPE frequencies as cis values."""
    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(max_seq_len, device=device).float()
    freqs = torch.outer(t, freqs)  # [T, head_dim/2]
    return torch.polar(torch.ones_like(freqs), freqs)  # complex64


def apply_rope(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    """
    x: [B, H, T, D]
    freqs_cis: [T, D/2] complex
    """
    B, H, T, D = x.shape
    # [B,H,T,D/2] complex
    x_complex = torch.view_as_complex(x.float().reshape(B, H, T, D // 2, 2))
    freqs_cis = freqs_cis[:T].view(1, 1, T, D // 2)
    x_rotated = x_complex * freqs_cis
    return torch.view_as_real(x_rotated).reshape(B, H, T, D).type_as(x)


class MLP(nn.Module):
    def __init__(self, config: EveConfig, intermediate_size: Optional[int] = None):
        super().__init__()
        hidden_dim = intermediate_size or config.expert_intermediate_size
        self.w1 = nn.Linear(config.n_embd, hidden_dim, bias=False)
        self.w2 = nn.Linear(config.n_embd, hidden_dim, bias=False)
        self.c_proj = nn.Linear(hidden_dim, config.n_embd, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.c_proj(F.silu(self.w1(x)) * self.w2(x))


class SharedMoE(nn.Module):
    """
    Simple top-k MoE:
    - One shared expert always applied
    - N routed experts mixed by router weights
    - Aux loss encourages balanced expert usage (simple squared-mean heuristic)
    """

    def __init__(self, config: EveConfig):
        super().__init__()
        self.config = config
        self.top_k = config.top_k
        self.shared_expert = MLP(config, config.shared_expert_intermediate_size)
        self.experts = nn.ModuleList([MLP(config) for _ in range(config.num_experts)])
        self.router = nn.Linear(config.n_embd, config.num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        B, T, C = x.shape
        if self.top_k < 1 or self.top_k > self.config.num_experts:
            raise ValueError(f"Invalid MoE top_k={self.top_k}; must be in [1, {self.config.num_experts}]")

        shared_out = self.shared_expert(x)

        logits = self.router(x)  # [B,T,E]
        probs = F.softmax(logits, dim=-1)  # [B,T,E]
        top_k_weights, top_k_indices = torch.topk(probs, self.top_k, dim=-1)  # [B,T,K]
        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)

        # Aux loss: encourage balanced usage across experts
        flat_probs = probs.view(-1, self.config.num_experts)  # [B*T,E]
        expert_usage = flat_probs.mean(dim=0)  # [E]
        aux_loss = torch.sum(expert_usage * expert_usage) * self.config.num_experts

        routed_out = torch.zeros_like(x)
        flat_x = x.view(-1, C)  # [B*T,C]
        flat_indices = top_k_indices.view(-1, self.top_k)  # [B*T,K]
        flat_weights = top_k_weights.view(-1, self.top_k)  # [B*T,K]

        # NOTE: This routing loop is simple but not optimal.
        for i, expert in enumerate(self.experts):
            mask = flat_indices == i  # [B*T,K]
            batch_idx, rank_idx = torch.where(mask)
            if batch_idx.numel() > 0:
                expert_input = flat_x[batch_idx]
                expert_output = expert(expert_input)
                weight = flat_weights[batch_idx, rank_idx].unsqueeze(-1)
                routed_out.view(-1, C).index_add_(0, batch_idx, expert_output * weight)

        return shared_out + routed_out, aux_loss


class CausalSelfAttention(nn.Module):
    def __init__(self, config: EveConfig):
        super().__init__()
        self.n_head = config.n_head
        self.head_dim = config.head_dim
        self.n_embd = config.n_embd

        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)

    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape

        qkv = self.c_attn(x)
        q, k, v = qkv.split(self.n_embd, dim=2)

        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)  # [B,H,T,D]
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        q = apply_rope(q, freqs_cis)
        k = apply_rope(k, freqs_cis)

        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y)


class Block(nn.Module):
    def __init__(self, config: EveConfig):
        super().__init__()
        self.ln_1 = RMSNorm(config.n_embd)
        self.ln_2 = RMSNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.mlp = SharedMoE(config)

    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        x = x + self.attn(self.ln_1(x), freqs_cis)
        mlp_out, aux_loss = self.mlp(self.ln_2(x))
        x = x + mlp_out
        return x, aux_loss


class DeepSeekMoE(PreTrainedModel, GenerationMixin):
    config_class = EveConfig
    _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}

    # _tied_weights_keys = ["lm_head.weight"]  # <--- Removed to avoid conflict with PreTrainedModel internals

    def __init__(self, config: EveConfig):
        super().__init__(config)
        self.config = config

        self.transformer = nn.ModuleDict(
            dict(
                wte=nn.Embedding(config.vocab_size, config.n_embd),
                h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
                ln_f=RMSNorm(config.n_embd),
            )
        )
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # Tie weights (Embedding and LM head share the same base parameter)
        self.transformer.wte.weight = self.lm_head.weight

        freqs_cis = precompute_rope_freqs(config.head_dim, config.block_size, config.rope_theta)
        self.register_buffer("freqs_cis", freqs_cis, persistent=False)

        # Initialize weights and apply final processing
        self.post_init()

        # Harden generation_config to avoid invalid configs blocking save_pretrained()
        if hasattr(self, "generation_config") and self.generation_config is not None:
            g = self.generation_config
            # If not sampling, sampling-only knobs must be neutral.
            if not getattr(g, "do_sample", False):
                if getattr(g, "top_k", 0):
                    g.top_k = None
                if getattr(g, "top_p", 1.0) != 1.0:
                    g.top_p = None
                if getattr(g, "temperature", 1.0) != 1.0:
                    g.temperature = None

    # --- PEFT / HF compatibility hooks ---
    def get_input_embeddings(self) -> nn.Module:
        return self.transformer.wte

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.transformer.wte = value

    def get_output_embeddings(self) -> nn.Module:
        return self.lm_head

    def set_output_embeddings(self, value: nn.Module) -> None:
        self.lm_head = value

    # --- Forward ---
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        idx: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,  # accept + ignore
        labels: Optional[torch.LongTensor] = None,
        targets: Optional[torch.LongTensor] = None,
        **kwargs: Any,
    ) -> CausalLMOutputWithPast:
        """
        If labels/targets are provided, computes *shifted* causal LM loss:
            loss = CE(logits[:, :-1], labels[:, 1:])
        """
        if idx is None:
            if input_ids is None:
                raise ValueError("Must provide input_ids or idx.")
            idx = input_ids
        if targets is None:
            targets = labels

        B, T = idx.shape
        x = self.transformer.wte(idx)

        total_aux_loss: Optional[torch.Tensor] = None
        freqs_cis = self.freqs_cis.to(x.device)

        for block in self.transformer.h:
            x, aux_loss = block(x, freqs_cis[:T])
            total_aux_loss = aux_loss if total_aux_loss is None else (total_aux_loss + aux_loss)

        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)  # [B,T,V]

        loss = None
        if targets is not None:
            # Shift for causal LM
            if T < 2:
                # Nothing to predict; return aux-only if desired
                shift_logits = logits[:, :0, :]
                shift_labels = targets[:, :0]
            else:
                shift_logits = logits[:, :-1, :].contiguous()
                shift_labels = targets[:, 1:].contiguous()

            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)).to(torch.float32),
                shift_labels.view(-1),
                ignore_index=-100,
            )



            if total_aux_loss is not None and self.config.router_aux_loss_coef:
                loss = loss + (self.config.router_aux_loss_coef * total_aux_loss)

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=None,
        )

    # --- Generation ---
    def prepare_inputs_for_generation(self, input_ids: torch.LongTensor, **kwargs: Any) -> Dict[str, Any]:
        # No kv-cache support; always feed full sequence.
        out = {"input_ids": input_ids}
        # HF generate() may pass attention_mask; accept it even if we don't apply it.
        if "attention_mask" in kwargs and kwargs["attention_mask"] is not None:
            out["attention_mask"] = kwargs["attention_mask"]
        return out