d9d 0.1.0__py3-none-any.whl

d9d/module/block/moe/router.py

@@ -0,0 +1,103 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from d9d.module.base import ModuleLateInit
+
+
+class TopKRouter(nn.Module, ModuleLateInit):
+    """
+    Selects the top-K experts based on a learned gating mechanism.
+
+    This router:
+
+    1. Projects input tokens into expert space
+    2. Applies softmax, optionally adds expert bias to influence selection
+    3. Selects the experts with the highest probabilities
+    4. Selected probabilities are then re-normalized to sum to 1 if needed.
+    """
+
+    def __init__(
+            self,
+            dim: int,
+            num_experts: int,
+            top_k: int,
+            renormalize_probabilities: bool,
+            enable_expert_bias: bool = False
+    ):
+        """
+        Constructs the TopKRouter.
+
+        Args:
+            dim: Input feature dimensionality.
+            num_experts: Total number of experts to choose from.
+            top_k: Number of experts to select for each token.
+            renormalize_probabilities: If True, probabilities of selected experts will be renormalized to sum up to 1
+            enable_expert_bias: If True, adds a bias term to the routing scores before top-k selection. This can be
+                used for loss-free load balancing.
+        """
+
+        super().__init__()
+        self.gate = nn.Linear(dim, num_experts, bias=False)
+
+        self.expert_bias: nn.Buffer | None
+        if enable_expert_bias:
+            self.expert_bias = nn.Buffer(
+                torch.empty(num_experts, dtype=torch.float32),
+                persistent=True,
+            )
+        else:
+            self.expert_bias = None
+
+        self._num_experts = num_experts
+        self._top_k = top_k
+        self._renormalize_probabilities = renormalize_probabilities
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Calculates routing decisions for the input tokens.
+
+        Args:
+            hidden_states: Input tokens. Shape: `(num_tokens, dim)`.
+
+        Returns:
+            A tuple containing:
+
+            - Selected expert indices. Shape: `(num_tokens, top_k)`.
+            - Normalized routing weights for the selected experts. Shape: `(num_tokens, top_k)`.
+        """
+
+        # scores shape (bs*slen, num_experts)
+
+        # gate
+        scores = self.gate(hidden_states)
+
+        # and now do softmax (before top-k to be able to apply expert bias)
+        scores = F.softmax(scores, dim=-1, dtype=torch.float32)
+
+        # select top-k
+        if self.expert_bias is None:
+            scores, selected_experts_indices = torch.topk(
+                scores, k=self._top_k, dim=-1
+            )
+        else:
+            _, selected_experts_indices = torch.topk(
+                scores + self.expert_bias, k=self._top_k, dim=-1
+            )
+            scores = scores.gather(dim=-1, index=selected_experts_indices)
+
+        # re-normalize scores
+        denominator = scores.sum(dim=-1, keepdim=True) + 1e-20
+        scores = scores / denominator
+
+        return selected_experts_indices, scores
+
+    def reset_parameters(self):
+        """Resets module parameters."""
+        if self.expert_bias is not None:
+            nn.init.zeros_(self.expert_bias)
+
+        self.gate.reset_parameters()

d9d/module/block/positional/__init__.py

@@ -0,0 +1,8 @@
+"""Provides modules for positional embeddings, such as Rotary Positional Embeddings."""
+
+from .rope import RotaryEmbeddingApplicator, RotaryEmbeddingProvider
+
+__all__ = [
+    "RotaryEmbeddingApplicator",
+    "RotaryEmbeddingProvider"
+]

d9d/module/block/positional/rope.py

@@ -0,0 +1,150 @@
+import torch
+from torch import nn
+
+from d9d.module.base import ModuleLateInit
+
+
+def _prepare_rope_inverse_frequencies(
+        rope_base: int,
+        inside_dim: int
+) -> torch.Tensor:
+    """
+    Calculates inverse frequencies for RoPE calculation.
+
+    Args:
+        rope_base: Base for the geometric progression.
+        inside_dim: Dimension of the attention head (must be even).
+
+    Returns:
+        A tensor containing the inverse frequencies.
+    """
+
+    power = torch.arange(0, inside_dim, 2, dtype=torch.int64).to(dtype=torch.float) / inside_dim
+    freq = rope_base ** power
+    inv_freq = 1.0 / freq
+    return inv_freq
+
+
+def prepare_rotary_cos_sin_emb(
+        rope_base: int,
+        head_dim: int,
+        max_position_ids: int,
+        device: torch.device,
+        dtype: torch.dtype
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Precomputes rotary cosine and sine embeddings.
+
+    Args:
+        rope_base: Base frequency for calculation.
+        head_dim: Dimensionality of the attention head (E).
+        max_position_ids: Maximum sequence length supported (S).
+        device: Target device for the tensors.
+        dtype: Target data type for the tensors.
+
+    Returns:
+        A tuple containing cosine and sine tensors, both of shapes [S, E].
+    """
+
+    position_ids = torch.arange(0, max_position_ids, dtype=torch.long)
+    freqs = _prepare_rope_inverse_frequencies(rope_base, head_dim)
+
+    arguments = (freqs[:, None] @ position_ids[None, :].float()).T
+
+    emb = torch.cat((arguments, arguments), dim=-1)
+    cos = emb.cos()
+    sin = emb.sin()
+    return cos.to(device=device, dtype=dtype), sin.to(device=device, dtype=dtype)
+
+
+class RotaryEmbeddingProvider(nn.Module, ModuleLateInit):
+    """Module that manages and provides Rotary Positional Embeddings."""
+
+    def __init__(self, rope_base: int, head_dim: int, max_position_ids: int):
+        """Constructs the RotaryEmbeddingProvider."""
+
+        super().__init__()
+        self._rope_base = rope_base
+        self._head_dim = head_dim
+        self._max_position_ids = max_position_ids
+        self.cos_emb = nn.Buffer(torch.empty(max_position_ids, head_dim), persistent=False)
+        self.sin_emb = nn.Buffer(torch.empty(max_position_ids, head_dim), persistent=False)
+
+    def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Retrieves cached cosine and sine embeddings for specific positions.
+
+        Args:
+            position_ids: Tensor of position indices.
+
+        Returns:
+            A tuple of (cos, sin) tensors aligned with the input positions.
+        """
+
+        return self.cos_emb[position_ids], self.sin_emb[position_ids]
+
+    def reset_parameters(self):
+        with torch.no_grad():
+            cos, sin = prepare_rotary_cos_sin_emb(
+                rope_base=self._rope_base,
+                head_dim=self._head_dim,
+                max_position_ids=self._max_position_ids,
+                device=self.cos_emb.device,
+                dtype=self.cos_emb.dtype
+            )
+            self.cos_emb.data = cos
+            self.sin_emb.data = sin
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary_pos_emb(
+        q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    q_embed = (q * cos) + (_rotate_half(q) * sin)
+    k_embed = (k * cos) + (_rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class RotaryEmbeddingApplicator(nn.Module):
+    """Applies Rotary Positional Embeddings (RoPE) to Q and K projections."""
+
+    def __init__(self):
+        """
+        Constructs RotaryEmbeddingApplicator object.
+        """
+
+        super().__init__()
+
+    def forward(
+            self,
+            query_states: torch.Tensor,
+            key_states: torch.Tensor,
+            position_embedding_cos: torch.Tensor,
+            position_embedding_sin: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Rotates query and key states using provided cosine and sine embeddings.
+
+        Args:
+            query_states: Query tensor. Shape: `(batch, n_heads, seq_len, head_dim)`.
+            key_states: Key tensor. Shape: `(batch, n_kv_heads, seq_len, head_dim)`.
+            position_embedding_cos: Cosine values for positions.
+                Shape: `(batch, seq_len, head_dim)`.
+            position_embedding_sin: Sine values for positions.
+                Shape: `(batch, seq_len, head_dim)`.
+
+        Returns:
+            A tuple containing the rotated query and key tensors.
+        """
+
+        query_states, key_states = _apply_rotary_pos_emb(query_states, key_states,
+                                                         position_embedding_cos, position_embedding_sin)
+
+        return query_states, key_states

