feat : GPT-2 residual scaling weight init

- Linear: normal(0, init_std) replaces kaiming_uniform_(a=sqrt(5)) - o_proj / mlp.down: init_std = 0.02 / sqrt(2 * n_layers) - MoE: expert down scaled by 1/sqrt(1/n_shared + 1/K) - Embedding: normal(0, 0.02), unchanged
2026-06-25 15:08:31 +08:00
parent 01d2da2893
commit 88ec63121d
4 changed files with 25 additions and 8 deletions
@@ -14,6 +14,7 @@ class DecoderBlock(nn.Module):
    def __init__(self, config, layer_id: int):
        super().__init__()
        cfg = asdict(config)
+        cfg["down_init_std"] = 0.02 / (2 * config.n_layers) ** 0.5
        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)