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

- DecoderBlock.__init__ 改为 (config, layer_id),内部用 asdict
  展开字段给 AttnFactory/FFNFactory,factory 按 __init__ 签名自动过滤
- EncoderConfig 补充 attn_type 和 ffn_type 字段
- 314 个测试全部通过
This commit is contained in:
2026-06-19 14:15:33 +08:00
parent 7348bac6ab
commit b1adc40cfb
4 changed files with 11 additions and 63 deletions
+7 -28
View File
@@ -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,