d9d/module/model/qwen3_moe/__init__.py

@@ -0,0 +1,16 @@
+from .decoder_layer import Qwen3MoELayer
+from .model import Qwen3MoEForCausalLM, Qwen3MoEModel
+from .params import (
+    Qwen3MoEForCausalLMParameters,
+    Qwen3MoELayerParameters,
+    Qwen3MoEParameters,
+)
+
+__all__ = [
+    "Qwen3MoEForCausalLM",
+    "Qwen3MoEForCausalLMParameters",
+    "Qwen3MoELayer",
+    "Qwen3MoELayerParameters",
+    "Qwen3MoEModel",
+    "Qwen3MoEParameters"
+]

d9d/module/model/qwen3_moe/decoder_layer.py

@@ -0,0 +1,110 @@
+import torch
+from torch import nn
+
+from d9d.module.base import ModuleLateInit
+from d9d.module.block.attention import GroupedQueryAttention
+from d9d.module.block.moe import MoELayer
+
+from .params import Qwen3MoELayerParameters
+
+
+class Qwen3MoELayer(nn.Module, ModuleLateInit):
+    """
+    Implements a single Qwen3 Mixture-of-Experts (MoE) transformer layer.
+
+    This layer consists of a Grouped Query Attention mechanism followed by an MoE
+    MLP block, with pre-RMSNorm applied before each sub-layer.
+    """
+
+    def __init__(
+            self,
+            params: Qwen3MoELayerParameters
+    ):
+        """
+        Constructs a Qwen3MoELayer object.
+
+        Args:
+            params: Configuration parameters for the layer.
+        """
+
+        super().__init__()
+
+        self.self_attn = GroupedQueryAttention(
+            hidden_size=params.hidden_size,
+            num_attention_heads=params.num_attention_heads,
+            num_key_value_heads=params.num_key_value_heads,
+            is_causal=True,
+            qk_norm_eps=params.rms_norm_eps,
+            head_dim=params.head_dim
+        )
+
+        self.mlp = MoELayer(
+            hidden_dim=params.hidden_size,
+            num_grouped_experts=params.num_experts,
+            intermediate_dim_grouped=params.intermediate_size,
+            top_k=params.experts_top_k,
+            router_renormalize_probabilities=True
+        )
+
+        self.input_layernorm = nn.RMSNorm(params.hidden_size, eps=params.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(params.hidden_size, eps=params.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor]
+    ) -> torch.Tensor:
+        """
+        Performs the forward pass of the MoE layer.
+
+        Args:
+            hidden_states: Input tensor of shape `(batch, seq_len, hidden_dim)`.
+            position_embeddings: Tuple containing RoPE precomputed embeddings (cos, sin).
+
+        Returns:
+            Output tensor after attention and MoE blocks, shape `(batch, seq_len, hidden_dim)`.
+        """
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            position_embeddings=position_embeddings,
+            attention_mask=None  # no mask for moe decoder
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+    def reset_moe_stats(self):
+        """
+        Resets statistical counters inside the MoE router (e.g., token counts per expert).
+        """
+
+        self.mlp.reset_stats()
+
+    @property
+    def moe_tokens_per_expert(self) -> torch.Tensor:
+        """
+        Returns the number of tokens routed to each expert.
+        """
+
+        return self.mlp.tokens_per_expert
+
+    def reset_parameters(self):
+        """
+        Resets module parameters.
+        """
+
+        self.self_attn.reset_parameters()
+        self.mlp.reset_parameters()
+        self.input_layernorm.reset_parameters()
+        self.post_attention_layernorm.reset_parameters()