"""Sequence-routed mixture-of-experts causal language model.""" from __future__ import annotations import math from dataclasses import dataclass from pathlib import Path from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from safetensors import safe_open from transformers import PreTrainedModel from transformers.cache_utils import Cache, DynamicCache from transformers.generation.utils import GenerationMixin from transformers.utils.hub import cached_file from transformers.utils import ModelOutput try: from .configuration_slmoe import SLMoEConfig except ImportError: # Allows the standalone training script to import local code. from configuration_slmoe import SLMoEConfig def _resolve_expert_weights_path(config: SLMoEConfig) -> Path: """Resolve the SafeTensors file used as the disk-backed expert store.""" source = str(getattr(config, "_name_or_path", "")) source_path = Path(source).expanduser() if source else None if source_path is not None and source_path.is_dir(): candidate = source_path / config.expert_weights_file if candidate.is_file(): return candidate.resolve() if source: resolved = cached_file(source, config.expert_weights_file) if resolved is not None: return Path(resolved) raise FileNotFoundError( "Disk-backed experts require the checkpoint's " f"{config.expert_weights_file!r} file" ) @dataclass class SLMoECausalLMOutputWithPast(ModelOutput): loss: Optional[torch.Tensor] = None logits: Optional[torch.Tensor] = None past_key_values: Optional[Cache] = None router_aux_loss: Optional[torch.Tensor] = None router_z_loss: Optional[torch.Tensor] = None expert_indices: Optional[torch.LongTensor] = None expert_weights: Optional[torch.Tensor] = None class SLMoERMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: states = hidden_states.float() states = states * torch.rsqrt(states.square().mean(-1, keepdim=True) + self.eps) return (states * self.weight.float()).to(hidden_states.dtype) def _rope_cos_sin( head_dim: int, positions: torch.Tensor, theta: float, ) -> tuple[torch.Tensor, torch.Tensor]: inv_freq = 1.0 / ( theta ** ( torch.arange(0, head_dim, 2, dtype=torch.float32, device=positions.device) / head_dim ) ) frequencies = torch.outer(positions.float(), inv_freq) return frequencies.cos(), frequencies.sin() def _apply_rope( query: torch.Tensor, key: torch.Tensor, cosine: torch.Tensor, sine: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: query_dtype = query.dtype key_dtype = key.dtype cosine = cosine[None, None, :, :] sine = sine[None, None, :, :] query_pairs = query.float().reshape(*query.shape[:-1], -1, 2) key_pairs = key.float().reshape(*key.shape[:-1], -1, 2) query_even, query_odd = query_pairs.unbind(-1) key_even, key_odd = key_pairs.unbind(-1) query = torch.stack( (query_even * cosine - query_odd * sine, query_even * sine + query_odd * cosine), dim=-1, ).flatten(-2) key = torch.stack( (key_even * cosine - key_odd * sine, key_even * sine + key_odd * cosine), dim=-1, ).flatten(-2) return query.to(query_dtype), key.to(key_dtype) class SLMoECache(DynamicCache): """K/V cache carrying the one routing decision for the whole response.""" def __init__(self, config: SLMoEConfig): try: super().__init__(config=config) except TypeError: super().__init__() self.expert_indices: torch.LongTensor | None = None self.expert_weights: torch.Tensor | None = None self.active_expert_weights: dict[ int, tuple[torch.Tensor, torch.Tensor, torch.Tensor], ] = {} def set_routing( self, expert_indices: torch.LongTensor, expert_weights: torch.Tensor, ) -> None: if self.expert_indices is not None: raise RuntimeError("The sequence routing plan may only be set once") self.expert_indices = expert_indices self.expert_weights = expert_weights def get_active_experts( self, layer_idx: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: return self.active_expert_weights.get(layer_idx) def set_active_experts( self, layer_idx: int, weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor], ) -> None: if layer_idx in self.active_expert_weights: raise RuntimeError(f"Active experts for layer {layer_idx} were already cached") self.active_expert_weights[layer_idx] = weights def reorder_cache(self, beam_idx: torch.LongTensor): super().reorder_cache(beam_idx) if self.expert_indices is not None: beam_idx = beam_idx.to(self.expert_indices.device) self.expert_indices = self.expert_indices.index_select(0, beam_idx) self.expert_weights = self.expert_weights.index_select(0, beam_idx) self.active_expert_weights = { layer_idx: tuple( tensor.index_select(0, beam_idx.to(tensor.device)) for tensor in weights ) for layer_idx, weights in self.active_expert_weights.items() } def batch_repeat_interleave(self, repeats: int): super().batch_repeat_interleave(repeats) if self.expert_indices is not None: self.expert_indices = self.expert_indices.repeat_interleave(repeats, dim=0) self.expert_weights = self.expert_weights.repeat_interleave(repeats, dim=0) self.active_expert_weights = { layer_idx: tuple( tensor.repeat_interleave(repeats, dim=0) for tensor in weights ) for layer_idx, weights in self.active_expert_weights.items() } def batch_select_indices(self, indices: torch.Tensor): super().batch_select_indices(indices) if self.expert_indices is not None: indices = indices.to(self.expert_indices.device) self.expert_indices = self.expert_indices.index_select(0, indices) self.expert_weights = self.expert_weights.index_select(0, indices) self.active_expert_weights = { layer_idx: tuple( tensor.index_select(0, indices.to(tensor.device)) for tensor in weights ) for layer_idx, weights in self.active_expert_weights.items() } class SLMoEAttention(nn.Module): def __init__(self, config: SLMoEConfig, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads self.head_dim = config.head_dim self.num_kv_groups = self.num_heads // self.num_kv_heads self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False) self.o_proj.SLMOE_SCALE_INIT = True self.q_norm = SLMoERMSNorm(self.head_dim, config.rms_norm_eps) self.k_norm = SLMoERMSNorm(self.head_dim, config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, cosine: torch.Tensor, sine: torch.Tensor, attention_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, ) -> torch.Tensor: batch_size, query_length, _ = hidden_states.shape query = self.q_proj(hidden_states).view( batch_size, query_length, self.num_heads, self.head_dim ).transpose(1, 2) key = self.k_proj(hidden_states).view( batch_size, query_length, self.num_kv_heads, self.head_dim ).transpose(1, 2) value = self.v_proj(hidden_states).view( batch_size, query_length, self.num_kv_heads, self.head_dim ).transpose(1, 2) query = self.q_norm(query) key = self.k_norm(key) query, key = _apply_rope(query, key, cosine, sine) past_length = 0 if past_key_values is not None: past_length = past_key_values.get_seq_length(self.layer_idx) key, value = past_key_values.update(key, value, self.layer_idx) key_length = key.size(-2) key = key.repeat_interleave(self.num_kv_groups, dim=1) value = value.repeat_interleave(self.num_kv_groups, dim=1) is_causal = query_length > 1 and past_length == 0 and attention_mask is None sdpa_mask = None if not is_causal and query_length > 1: query_positions = past_length + torch.arange(query_length, device=query.device) key_positions = torch.arange(key_length, device=query.device) sdpa_mask = (key_positions[None, :] <= query_positions[:, None])[None, None] if attention_mask is not None: key_padding = attention_mask.to(torch.bool) if key_padding.size(-1) < key_length: key_padding = F.pad(key_padding, (key_length - key_padding.size(-1), 0), value=True) else: key_padding = key_padding[:, -key_length:] key_padding = key_padding[:, None, None, :] sdpa_mask = key_padding if sdpa_mask is None else sdpa_mask & key_padding is_causal = False output = F.scaled_dot_product_attention( query, key, value, attn_mask=sdpa_mask, is_causal=is_causal, ) output = output.transpose(1, 2).contiguous().view( batch_size, query_length, self.num_heads * self.head_dim ) return self.o_proj(output) class SLMoESequenceRouter(nn.Module): """Choose one fixed expert set from a causal prefix of each sequence.""" def __init__(self, config: SLMoEConfig): super().__init__() self.num_experts = config.num_experts self.top_k = config.num_experts_per_sequence self.prefix_length = config.router_prefix_length self.jitter_noise = config.router_jitter_noise self.norm = SLMoERMSNorm(config.hidden_size, config.rms_norm_eps) self.proj = nn.Linear(config.hidden_size, config.num_experts, bias=False) def prefix_mask( self, token_embeddings: torch.Tensor, attention_mask: torch.Tensor | None, ) -> torch.Tensor: batch_size, sequence_length, _ = token_embeddings.shape if attention_mask is None: positions = torch.arange(sequence_length, device=token_embeddings.device) return (positions < self.prefix_length).expand(batch_size, -1) valid = attention_mask[:, -sequence_length:].to(torch.bool) valid_order = valid.long().cumsum(dim=-1) return valid & (valid_order <= self.prefix_length) def forward( self, token_embeddings: torch.Tensor, attention_mask: torch.Tensor | None, ) -> tuple[torch.LongTensor, torch.Tensor, torch.Tensor, torch.Tensor]: prefix_mask = self.prefix_mask(token_embeddings, attention_mask) normalized = self.norm(token_embeddings) mask = prefix_mask.unsqueeze(-1).to(normalized.dtype) pooled = (normalized * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0) if self.training and self.jitter_noise > 0: pooled = pooled * torch.empty_like(pooled).uniform_( 1.0 - self.jitter_noise, 1.0 + self.jitter_noise, ) router_logits = self.proj(pooled).float() router_probs = F.softmax(router_logits, dim=-1, dtype=torch.float32) top_probs, expert_indices = torch.topk( router_probs, k=self.top_k, dim=-1, sorted=True, ) expert_weights = top_probs / top_probs.sum(dim=-1, keepdim=True).clamp_min(1e-9) selected_fraction = F.one_hot( expert_indices, num_classes=self.num_experts, ).float().mean(dim=(0, 1)) probability_fraction = router_probs.mean(dim=0) auxiliary_loss = self.num_experts * torch.sum( selected_fraction * probability_fraction ) router_z_loss = torch.logsumexp(router_logits, dim=-1).square().mean() return expert_indices, expert_weights, auxiliary_loss, router_z_loss class SLMoEExpertBank(nn.Module): """Resident or disk-backed expert weights for one transformer layer.""" def __init__( self, config: SLMoEConfig, layer_idx: int, expert_weights_path: Path | None, ): super().__init__() experts = config.num_experts hidden = config.hidden_size intermediate = config.expert_intermediate_size self.layer_idx = layer_idx self.expert_storage = config.expert_storage self.expert_weights_path = expert_weights_path self.output_scale = config.expert_output_scale self.num_experts = experts self.hidden_size = hidden self.intermediate_size = intermediate if self.expert_storage == "ram": self.gate_weight = nn.Parameter(torch.empty(experts, intermediate, hidden)) self.up_weight = nn.Parameter(torch.empty(experts, intermediate, hidden)) self.down_weight = nn.Parameter(torch.empty(experts, hidden, intermediate)) nn.init.normal_(self.gate_weight, mean=0.0, std=config.initializer_range) nn.init.normal_(self.up_weight, mean=0.0, std=config.initializer_range) down_std = config.initializer_range * (2 * config.num_hidden_layers) ** -0.5 nn.init.normal_(self.down_weight, mean=0.0, std=down_std) else: if self.expert_weights_path is None: raise ValueError("Disk-backed experts require an expert weights path") self.register_parameter("gate_weight", None) self.register_parameter("up_weight", None) self.register_parameter("down_weight", None) def _key(self, weight_name: str) -> str: return f"transformer.h.{self.layer_idx}.experts.{weight_name}" def load_active_weights( self, expert_indices: torch.LongTensor, *, device: torch.device, dtype: torch.dtype, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Materialize only this batch's selected experts from SafeTensors.""" if self.expert_storage != "disk" or self.expert_weights_path is None: raise RuntimeError("load_active_weights is only valid for disk-backed experts") if self.training and torch.is_grad_enabled(): raise RuntimeError("expert_storage='disk' is inference-only") batch_size, top_k = expert_indices.shape flat_indices = expert_indices.detach().to(device="cpu", dtype=torch.long).reshape(-1) unique_indices, inverse = torch.unique( flat_indices, sorted=True, return_inverse=True, ) selected = [int(index) for index in unique_indices.tolist()] with safe_open(self.expert_weights_path, framework="pt", device="cpu") as handle: gate_unique = torch.stack( [handle.get_slice(self._key("gate_weight"))[index] for index in selected] ) up_unique = torch.stack( [handle.get_slice(self._key("up_weight"))[index] for index in selected] ) down_unique = torch.stack( [handle.get_slice(self._key("down_weight"))[index] for index in selected] ) def expand_selected(tensor: torch.Tensor) -> torch.Tensor: tensor = tensor.index_select(0, inverse) tensor = tensor.reshape(batch_size, top_k, *tensor.shape[1:]) return tensor.to(device=device, dtype=dtype) return ( expand_selected(gate_unique), expand_selected(up_unique), expand_selected(down_unique), ) def forward( self, hidden_states: torch.Tensor, expert_indices: torch.LongTensor, expert_weights: torch.Tensor, active_weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: if active_weights is not None: gate_weight, up_weight, down_weight = active_weights elif self.expert_storage == "ram": gate_weight = self.gate_weight[expert_indices] up_weight = self.up_weight[expert_indices] down_weight = self.down_weight[expert_indices] else: gate_weight, up_weight, down_weight = self.load_active_weights( expert_indices, device=hidden_states.device, dtype=hidden_states.dtype, ) gate = torch.einsum("bsh,bkih->bski", hidden_states, gate_weight) up = torch.einsum("bsh,bkih->bski", hidden_states, up_weight) activated = F.silu(gate) * up activated = activated * expert_weights[:, None, :, None].to(activated.dtype) output = torch.einsum("bski,bkhi->bsh", activated, down_weight) return output * self.output_scale class SLMoEBlock(nn.Module): def __init__( self, config: SLMoEConfig, layer_idx: int, expert_weights_path: Path | None, ): super().__init__() self.layer_idx = layer_idx self.input_norm = SLMoERMSNorm(config.hidden_size, config.rms_norm_eps) self.attention = SLMoEAttention(config, layer_idx) self.post_attention_norm = SLMoERMSNorm(config.hidden_size, config.rms_norm_eps) self.experts = SLMoEExpertBank(config, layer_idx, expert_weights_path) def forward( self, hidden_states: torch.Tensor, cosine: torch.Tensor, sine: torch.Tensor, expert_indices: torch.LongTensor, expert_weights: torch.Tensor, attention_mask: torch.Tensor | None, past_key_values: Cache | None, active_expert_weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: hidden_states = hidden_states + self.attention( self.input_norm(hidden_states), cosine, sine, attention_mask=attention_mask, past_key_values=past_key_values, ) return hidden_states + self.experts( self.post_attention_norm(hidden_states), expert_indices, expert_weights, active_weights=active_expert_weights, ) class SLMoEPreTrainedModel(PreTrainedModel): config_class = SLMoEConfig base_model_prefix = "transformer" supports_gradient_checkpointing = False _no_split_modules = ["SLMoEBlock"] _supports_sdpa = True _supports_cache_class = True def _init_weights(self, module: nn.Module): std = self.config.initializer_range if hasattr(module, "SLMOE_SCALE_INIT"): std *= (2 * self.config.num_hidden_layers) ** -0.5 if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=std) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=std) class SLMoEForCausalLM(SLMoEPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} _keys_to_ignore_on_load_unexpected = [ r"transformer\.h\.\d+\.experts\.(gate_weight|up_weight|down_weight)" ] @classmethod def _supports_default_dynamic_cache(cls) -> bool: return False def __init__(self, config: SLMoEConfig): super().__init__(config) self.expert_weights_path = ( _resolve_expert_weights_path(config) if config.expert_storage == "disk" else None ) self.router = SLMoESequenceRouter(config) self.transformer = nn.ModuleDict( { "wte": nn.Embedding(config.vocab_size, config.hidden_size), "h": nn.ModuleList( [ SLMoEBlock(config, index, self.expert_weights_path) for index in range(config.num_hidden_layers) ] ), "ln_f": SLMoERMSNorm(config.hidden_size, config.rms_norm_eps), } ) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.embedding_scale = math.sqrt(config.hidden_size) self.post_init() if config.tie_word_embeddings: self.tie_weights() def get_input_embeddings(self): return self.transformer["wte"] def set_input_embeddings(self, value): self.transformer["wte"] = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, value): self.lm_head = value def _route( self, token_embeddings: torch.Tensor, attention_mask: torch.Tensor | None, past_key_values: SLMoECache | None, ) -> tuple[torch.LongTensor, torch.Tensor, torch.Tensor, torch.Tensor]: if past_key_values is not None and past_key_values.expert_indices is not None: zero = token_embeddings.new_zeros((), dtype=torch.float32) return ( past_key_values.expert_indices, past_key_values.expert_weights, zero, zero, ) expert_indices, expert_weights, auxiliary_loss, router_z_loss = self.router( token_embeddings, attention_mask, ) if past_key_values is not None: past_key_values.set_routing(expert_indices, expert_weights) return expert_indices, expert_weights, auxiliary_loss, router_z_loss def forward( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, use_cache: Optional[bool] = None, only_use_active_experts: Optional[bool] = None, **kwargs, ) -> SLMoECausalLMOutputWithPast: if self.config.expert_storage == "disk" and only_use_active_experts is False: raise ValueError( "expert_storage='disk' requires only_use_active_experts=True; " "use expert_storage='ram' to keep the complete expert bank resident" ) if only_use_active_experts is None: only_use_active_experts = ( self.config.expert_storage == "disk" or self.config.only_use_active_experts ) if use_cache is None: use_cache = self.config.use_cache and labels is None if use_cache and past_key_values is None: past_key_values = SLMoECache(self.config) if use_cache and not isinstance(past_key_values, SLMoECache): raise TypeError("SLMoE requires SLMoECache to preserve sequence routing") if not use_cache: past_key_values = None past_length = past_key_values.get_seq_length() if past_key_values is not None else 0 sequence_length = input_ids.size(1) total_length = past_length + sequence_length if total_length > self.config.max_position_embeddings: raise ValueError( f"Sequence length {total_length} exceeds {self.config.max_position_embeddings}" ) token_embeddings = self.transformer["wte"](input_ids) expert_indices, expert_weights, router_aux_loss, router_z_loss = self._route( token_embeddings, attention_mask, past_key_values, ) hidden_states = token_embeddings * self.embedding_scale positions = torch.arange( past_length, total_length, dtype=torch.float32, device=input_ids.device, ) cosine, sine = _rope_cos_sin( self.config.head_dim, positions, self.config.rope_theta, ) for block in self.transformer["h"]: active_expert_weights = None if self.config.expert_storage == "disk": if past_key_values is not None: active_expert_weights = past_key_values.get_active_experts(block.layer_idx) if active_expert_weights is None: active_expert_weights = block.experts.load_active_weights( expert_indices, device=hidden_states.device, dtype=hidden_states.dtype, ) if past_key_values is not None: past_key_values.set_active_experts( block.layer_idx, active_expert_weights, ) hidden_states = block( hidden_states, cosine, sine, expert_indices, expert_weights, attention_mask, past_key_values, active_expert_weights=active_expert_weights, ) hidden_states = self.transformer["ln_f"](hidden_states) logits = self.lm_head(hidden_states) loss = None if labels is not None: shift_logits = logits[..., :-1, :].float().contiguous() shift_labels = labels[..., 1:].clone().contiguous() prefix_mask = self.router.prefix_mask(token_embeddings, attention_mask) sequence_positions = torch.arange(sequence_length, device=input_ids.device) last_prefix_position = torch.where( prefix_mask, sequence_positions[None, :], -1, ).amax(dim=-1) prediction_positions = sequence_positions[:-1][None, :] shift_labels[prediction_positions < last_prefix_position[:, None]] = -100 ce_loss = F.cross_entropy( shift_logits.reshape(-1, shift_logits.size(-1)), shift_labels.reshape(-1), ignore_index=-100, ) loss = ( ce_loss + self.config.router_aux_loss_coeff * router_aux_loss + self.config.router_z_loss_coeff * router_z_loss ) return SLMoECausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=past_key_values, router_aux_loss=router_aux_loss, router_z_loss=router_z_loss, expert_indices=expert_indices, expert_weights=expert_weights, ) SLMoEForCausalLM.register_for_auto_class("AutoModelForCausalLM") __all__ = [ "SLMoECache", "SLMoECausalLMOutputWithPast", "SLMoEForCausalLM", "SLMoEPreTrainedModel", ]