From 6d98bb4f9fca1f7e05c202d2033167d88082e56c Mon Sep 17 00:00:00 2001 From: Gaolingx <947770192@qq.com> Date: Sat, 1 Aug 2026 22:48:58 +0800 Subject: [PATCH 1/2] 20260801-moe model impl need to add aux loss for load balancing --- astrai/config/model_config.py | 10 ++ astrai/model/__init__.py | 3 +- astrai/model/components/__init__.py | 3 +- astrai/model/components/decoder_block.py | 15 ++- astrai/model/components/mlp.py | 16 ++- tests/module/test_forward_configs.py | 141 +++++++++++++++++++++++ 6 files changed, 182 insertions(+), 6 deletions(-) diff --git a/astrai/config/model_config.py b/astrai/config/model_config.py index 0aefffe..911ebd1 100644 --- a/astrai/config/model_config.py +++ b/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: diff --git a/astrai/model/__init__.py b/astrai/model/__init__.py index 5bdba1c..22b2283 100644 --- a/astrai/model/__init__.py +++ b/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 diff --git a/astrai/model/components/__init__.py b/astrai/model/components/__init__.py index 96e0cda..6205674 100644 --- a/astrai/model/components/__init__.py +++ b/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", diff --git a/astrai/model/components/decoder_block.py b/astrai/model/components/decoder_block.py index 2c1d880..69b4785 100644 --- a/astrai/model/components/decoder_block.py +++ b/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, diff --git a/astrai/model/components/mlp.py b/astrai/model/components/mlp.py index 08a42ab..95f7d8f 100644 --- a/astrai/model/components/mlp.py +++ b/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): diff --git a/tests/module/test_forward_configs.py b/tests/module/test_forward_configs.py index d659a82..a6f20d4 100644 --- a/tests/module/test_forward_configs.py +++ b/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 From d7db37a70fd6d00bd1969d0425709201b11ee054 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sun, 2 Aug 2026 05:30:26 +0800 Subject: [PATCH 2/2] fix: preserve MoE routing defaults --- astrai/config/model_config.py | 10 +++++++-- astrai/model/components/mlp.py | 12 +++++++--- tests/module/test_forward_configs.py | 33 ++++++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 5 deletions(-) diff --git a/astrai/config/model_config.py b/astrai/config/model_config.py index 911ebd1..b982769 100644 --- a/astrai/config/model_config.py +++ b/astrai/config/model_config.py @@ -65,7 +65,7 @@ class AutoRegressiveLMConfig(BaseModelConfig): 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. + norm_topk_prob (bool): Normalize top-k routing probabilities. Defaults to True. 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. """ @@ -94,7 +94,7 @@ class AutoRegressiveLMConfig(BaseModelConfig): topk_method: Optional[str] = None moe_intermediate_size: Optional[int] = None shared_expert_intermediate_size: Optional[int] = None - norm_topk_prob: bool = False + norm_topk_prob: bool = True decoder_sparse_step: int = 1 mlp_only_layers: Optional[list[int]] = None @@ -112,6 +112,12 @@ class AutoRegressiveLMConfig(BaseModelConfig): raise ValueError(f"ffn_type must be one of {sorted(_FFN_TYPES)}, got {v!r}") return v + @field_validator("decoder_sparse_step") + def _validate_decoder_sparse_step(cls, v: int) -> int: + if v < 1: + raise ValueError(f"decoder_sparse_step must be at least 1, got {v}") + return v + @dataclass @ConfigFactory.register("embedding") diff --git a/astrai/model/components/mlp.py b/astrai/model/components/mlp.py index 95f7d8f..e294fb8 100644 --- a/astrai/model/components/mlp.py +++ b/astrai/model/components/mlp.py @@ -40,7 +40,7 @@ class DeepSeekMoE(nn.Module): n_layers: int = 1, moe_intermediate_size: Optional[int] = None, shared_expert_intermediate_size: Optional[int] = None, - norm_topk_prob: bool = False, + norm_topk_prob: bool = True, ): super().__init__() self.dim = dim @@ -50,8 +50,14 @@ class DeepSeekMoE(nn.Module): 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 + 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 diff --git a/tests/module/test_forward_configs.py b/tests/module/test_forward_configs.py index a6f20d4..1213c48 100644 --- a/tests/module/test_forward_configs.py +++ b/tests/module/test_forward_configs.py @@ -246,3 +246,36 @@ def test_moe_custom_intermediate_shape(): assert expert.up.weight.shape[0] == 20 assert expert.gate.weight.shape[0] == 20 assert expert.down.weight.shape[1] == 20 + + +def test_moe_defaults_preserve_normalized_routing(): + from astrai.config.model_config import AutoRegressiveLMConfig + + config = AutoRegressiveLMConfig( + **TINY_CONFIG, + ffn_type="moe", + n_routed_experts=4, + n_shared_experts=1, + n_activated_experts=2, + topk_method="greedy", + ) + model = AutoRegressiveLM(config) + + assert config.norm_topk_prob is True + assert model.layers[0].mlp.norm_topk_prob is True + + +@pytest.mark.parametrize("decoder_sparse_step", [0, -1]) +def test_moe_rejects_invalid_decoder_sparse_step(decoder_sparse_step): + from pydantic import ValidationError + + from astrai.config.model_config import AutoRegressiveLMConfig + + with pytest.raises(ValidationError, match="decoder_sparse_step must be at least 1"): + AutoRegressiveLMConfig( + **TINY_CONFIG, + ffn_type="moe", + n_routed_experts=4, + n_activated_experts=2, + decoder_sparse_step=decoder_sparse_step, + )