""" Hugging Face Transformers Compatible Modeling Implementation for Mark-Engineer-V2-8B Self-contained causal language model architecture with Multi-Head Latent Attention (MLA), DeepSeekMoE (1 Shared + 24 Routed Experts, Top-2 active), and Multi-Token Prediction (MTP). """ import math from typing import List, Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.modeling_utils import PreTrainedModel try: from .configuration_mark_engineer_v2_8b import MarkEngineerV2_8BConfig except (ImportError, ValueError): from configuration_mark_engineer_v2_8b import MarkEngineerV2_8BConfig class MarkRMSNorm(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: variance = x.pow(2).mean(-1, keepdim=True) return x * torch.rsqrt(variance + self.eps) * self.weight def precompute_rope_freqs(head_dim: int, max_seq_len: int, theta: float = 500000.0) -> Tuple[torch.Tensor, torch.Tensor]: freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) t = torch.arange(max_seq_len, dtype=torch.float) freqs = torch.outer(t, freqs) cos = freqs.cos() sin = freqs.sin() return cos, sin def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: d = x.shape[-1] x1 = x[..., : d // 2] x2 = x[..., d // 2 :] rotated = torch.cat((-x2, x1), dim=-1) cos = cos.unsqueeze(0).unsqueeze(2) sin = sin.unsqueeze(0).unsqueeze(2) cos = torch.cat([cos, cos], dim=-1) sin = torch.cat([sin, sin], dim=-1) return (x * cos) + (rotated * sin) class MarkSwiGLUFFN(nn.Module): def __init__(self, dim: int, intermediate_size: int): super().__init__() self.w_gate = nn.Linear(dim, intermediate_size, bias=False) self.w_up = nn.Linear(dim, intermediate_size, bias=False) self.w_down = nn.Linear(intermediate_size, dim, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)) class MarkMLAAttention(nn.Module): def __init__(self, config: MarkEngineerV2_8BConfig, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.dim = config.hidden_size self.n_heads = config.num_attention_heads self.q_lora_rank = config.q_lora_rank self.kv_lora_rank = config.kv_lora_rank self.qk_nope_head_dim = config.qk_nope_head_dim self.qk_rope_head_dim = config.qk_rope_head_dim self.v_head_dim = config.v_head_dim self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False) self.q_norm = MarkRMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.qk_head_dim, bias=False) self.wkv_a = nn.Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim, bias=False) self.kv_norm = MarkRMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) self.wkv_b = nn.Linear( self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False, ) self.wo = nn.Linear(self.n_heads * self.v_head_dim, self.dim, bias=False) def forward( self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: bsz, seqlen, _ = x.shape q_latent = self.q_norm(self.wq_a(x)) q = self.wq_b(q_latent).view(bsz, seqlen, self.n_heads, self.qk_head_dim) q_nope = q[..., : self.qk_nope_head_dim] q_pe = q[..., self.qk_nope_head_dim :] kv_combined = self.wkv_a(x) kv_latent = self.kv_norm(kv_combined[..., : self.kv_lora_rank]) k_pe = kv_combined[..., self.kv_lora_rank :].view(bsz, seqlen, 1, self.qk_rope_head_dim) kv = self.wkv_b(kv_latent).view(bsz, seqlen, self.n_heads, self.qk_nope_head_dim + self.v_head_dim) k_nope = kv[..., : self.qk_nope_head_dim] v = kv[..., self.qk_nope_head_dim :] q_pe = apply_rotary_emb(q_pe, cos[:seqlen], sin[:seqlen]) k_pe = apply_rotary_emb(k_pe.expand(-1, -1, self.n_heads, -1), cos[:seqlen], sin[:seqlen]) k = torch.cat([k_nope, k_pe], dim=-1) q = torch.cat([q_nope, q_pe], dim=-1) q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.qk_head_dim) if attention_mask is not None: scores = scores + attention_mask attn_weights = F.softmax(scores, dim=-1) output = torch.matmul(attn_weights, v) output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1) return self.wo(output) class MarkMoELayer(nn.Module): def __init__(self, config: MarkEngineerV2_8BConfig): super().__init__() self.n_routed = config.n_routed_experts self.n_shared = config.n_shared_experts self.top_k = config.num_experts_per_tok self.dim = config.hidden_size self.moe_intermediate = config.moe_intermediate_size self.gate = nn.Linear(self.dim, self.n_routed, bias=False) self.register_buffer("router_bias", torch.zeros(self.n_routed)) self.router_bias_rate = config.router_bias_update_rate self.routed_experts = nn.ModuleList([ MarkSwiGLUFFN(self.dim, self.moe_intermediate) for _ in range(self.n_routed) ]) if self.n_shared > 0: self.shared_experts = MarkSwiGLUFFN(self.dim, self.moe_intermediate * self.n_shared) else: self.shared_experts = None self.last_router_z_loss = None def forward(self, x: torch.Tensor) -> torch.Tensor: bsz, seqlen, dim = x.shape flat_x = x.view(-1, dim) logits = self.gate(flat_x) if self.training: self.last_router_z_loss = torch.logsumexp(logits, dim=-1).pow(2).mean() biased_logits = logits + self.router_bias topk_weights, topk_indices = torch.topk(F.softmax(biased_logits, dim=-1), self.top_k, dim=-1) topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) out = torch.zeros_like(flat_x) for i, expert in enumerate(self.routed_experts): mask = (topk_indices == i).any(dim=-1) if mask.any(): idx_in_topk = (topk_indices[mask] == i).nonzero(as_tuple=True)[1] w = topk_weights[mask, idx_in_topk].unsqueeze(-1) out[mask] += w * expert(flat_x[mask]) if self.shared_experts is not None: out += self.shared_experts(flat_x) return out.view(bsz, seqlen, dim) class MarkEngineerV2Block(nn.Module): def __init__(self, config: MarkEngineerV2_8BConfig, layer_idx: int): super().__init__() self.attn_norm = MarkRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.attn = MarkMLAAttention(config, layer_idx) self.ffn_norm = MarkRMSNorm(config.hidden_size, eps=config.rms_norm_eps) if layer_idx < config.n_dense_layers: self.ffn = MarkSwiGLUFFN(config.hidden_size, config.intermediate_size) self.is_moe = False else: self.ffn = MarkMoELayer(config) self.is_moe = True def forward( self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: h = x + self.attn(self.attn_norm(x), cos, sin, attention_mask) out = h + self.ffn(self.ffn_norm(h)) return out class MarkEngineerV2ForCausalLM(PreTrainedModel): config_class = MarkEngineerV2_8BConfig base_model_prefix = "model" _supports_flash_attn_2 = False def __init__(self, config: MarkEngineerV2_8BConfig): super().__init__(config) self.config = config self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = nn.ModuleList([ MarkEngineerV2Block(config, i) for i in range(config.num_hidden_layers) ]) self.norm = MarkRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) if config.mtp_depth > 0: self.mtp_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) self.mtp_norm = MarkRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.mtp_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) else: self.mtp_proj = None if config.tie_word_embeddings: self.lm_head.weight = self.embed_tokens.weight cos, sin = precompute_rope_freqs( config.qk_rope_head_dim, config.max_position_embeddings, theta=config.rope_theta, ) self.register_buffer("rope_cos", cos, persistent=False) self.register_buffer("rope_sin", sin, persistent=False) self.post_init() def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value def get_output_embeddings(self): return self.lm_head def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.LongTensor] = None, return_mtp: bool = False, **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: bsz, seqlen = input_ids.shape x = self.embed_tokens(input_ids) if attention_mask is None: causal_mask = torch.triu( torch.full((seqlen, seqlen), float("-inf"), device=input_ids.device), diagonal=1, ).unsqueeze(0).unsqueeze(0) else: causal_mask = attention_mask for layer in self.layers: x = layer(x, self.rope_cos, self.rope_sin, causal_mask) h = self.norm(x) logits = self.lm_head(h) mtp_logits = None if return_mtp and self.mtp_proj is not None and seqlen > 1: next_emb = self.embed_tokens(input_ids[:, 1:]) h_prev = h[:, :-1] cat_h = torch.cat([h_prev, next_emb], dim=-1) h_mtp = self.mtp_norm(self.mtp_proj(cat_h)) mtp_logits = self.mtp_head(h_mtp) loss = None if labels is not None: shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss_fct = nn.CrossEntropyLoss() loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1)) if mtp_logits is not None and seqlen > 2: shift_mtp_logits = mtp_logits[..., :-1, :].contiguous() shift_mtp_labels = labels[..., 2:].contiguous() loss_mtp = loss_fct(shift_mtp_logits.view(-1, self.config.vocab_size), shift_mtp_labels.view(-1)) loss = loss + (self.config.mtp_loss_weight * loss_mtp) return CausalLMOutputWithPast( loss=loss, logits=logits, )