20260801-moe model impl

need to add aux loss for load balancing

astrai/config/model_config.py

@@ -63,6 +63,11 @@ class AutoRegressiveLMConfig(BaseModelConfig):
         n_shared_experts (Optional[int]): Number of shared experts, MoE only. Defaults to None.
         n_activated_experts (Optional[int]): Number of activated experts per token, MoE only. Defaults to None.
         topk_method (Optional[str]): Top-k routing method, MoE only. Defaults to None.
+        moe_intermediate_size (Optional[int]): Expert hidden dim, defaults to intermediate_size if None. MoE only.
+        shared_expert_intermediate_size (Optional[int]): Shared expert hidden dim, defaults to intermediate_size if None. MoE only.
+        norm_topk_prob (bool): Normalize top-k routing probabilities. Defaults to False.
+        decoder_sparse_step (int): Frequency of MoE layers, 1=every layer. Defaults to 1.
+        mlp_only_layers (Optional[list[int]]): Layer indices using dense MLP instead of MoE. Defaults to None.
     """
 
     vocab_size: Optional[int] = None
@@ -87,6 +92,11 @@ class AutoRegressiveLMConfig(BaseModelConfig):
     n_shared_experts: Optional[int] = None
     n_activated_experts: Optional[int] = None
     topk_method: Optional[str] = None
+    moe_intermediate_size: Optional[int] = None
+    shared_expert_intermediate_size: Optional[int] = None
+    norm_topk_prob: bool = False
+    decoder_sparse_step: int = 1
+    mlp_only_layers: Optional[list[int]] = None
 
     @field_validator("attn_type")
     def _validate_attn_type(cls, v: str) -> str:

astrai/model/__init__.py

@@ -9,7 +9,7 @@ from astrai.model.components.lora import (
     merge_lora,
     save_lora,
 )
-from astrai.model.components.mlp import MLP
+from astrai.model.components.mlp import MLP, DeepSeekMoE
 from astrai.model.components.norm import RMSNorm
 from astrai.model.encoder import EmbeddingEncoder
 from astrai.model.transformer import AutoRegressiveLM
@@ -19,6 +19,7 @@ __all__ = [
     "Linear",
     "RMSNorm",
     "MLP",
+    "DeepSeekMoE",
     "GQA",
     "DecoderBlock",
     # Models

astrai/model/components/__init__.py

@@ -3,7 +3,7 @@ from astrai.model.components.attention import GQA, MLA
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding
 from astrai.model.components.linear import Linear
-from astrai.model.components.mlp import MLP
+from astrai.model.components.mlp import MLP, DeepSeekMoE
 from astrai.model.components.norm import RMSNorm
 from astrai.model.components.rope import (
     RotaryEmbedding,
@@ -14,6 +14,7 @@ __all__ = [
     "Linear",
     "RMSNorm",
     "MLP",
+    "DeepSeekMoE",
     "Embedding",
     "GQA",
     "MLA",

astrai/model/components/decoder_block.py

@@ -26,7 +26,20 @@ class DecoderBlock(nn.Module):
         self.attention = AttnFactory.create(config.attn_type, **cfg, layer_id=layer_id)
         self.input_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
         self.post_attention_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.mlp = FFNFactory.create(config.ffn_type, **cfg)
+        ffn_type = self._resolve_ffn_type(config, layer_id)
+        self.mlp = FFNFactory.create(ffn_type, **cfg)
+
+    @staticmethod
+    def _resolve_ffn_type(config, layer_id: int) -> str:
+        if config.ffn_type != "moe":
+            return config.ffn_type
+        mlp_only = config.mlp_only_layers or []
+        if layer_id in mlp_only:
+            return "mlp"
+        if config.decoder_sparse_step > 1:
+            if (layer_id + 1) % config.decoder_sparse_step != 0:
+                return "mlp"
+        return "moe"
 
     def forward(
         self,

astrai/model/components/mlp.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -36,6 +38,9 @@ class DeepSeekMoE(nn.Module):
         n_activated_experts: int = 2,
         topk_method: str = "greedy",
         n_layers: int = 1,
+        moe_intermediate_size: Optional[int] = None,
+        shared_expert_intermediate_size: Optional[int] = None,
+        norm_topk_prob: bool = False,
     ):
         super().__init__()
         self.dim = dim
@@ -43,6 +48,10 @@ class DeepSeekMoE(nn.Module):
         self.n_shared_experts = n_shared_experts
         self.n_activated_experts = n_activated_experts
         self.topk_method = topk_method
+        self.norm_topk_prob = norm_topk_prob
+
+        expert_dim_ffn = moe_intermediate_size if moe_intermediate_size is not None else dim_ffn
+        shared_dim_ffn = shared_expert_intermediate_size if shared_expert_intermediate_size is not None else dim_ffn
 
         self.router = Linear(dim, n_routed_experts, bias=False)
         moe_scale = 1 / max(n_shared_experts, 1) + 1 / n_activated_experts
@@ -50,13 +59,13 @@ class DeepSeekMoE(nn.Module):
 
         self.shared_experts = nn.ModuleList(
             [
-                MLP(dim, dim_ffn, down_init_std=down_init_std)
+                MLP(dim, shared_dim_ffn, down_init_std=down_init_std)
                 for _ in range(n_shared_experts)
             ]
         )
         self.routed_experts = nn.ModuleList(
             [
-                MLP(dim, dim_ffn, down_init_std=down_init_std)
+                MLP(dim, expert_dim_ffn, down_init_std=down_init_std)
                 for _ in range(n_routed_experts)
             ]
         )
@@ -84,7 +93,8 @@ class DeepSeekMoE(nn.Module):
         router_probs = torch.softmax(router_logits.float(), dim=-1).to(x.dtype)
 
         topk_weights, topk_indices = torch.topk(router_probs, K, dim=-1)
-        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+        if self.norm_topk_prob:
+            topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
 
         output = torch.zeros(N, D, device=x.device, dtype=x.dtype)
         for expert_idx in range(self.n_routed_experts):

tests/module/test_forward_configs.py

@@ -1,6 +1,7 @@
 import pytest
 import torch
 
+from astrai.model.components.mlp import MLP, DeepSeekMoE
 from astrai.model.transformer import AutoRegressiveLM
 from tests.helpers import TINY_CONFIG
 
@@ -32,6 +33,59 @@ CONFIGS = [
         },
         id="gqa_moe",
     ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "topk_method": "greedy",
+            "mlp_only_layers": [0],
+        },
+        id="gqa_moe_dense_first",
+    ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "topk_method": "greedy",
+            "decoder_sparse_step": 2,
+        },
+        id="gqa_moe_sparse_step",
+    ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "topk_method": "greedy",
+            "norm_topk_prob": True,
+        },
+        id="gqa_moe_norm_topk",
+    ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "topk_method": "greedy",
+            "moe_intermediate_size": 24,
+            "shared_expert_intermediate_size": 20,
+        },
+        id="gqa_moe_custom_intermediate",
+    ),
     pytest.param(
         {
             **TINY_CONFIG,
@@ -105,3 +159,90 @@ def test_model_forward_with_padding(config_kwargs, device):
 
     assert output["logits"].shape == (batch_size, seq_len, config.vocab_size)
     assert not torch.isnan(output["logits"]).any()
+
+
+def test_moe_per_layer_ffn_resolution():
+    """Verify that mlp_only_layers and decoder_sparse_step resolve FFN types correctly."""
+    from astrai.config.model_config import AutoRegressiveLMConfig
+
+    # mlp_only_layers: first layer dense, rest MoE
+    config = AutoRegressiveLMConfig(
+        **{
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "mlp_only_layers": [0],
+        }
+    )
+    model = AutoRegressiveLM(config)
+    assert isinstance(model.layers[0].mlp, MLP)
+    assert not isinstance(model.layers[0].mlp, DeepSeekMoE)
+    assert isinstance(model.layers[1].mlp, DeepSeekMoE)
+
+    # decoder_sparse_step=2: every other layer is MoE
+    config2 = AutoRegressiveLMConfig(
+        **{
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "decoder_sparse_step": 2,
+        }
+    )
+    model2 = AutoRegressiveLM(config2)
+    # layer 0 (id=0): (0+1)%2=1 != 0 -> MLP
+    assert isinstance(model2.layers[0].mlp, MLP)
+    assert not isinstance(model2.layers[0].mlp, DeepSeekMoE)
+    # layer 1 (id=1): (1+1)%2=0 -> MoE
+    assert isinstance(model2.layers[1].mlp, DeepSeekMoE)
+
+    # decoder_sparse_step=1 (default): all layers MoE
+    config3 = AutoRegressiveLMConfig(
+        **{
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+        }
+    )
+    model3 = AutoRegressiveLM(config3)
+    for layer in model3.layers:
+        assert isinstance(layer.mlp, DeepSeekMoE)
+
+
+def test_moe_custom_intermediate_shape():
+    """Verify MoE uses custom intermediate sizes when specified."""
+    from astrai.config.model_config import AutoRegressiveLMConfig
+
+    config = AutoRegressiveLMConfig(
+        **{
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "moe_intermediate_size": 24,
+            "shared_expert_intermediate_size": 20,
+        }
+    )
+    model = AutoRegressiveLM(config)
+    moe_layer = model.layers[0].mlp
+    assert isinstance(moe_layer, DeepSeekMoE)
+    # routed experts use moe_intermediate_size
+    for expert in moe_layer.routed_experts:
+        assert expert.up.weight.shape[0] == 24
+        assert expert.gate.weight.shape[0] == 24
+        assert expert.down.weight.shape[1] == 24
+    # shared experts use shared_expert_intermediate_size
+    for expert in moe_layer.shared_experts:
+        assert expert.up.weight.shape[0] == 20
+        assert expert.gate.weight.shape[0] == 20
+        assert expert.down.weight.shape[1] == 20