From b1adc40cfb53c2373fa708d1773fa228a85f43fc Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Fri, 19 Jun 2026 14:15:20 +0800
Subject: [PATCH] =?UTF-8?q?refactor=20:=20=E5=B0=86=20config=20=E5=AF=B9?=
 =?UTF-8?q?=E8=B1=A1=E7=9B=B4=E6=8E=A5=E4=BC=A0=E7=BB=99=20DecoderBlock?=
 =?UTF-8?q?=EF=BC=8C=E6=9B=BF=E4=BB=A3=2016=20=E4=B8=AA=E7=8B=AC=E7=AB=8B?=
 =?UTF-8?q?=E5=8F=82=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- DecoderBlock.__init__ 改为 (config, layer_id)，内部用 asdict
  展开字段给 AttnFactory/FFNFactory，factory 按 __init__ 签名自动过滤
- EncoderConfig 补充 attn_type 和 ffn_type 字段
- 314 个测试全部通过
---
 astrai/config/model_config.py            |  2 ++
 astrai/model/components/decoder_block.py | 35 +++++-------------------
 astrai/model/encoder.py                  | 14 +---------
 astrai/model/transformer.py              | 23 +---------------
 4 files changed, 11 insertions(+), 63 deletions(-)

diff --git a/astrai/config/model_config.py b/astrai/config/model_config.py
index ceb12db..d0c1cf2 100644
--- a/astrai/config/model_config.py
+++ b/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
diff --git a/astrai/model/components/decoder_block.py b/astrai/model/components/decoder_block.py
index 60b263f..d5c71aa 100644
--- a/astrai/model/components/decoder_block.py
+++ b/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__(
-        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,
-    ):
+    def __init__(self, config, layer_id: int):
         super().__init__()
-        self.attention = AttnFactory.create(
-            attn_type,
-            dim=dim,
-            n_heads=n_heads,
-            n_kv_heads=n_kv_heads,
-            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)
+        cfg = asdict(config)
+        self.attention = AttnFactory.create(config.attn_type, **cfg, layer_id=layer_id)
+        self.input_norm = RMSNorm(config.dim, config.norm_eps)
+        self.post_attention_norm = RMSNorm(config.dim, config.norm_eps)
+        self.mlp = FFNFactory.create(config.ffn_type, **cfg)
 
     def forward(
         self,
diff --git a/astrai/model/encoder.py b/astrai/model/encoder.py
index 2957be7..1ea1b2c 100644
--- a/astrai/model/encoder.py
+++ b/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.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)
-            ]
+            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
         )
 
         self.norm = RMSNorm(config.dim, config.norm_eps)
diff --git a/astrai/model/transformer.py b/astrai/model/transformer.py
index 9a2a7e2..44655f8 100644
--- a/astrai/model/transformer.py
+++ b/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.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)
-            ]
+            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
         )
 
         self.norm = RMSNorm(config.dim, config.norm_eps)