refactor : 将 config 对象直接传给 DecoderBlock，替代 16 个独立参数

- DecoderBlock.__init__ 改为 (config, layer_id)，内部用 asdict
  展开字段给 AttnFactory/FFNFactory，factory 按 __init__ 签名自动过滤
- EncoderConfig 补充 attn_type 和 ffn_type 字段
- 314 个测试全部通过

astrai/config/model_config.py

@@ -70,10 +70,12 @@ class EncoderConfig(BaseModelConfig):
     rope_theta: Optional[float] = None
     rope_scaling: Optional[dict] = None
 
+    attn_type: str = "gqa"
     n_heads: Optional[int] = None
     n_kv_heads: Optional[int] = None
     use_qk_norm: Optional[bool] = None
     use_gated_attention: Optional[bool] = None
 
+    ffn_type: str = "mlp"
     pooling_type: Optional[str] = None
     normalize_embeddings: Optional[bool] = None

astrai/model/components/decoder_block.py

@@ -1,3 +1,4 @@
+from dataclasses import asdict
 from typing import Optional
 
 import torch.nn as nn
@@ -10,35 +11,13 @@ from astrai.model.components.norm import RMSNorm
 
 
 class DecoderBlock(nn.Module):
-    def __init__(
+    def __init__(self, config, layer_id: int):
-        self,
-        dim: int,
-        n_heads: int,
-        dim_ffn: int,
-        n_kv_heads: int,
-        norm_eps: float,
-        use_qk_norm: bool,
-        use_gated_attention: bool,
-        layer_id: int,
-        attn_type: str = "gqa",
-        ffn_type: str = "mlp",
-        **kwargs,
-    ):
         super().__init__()
-        self.attention = AttnFactory.create(
+        cfg = asdict(config)
-            attn_type,
+        self.attention = AttnFactory.create(config.attn_type, **cfg, layer_id=layer_id)
-            dim=dim,
+        self.input_norm = RMSNorm(config.dim, config.norm_eps)
-            n_heads=n_heads,
+        self.post_attention_norm = RMSNorm(config.dim, config.norm_eps)
-            n_kv_heads=n_kv_heads,
+        self.mlp = FFNFactory.create(config.ffn_type, **cfg)
-            use_qk_norm=use_qk_norm,
-            norm_eps=norm_eps,
-            use_gated_attention=use_gated_attention,
-            layer_id=layer_id,
-            **kwargs,
-        )
-        self.input_norm = RMSNorm(dim, norm_eps)
-        self.post_attention_norm = RMSNorm(dim, norm_eps)
-        self.mlp = FFNFactory.create(ffn_type, dim, dim_ffn, **kwargs)
 
     def forward(
         self,

astrai/model/encoder.py

@@ -26,19 +26,7 @@ class EmbeddingEncoder(AutoModel):
         self.embed_tokens = Embedding(config.vocab_size, config.dim)
 
         self.layers = nn.ModuleList(
-            [
+            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
-                DecoderBlock(
-                    config.dim,
-                    config.n_heads,
-                    config.dim_ffn,
-                    config.n_kv_heads,
-                    config.norm_eps,
-                    config.use_qk_norm,
-                    config.use_gated_attention,
-                    layer_id,
-                )
-                for layer_id in range(config.n_layers)
-            ]
         )
 
         self.norm = RMSNorm(config.dim, config.norm_eps)

astrai/model/transformer.py

@@ -62,28 +62,7 @@ class AutoRegressiveLM(AutoModel):
         self.embed_tokens = Embedding(config.vocab_size, config.dim)
 
         self.layers = nn.ModuleList(
-            [
+            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
-                DecoderBlock(
-                    config.dim,
-                    config.n_heads,
-                    config.dim_ffn,
-                    config.n_kv_heads,
-                    config.norm_eps,
-                    config.use_qk_norm,
-                    config.use_gated_attention,
-                    layer_id,
-                    attn_type=config.attn_type,
-                    ffn_type=config.ffn_type,
-                    n_routed_experts=config.n_routed_experts,
-                    n_shared_experts=config.n_shared_experts,
-                    n_activated_experts=config.n_activated_experts,
-                    topk_method=config.topk_method,
-                    kv_lora_rank=config.kv_lora_rank,
-                    qk_nope_head_dim=config.qk_nope_head_dim,
-                    qk_rope_head_dim=config.qk_rope_head_dim,
-                )
-                for layer_id in range(config.n_layers)
-            ]
         )
 
         self.norm = RMSNorm(config.dim, config.norm_eps)