docs/core/block.md

core/block.py · Block

代码

class Block(nn.Module):
    def __init__(self, layer_id, config: LMConfig):
        self.self_attn = Attention(config)
        self.input_layernorm = RMSNorm(hidden_size, eps)
        self.post_attention_layernorm = RMSNorm(hidden_size, eps)
        self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)

    def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
        residual = hidden_states
        h, present = self.self_attn(self.input_layernorm(hidden_states), position_embeddings, past_key_value, use_cache, attention_mask)
        h = h + residual
        h = h + self.mlp(self.post_attention_layernorm(h))
        return h, present

结构

Pre-Norm Transformer 块（LLaMA 风格）：

x ─┬─► RMSNorm ─► Attention ─────────────┐
  │                                       + ─► x
  └───────────────────────────────────────┘
x ─┬─► RMSNorm ─► MLP ───────────────────┐
  │                                       + ─► 输出
  └───────────────────────────────────────┘

• 残差直接相加，归一化在子层之前（Pre-Norm），深层训练更稳定。
• mlp 依据 config.use_moe 在稠密 / 稀疏间切换，对上层透明。
• present 返回当前层 KV（用于 cache）。

要点（面试）

• Pre-Norm vs Post-Norm：Pre-Norm 残差路径无归一化，梯度更易回传，现代 LLM 主流。
• 为什么是 input_layernorm / post_attention_layernorm 两个？分别归一化 attention 与 mlp 的输入。
• MoE 在 Block 层切换，意味着同一层要么全稠密、要么全 MoE，没有混合层。