| """ |
| 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, |
| ) |
|
|
