File size: 1,973 Bytes
c6b1b88 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | # core/block.py · Block
## 代码
```python
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**，没有混合层。
|