""" TRM MoE modeling. Dense SwiGLU MLP -> Top-K Sparse MoE. """ import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.generation import GenerationMixin from transformers.modeling_outputs import CausalLMOutputWithPast try: from .configuration_trm_moe import TRMMoEConfig except ImportError: from configuration_trm_moe import TRMMoEConfig def apply_rope(x, cos, sin): S = x.shape[2] cos = cos[:, :, :S, :].to(dtype=x.dtype, device=x.device) sin = sin[:, :, :S, :].to(dtype=x.dtype, device=x.device) x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) class SwiGLUExpert(nn.Module): def __init__(self, dim, hidden): super().__init__() self.gate_proj = nn.Linear(dim, hidden, bias=False) self.up_proj = nn.Linear(dim, hidden, bias=False) self.down_proj = nn.Linear(hidden, dim, bias=False) def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class SparseMoELayer(nn.Module): def __init__(self, config): super().__init__() self.num_experts = config.num_experts self.top_k = config.num_experts_per_token self.aux_loss_coeff = config.moe_aux_loss_coeff self.z_loss_coeff = config.moe_z_loss_coeff hidden = config.mlp_hidden_size or int(config.dim * config.mlp_ratio) self.router = nn.Linear(config.dim, self.num_experts, bias=False) self.experts = nn.ModuleList([ SwiGLUExpert(config.dim, hidden) for _ in range(self.num_experts) ]) def forward(self, x): B, S, D = x.shape x_flat = x.reshape(-1, D) N = x_flat.shape[0] router_logits = self.router(x_flat) router_logits_fp32 = router_logits.float() z_loss = torch.logsumexp(router_logits_fp32, dim=-1).square().mean() router_probs = F.softmax(router_logits_fp32, dim=-1) topk_weights, topk_indices = torch.topk(router_probs, self.top_k, dim=-1) topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True).clamp_min(1e-9) topk_weights = topk_weights.to(dtype=x.dtype) expert_mask = F.one_hot(topk_indices, num_classes=self.num_experts).float() expert_mask_flat = expert_mask.sum(dim=1) f = expert_mask_flat.mean(dim=0) p = router_probs.mean(dim=0) aux_loss = self.num_experts * (f * p).sum() out = torch.zeros_like(x_flat) for k in range(self.top_k): expert_idx_k = topk_indices[:, k] weights_k = topk_weights[:, k] for e in range(self.num_experts): token_mask = expert_idx_k == e if not token_mask.any(): continue y = self.experts[e](x_flat[token_mask]) out[token_mask] += weights_k[token_mask].unsqueeze(-1) * y total_aux = self.aux_loss_coeff * aux_loss + self.z_loss_coeff * z_loss return out.view(B, S, D), total_aux class TRMAttention(nn.Module): def __init__(self, config): super().__init__() self.n_heads = config.n_heads self.head_dim = config.head_dim self.qkv = nn.Linear(config.dim, 3 * config.dim, bias=False) self.out = nn.Linear(config.dim, config.dim, bias=False) def forward(self, x, mask, cos, sin): B, S, D = x.shape q, k, v = self.qkv(x).chunk(3, dim=-1) q = q.view(B, S, self.n_heads, self.head_dim).transpose(1, 2) k = k.view(B, S, self.n_heads, self.head_dim).transpose(1, 2) v = v.view(B, S, self.n_heads, self.head_dim).transpose(1, 2) q = apply_rope(q, cos, sin) k = apply_rope(k, cos, sin) attn_mask = mask[:, None, :, :] y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) y = y.transpose(1, 2).contiguous().view(B, S, D) return self.out(y) class TRMMoEBlock(nn.Module): def __init__(self, config): super().__init__() self.res = config.residual_scale self.norm1 = nn.RMSNorm(config.dim) self.attn = TRMAttention(config) self.norm2 = nn.RMSNorm(config.dim) self.moe = SparseMoELayer(config) self.attn_gate = nn.Parameter(torch.ones(config.dim)) self.mlp_gate = nn.Parameter(torch.ones(config.dim)) def forward(self, x, mask, cos, sin): x = x + self.res * torch.sigmoid(self.attn_gate).view(1, 1, -1) * self.attn(self.norm1(x), mask, cos, sin) moe_out, aux_loss = self.moe(self.norm2(x)) x = x + self.res * torch.sigmoid(self.mlp_gate).view(1, 1, -1) * moe_out return x, aux_loss class TRMMoEForCausalLM(PreTrainedModel, GenerationMixin): config_class = TRMMoEConfig supports_gradient_checkpointing = True def __init__(self, config): super().__init__(config) self.token_emb = nn.Embedding(config.vocab_size, config.dim) self.block = TRMMoEBlock(config) self.norm = nn.RMSNorm(config.dim) self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False) pos = torch.arange(config.max_seq_len).float() theta = 1.0 / ( 10000.0 ** (torch.arange(0, config.head_dim // 2).float() / (config.head_dim // 2)) ) freqs = torch.outer(pos, theta) self.register_buffer("rope_cos", freqs.cos().view(1, 1, config.max_seq_len, -1), persistent=True) self.register_buffer("rope_sin", freqs.sin().view(1, 1, config.max_seq_len, -1), persistent=True) self.post_init() def get_input_embeddings(self): return self.token_emb def set_input_embeddings(self, value): self.token_emb = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, value): self.lm_head = value def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs): return { "input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False, } def forward(self, input_ids, attention_mask=None, labels=None, **kwargs): B, S = input_ids.shape if S > self.config.max_seq_len: raise ValueError(f"Sequence length {S} > max_seq_len {self.config.max_seq_len}") x = self.token_emb(input_ids) causal = torch.tril(torch.ones(S, S, device=input_ids.device, dtype=torch.bool)) mask = causal.unsqueeze(0).expand(B, -1, -1) if attention_mask is not None: key_mask = attention_mask[:, None, :].to(torch.bool) mask = mask & key_mask total_aux_loss = x.new_tensor(0.0) for _ in range(self.config.recurrence_steps): x, aux_loss = self.block(x, mask, self.rope_cos, self.rope_sin) total_aux_loss = total_aux_loss + aux_loss.to(dtype=x.dtype) logits = self.lm_head(self.norm(x)) loss = None if labels is not None: shift_logits = logits[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() lm_loss = F.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)).float(), shift_labels.view(-1), ignore_index=-100, ) loss = lm_loss + total_aux_loss.float() / max(1, self.config.recurrence_steps) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=None, hidden_states=None, attentions=None, )