""" Mixture of Experts (MoE) Layer - Cốt lõi của Nexus Coder ========================================================= 24 chuyên gia (experts) tổng cộng, chỉ 3 chuyên gia được kích hoạt mỗi token. Đạt được 10B tổng tham số với chỉ 1.5B tham số active. Tính năng: - Top-K routing với noise (load balancing) - Aux loss cho load balancing giữa các expert - Hỗ trợ SwiGLU experts """ import torch import torch.nn as nn import torch.nn.functional as F from typing import Tuple, Optional from .layers import SwiGLU class Expert(nn.Module): """Một chuyên gia (expert) - thực chất là một SwiGLU FFN. v0.3: hỗ trợ MLP-parallel (gate/up concat thành 1 matmul). """ def __init__(self, hidden_size: int, intermediate_size: int, parallel: bool = True): super().__init__() self.ffn = SwiGLU(hidden_size, intermediate_size, parallel=parallel) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.ffn(x) class Router(nn.Module): """Router/Gating network: quyết định token nào đi đến expert nào.""" def __init__(self, hidden_size: int, num_experts: int): super().__init__() self.gate = nn.Linear(hidden_size, num_experts, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.gate(x) def load_balancing_loss_func( gate_logits: torch.Tensor, num_experts: int, top_k: int, attention_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Tính auxiliary loss cho load balancing (Switch Transformer). attention_mask có thể là: - None: tất cả token đều valid - 2D bool [B, T]: True = valid token - 2D int [B, T]: 1 = valid, 0 = padding - 4D float [B, 1, 1, T]: 0 = valid, large_negative = padding """ if gate_logits is None: # gate_logits is None → cannot compute; return 0 on proper device return torch.tensor(0.0) # Normalize attention_mask → 1D bool [N_valid] if attention_mask is None: tokens_per_expert = gate_logits.shape[0] * gate_logits.shape[1] # 2D shape: [B, T] already flattened by caller, so gate_logits.shape[0] is N if gate_logits.dim() == 2: tokens_per_expert = gate_logits.shape[0] else: # Convert 4D mask to 2D bool if attention_mask.dim() == 4: # [B, 1, 1, T] with 0 / -inf values mask_2d = attention_mask.squeeze(1).squeeze(1) # [B, T] mask_bool = mask_2d > -1e9 elif attention_mask.dim() == 3: mask_bool = attention_mask.squeeze(1) > 0 elif attention_mask.dim() == 2: if attention_mask.dtype == torch.bool: mask_bool = attention_mask else: # 0/1 or 0/-inf if attention_mask.dtype.is_floating_point: mask_bool = attention_mask > -1e9 else: mask_bool = attention_mask > 0 else: mask_bool = None if mask_bool is None: tokens_per_expert = gate_logits.shape[0] else: tokens_per_expert = mask_bool.sum().item() if tokens_per_expert < 1: tokens_per_expert = gate_logits.shape[0] routing_weights = F.softmax(gate_logits, dim=-1) _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) expert_mask = F.one_hot(selected_experts, num_classes=num_experts) expert_mask = expert_mask.sum(dim=-2).float() tokens_per_expert_normalized = expert_mask.mean(dim=-2) router_prob_per_expert = routing_weights.mean(dim=-2) aux_loss = ( num_experts * (tokens_per_expert_normalized * router_prob_per_expert).sum() ) / max(tokens_per_expert, 1) return aux_loss class MixtureOfExperts(nn.Module): """MoE Layer với Top-K routing và load balancing.""" def __init__(self, config): super().__init__() self.config = config self.num_experts = config.num_experts self.num_active_experts = config.num_active_experts self.router_jitter_noise = config.router_jitter_noise self.aux_loss_coef = config.router_aux_loss_coef # Router self.router = Router(config.hidden_size, self.num_experts) # Experts (v0.3: MLP-parallel by default) mlp_parallel = getattr(config, "mlp_parallel", True) self.experts = nn.ModuleList([ Expert(config.hidden_size, config.intermediate_size, parallel=mlp_parallel) for _ in range(self.num_experts) ]) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: bsz, seq_len, hidden = hidden_states.shape flat_hidden = hidden_states.view(-1, hidden) # [N, H] # Router logits router_logits = self.router(flat_hidden) # [N, E] # Thêm noise trong training để encourage exploration if self.training and self.router_jitter_noise > 0: router_logits = router_logits + torch.randn_like(router_logits) * self.router_jitter_noise # Top-K routing routing_weights = F.softmax(router_logits, dim=-1) top_k_weights, top_k_indices = torch.topk( routing_weights, self.num_active_experts, dim=-1 ) top_k_weights = top_k_weights / (top_k_weights.sum(dim=-1, keepdim=True) + 1e-9) # Dispatch tokens to experts final_hidden = torch.zeros_like(flat_hidden) # Vectorized: iterate through experts for expert_idx in range(self.num_experts): # Find tokens that go to this expert expert_mask = (top_k_indices == expert_idx).any(dim=-1) # [N] if not expert_mask.any(): continue # Get token indices token_indices = expert_mask.nonzero(as_tuple=True)[0] # Get the corresponding weights expert_weights = top_k_weights[token_indices] # [num_tokens, top_k] expert_weight_for_this = (top_k_indices[token_indices] == expert_idx).float() * expert_weights expert_weight_for_this = expert_weight_for_this.sum(dim=-1) # [num_tokens] # Run expert expert_input = flat_hidden[token_indices] expert_output = self.experts[expert_idx](expert_input) expert_output = expert_output * expert_weight_for_this.unsqueeze(-1) final_hidden[token_indices] += expert_output # Load balancing loss aux_loss = load_balancing_loss_func( router_logits, self.num_experts, self.num_active_experts, attention_mask, ) final_hidden = final_hidden.view(bsz, seq_len, hidden) return final_hidden, aux_loss