ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

20260801-moe model impl

Browse Source 
need to add aux loss for load balancing
This commit is contained in:
Gaolingx
2026-08-01 22:48:58 +08:00
parent 925cbedc93
commit 6d98bb4f9f
6 changed files with 182 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/config/model_config.py
+10
View File
@@ -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
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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
+14 -1
View File
@@ -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
+13 -3
View File
@@ -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
+141
View File
@@ -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