fix: preserve MoE routing defaults

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

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