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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user