← 全部作业
← → 翻页 · O 总览 · ? 帮助
Assignment 3
A3

Modeling MLP

Dense GLU-MLP · LoRA · Sparse MoE

当前主战场 a3/src/modeling/mlp.py
目标

这一作业你在练什么

  • 实现 Transformer 的 前馈核心（MLP / FFN）
  • 理解 GLU 风格门控结构
  • 掌握 PEFT / LoRA 的最小工程形态
  • 理解 MoE：路由 + 多 expert 稀疏计算
Blog：A3 Modeling MLP · A5 会复用本模块
任务地图

三大块，建议串行

1 · Dense

GLU-MLP + 初始化

2 · LoRA

挂在 Dense 上，rank=0 可关

3 · Sparse

MoE：gate + top-k + experts

原则

Dense 过测 → LoRA → Sparse

Dense · 结构

GLU 前向四步

X: [b, s, d] up = X @ W_up # W_up [d, ffh] gate = φ(X @ W_gate) # W_gate [d, ffh] h = up ⊙ gate # element-wise O = h @ W_down # W_down [ffh, d] # O 与 X 同 shape / dtype / device · 无 bias

φ ← activation_type：SILU / GELU / RELU / SIGMOID / BILINEAR …

Dense · Shape

权重形状（作业约定）

W_up / W_gate

[d, ffh] = [hidden, ffh]

W_down

[ffh, d] = [ffh, hidden]

配合 X @ W 的 [in, out] 存法；不是 nn.Linear 默认的 [out, in]。
Dense · Learnable

哪些参数要学

  • 基座：W_up · W_gate · W_down
  • LoRA（rank>0）：低秩 A、B
  • Sparse：各 expert 的 Dense 参数 + gate G
不是 Parameter：激活本身、top-k 逻辑、routing 规则。
Dense · 初始化

reset_parameters 怎么调

__init__: empty 出三个 Parameter self.reset_parameters() # 只调一次 reset_parameters: 内部把全部权重 init 完 # ≠ 每个 weight 对外各调一次 reset

与 nn.Linear 同一习惯：创建结构 → 一次初始化。

Dense · Init 规则

Xavier 还是 Kaiming？

Xavier normal

  • SIGMOID
  • BILINEAR

Kaiming normal

  • RELU / GELU / SILU …
  • mode = fan_in
seed：up +1 · gate +2 · down +3 · 每个矩阵单独 manual_seed
Dense · 骨架

reset 伪代码（自己填）

use_xavier = act in {SIGMOID, BILINEAR} for weight, offset in [(W_up,1),(W_gate,2),(W_down,3)]: torch.manual_seed(base + offset) if use_xavier: xavier_normal_(weight) else: kaiming_normal_(weight, mode="fan_in", ...) # 自检 std ≈ √(2/fan_in) # 注意 [in,out] layout 对 fan 的影响
你的进度

A3 已走到哪

  • ✅ args → self.xxx
  • ✅ empty 三权重 shape 正确
  • 🟡 reset：分支有了，待 seed + 三次 init_
  • ⬜ __init__ 调用 reset_parameters()
  • ⬜ forward GLU（先无 LoRA）
  • ⬜ LoRA · SparseMLP
下一刀：写完 reset → 实例化 → print(std) → 再写 forward
LoRA

低秩增量，不是重训全家桶

ΔW ≈ (α / r) · A @ B # 作业简化：整段 MLP 只挂一次 LoRA rank == 0 → 完全跳过 α is None → α = rank A,B → uniform init seed → lora_init_base +1 / +2 dropout → lora_dropout_seed 可复现
LoRA · 注意

实现时别踩的点

  • 先 Dense 过测，再开 LoRA 分支
  • rank=0 时不要创建/计算 A、B
  • 输出仍与 X 同 dtype/device
  • 作业约定可能与 HF PEFT 不同——以 blog 为准
Sparse · MoE

稀疏 = 路由 + 多小 Dense

  1. G:[d, ne] → logits → softmax → P
  2. 每 token 取 top-k，再 renormalize
  3. expert ≈ DenseMLP(ffh = ffh // ne)
  4. 只算本地 rank 分片上的 experts
  5. 未路由到本地 → 输出 0（partial sum）
Sparse · 细节

seed 与精度

  • G 的 dtype 固定 float32
  • G：normal_，seed = init_base_seed（无偏移）
  • 每个 expert 的 seed 按 全局 expert index 偏移
  • ffh 应被 ne 整除（建议 assert）
踩坑

A3 高频扣分点

  • 多矩阵共用一个 seed
  • Xavier/Kaiming 分支反了
  • fan_in 与 layout 不一致
  • 输出 dtype/device 漂了
  • LoRA rank=0 没短路
  • α=None 没当 rank
  • MoE 忘 renorm
  • 本地 expert 索引算错
工位清单

建议完成顺序

  1. reset 三权重 + 调用 reset
  2. forward GLU，无 LoRA
  3. pytest test_toy.py
  4. 加 LoRA（含 rank=0）
  5. SparseMLP
  6. push → score-feedback
Action

现在只做这一刀

写完 reset_parameters · 挂上调用 · 打印 std · 再写 forward

一次只改一个点 pytest test_toy.py
链接

A3 资料

  • Blog：LLM-Blog → A3 Modeling MLP
  • 本地：~/Documents/nju-llm/a3
  • 核心：src/modeling/mlp.py
  • 下一份：A4 Attention
1 / 1

快捷键