modeling_slmoe.py

"""Sequence-routed mixture-of-experts causal language model."""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation.utils import GenerationMixin
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


@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

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

    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)

    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)


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):
    """Batched expert weights; only the sequence-selected slices are evaluated."""

    def __init__(self, config: SLMoEConfig):
        super().__init__()
        experts = config.num_experts
        hidden = config.hidden_size
        intermediate = config.expert_intermediate_size
        self.output_scale = config.expert_output_scale
        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)

    def forward(
        self,
        hidden_states: torch.Tensor,
        expert_indices: torch.LongTensor,
        expert_weights: torch.Tensor,
    ) -> torch.Tensor:
        gate_weight = self.gate_weight[expert_indices]
        up_weight = self.up_weight[expert_indices]
        down_weight = self.down_weight[expert_indices]
        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):
        super().__init__()
        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)

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


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"}

    @classmethod
    def _supports_default_dynamic_cache(cls) -> bool:
        return False

    def __init__(self, config: SLMoEConfig):
        super().__init__(config)
        self.router = SLMoESequenceRouter(config)
        self.transformer = nn.ModuleDict(
            {
                "wte": nn.Embedding(config.vocab_size, config.hidden_size),
                "h": nn.ModuleList(
                    [SLMoEBlock(config, index) 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,
        **kwargs,
    ) -> SLMoECausalLMOutputWithPast:
        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"]:
            hidden_states = block(
                hidden_states,
                cosine,
                sine,
                expert_indices,
                expert_weights,
                attention_mask,
                past_key_values,
            )
        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",
]