Eve-2-MoE-IT-272M: heavy IT patch (open-perfectblend, LoRA r=128, merged)

config.json

@@ -0,0 +1,27 @@
+{
+  "architecture": "Eve-2-MoE",
+  "architectures": [
+    "DeepSeekMoE"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_eve.EveConfig",
+    "AutoModelForCausalLM": "modeling_eve.DeepSeekMoE"
+  },
+  "block_size": 2048,
+  "dtype": "float32",
+  "expert_intermediate_size": 1408,
+  "head_dim": 64,
+  "model_type": "eve_moe",
+  "n_embd": 512,
+  "n_head": 8,
+  "n_layer": 12,
+  "num_experts": 8,
+  "rope_theta": 10000.0,
+  "router_aux_loss_coef": 0.01,
+  "shared_expert_intermediate_size": 1408,
+  "top_k": 2,
+  "transformers_version": "5.1.0",
+  "use_cache": false,
+  "use_checkpointing": false,
+  "vocab_size": 50304
+}

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "do_sample": true,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "transformers_version": "5.1.0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6a93597822aa69152364ff5f15d6b52efc1b60303e371c6e27334007af748342
+size 1190949984

modeling_eve.py

@@ -0,0 +1,311 @@
+# 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
+
+
+class EveConfig(PretrainedConfig):
+    model_type = "eve_moe"
+    # Mapping for Transformers compatibility
+    attribute_map = {
+        "num_hidden_layers": "n_layer",
+        "num_attention_heads": "n_head",
+        "hidden_size": "n_embd",
+        "max_position_embeddings": "block_size",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 50304,
+        n_layer: int = 12,
+        n_embd: int = 512,
+        n_head: int = 8,
+        head_dim: int = 64,
+        block_size: int = 2048,
+        num_experts: int = 8,
+        top_k: int = 2,
+        expert_intermediate_size: int = 1408,
+        shared_expert_intermediate_size: int = 1408,
+        router_aux_loss_coef: float = 0.01,
+        use_checkpointing: bool = False,
+        rope_theta: float = 10000.0,
+        **kwargs: Any,
+    ):
+        self.vocab_size = vocab_size
+        self.n_layer = n_layer
+        self.n_embd = n_embd
+        self.n_head = n_head
+        self.head_dim = head_dim
+        self.block_size = block_size
+
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.expert_intermediate_size = expert_intermediate_size
+        self.shared_expert_intermediate_size = shared_expert_intermediate_size
+        self.router_aux_loss_coef = router_aux_loss_coef
+
+        self.use_checkpointing = use_checkpointing
+        self.rope_theta = rope_theta
+        super().__init__(**kwargs)
+
+
+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
+
+        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"]
+
+    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)
+
+    # --- 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,
+        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.float().view(-1, shift_logits.size(-1)),
+                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.
+        return {"input_ids": input_ids}

tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": false,
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}