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

Transformer Block

KVCache · Decoder Layer · 完整 Decoder Block

拼装收官 a5/src/modeling/transformer.py
定位

把零件装成汉堡

  • 上层/下层「面包」：Embedding · LM Head
  • 中间「牛肉饼」：多层 Decoder Layer
  • 每层：Self-Attn + MLP + Norm + 残差
  • 推理：KVCache 避免重复算历史 K/V
A5 是综合题：前面模块接口要对、seed 表要严。
依赖

你会用到哪些旧零件

A2 Norm+ A2 Emb+ A2 RoPE
A3 MLP+ A4 Attn→ A5 Layer/Block
分仓！把通过测的实现同步进 a5，不要留 stub。
任务地图

三个 Task

T1 KVCache

存 / 取 / 追加历史 K,V

T2 DecoderLayer

一层完整前向 + cache

T3 DecoderBlock

Emb + N 层 + LMHead

附加 API

参数量 / 显存脚印统计

推理两阶段

Prefill 与 Decode

Prefilling

整段 query 一次前向
写出首 token 分布
缓存各层 K,V

Decoding

每次只进 1 个新 token
复用 cache 的 K,V
append 再算 attention

Task1

KVCache API 心智

__init__(qkv_layout, num_layers) has(layer_idx) get(layer_idx) → k, v[, cu_seqlens] set(layer_idx, k, v, cu_seqlens=None) append(layer_idx, k, v, cu_seqlens=None) reset() # THD/varlen 时 cu_seqlens 必带 # append：沿序列维拼接；层不存在则同 set
Task1 · 注意

实现自由度 vs 一致性

  • 无强制时空复杂度 → 结构可简单
  • 仍应检查 dtype/device/layout 一致性
  • 真实系统还有 PagedAttention 等优化（了解即可）
Task2

Decoder Layer 数据流

X [b,s,h] (或 THD + cu_seqlens) # Self-Attn 子层（示意 Pre-Norm） X1 = AttnNorm(X) QKV = proj(X1) + RoPE(Q,K) Oa = Attn(Q,K,V; kv_cache?) X = X + Oa # residual # MLP 子层 X2 = MLPNorm(X) Om = MLP(X2) # Dense or Sparse X = X + Om → 输出 O 与输入 X 同 shape
Task2 · Cache

层里怎么用 KVCache

  • 从 cache get 本层历史 K,V
  • 与当前 token 的 K,V 拼接后再 Attn
  • append/set 写回本层
  • mask 由 Attn 模块按 bottom-right 处理
RoPE 在 decode 时位置不再从 0 起——可用 offset（作业可不强制）。
Task2 · 组件选择

配置决定用哪套实现

  • Norm = GroupRMSNorm
  • RoPE = NTKAwareRoPE（或封装）
  • Attn = Offline 或 Online（条件满足才 Online）
  • MLP = Dense 或 Sparse（num_experts）
Online 仅当：定长 max_seq、无 cache、BSHD、Q_K_V 等条件全满足。
Task3

Decoder Block 全图

I [b,s] token ids → ParallelVocabEmbedding → X_ini → DecoderLayer × L → X_fin → FinalNorm → LMHead (nn.Linear) → Logits [b,s,vocab] 推理 (not training)：把同一 KVCache 传给每层 训练：通常不用 cache
Task3 · API

Block 还要提供

  • get / set / reset_kv_cache
  • num_parameters(learnable_only, unit)
  • num_memory_footprint(unit)
  • lm_head_tied：是否与 embedding 共享权重
Seed 表

按层偏移（必记思想）

第 i 层示例（见作业附表1）： qkv_proj : proj_init_seed + i+1 o_proj : proj_init_seed + i+2 attn_norm : init_base_seed + i+1 attn : init_base_seed + i+2 mlp_norm : init_base_seed + i+3 mlp : init_base_seed + i+4 softmax_drop : softmax_dropout_seed + i lora* : lora_*_seed + i vocab_embed : init_base_seed + 0 lm_head : proj_init_seed + 0
测例会抠可复现性——偏移写错就全挂。
Config

TransformerConfig 记什么

  • 结构：num_layers · hidden · ffh · heads · head_dim
  • 词表：vocab_size · emb init
  • Attn：window · causal · pack · layout · softmax_*
  • MLP：activation · lora_* · num_experts · moe_topk
  • 全局：param_dtype/device · 各种 seed

初始化时一个 config 打天下，按层传 layer_idx。

踩坑

A5 高频扣分点

  • 残差顺序 / PreNorm 写反
  • cache append 维错
  • 训练路径误用 cache
  • seed 偏移表抄错
  • lm_head_tied 没处理
  • 输出 device 跟 token 不一致
  • 旧模块仍是 NotImplemented
  • 参数量单位换算错
节奏

建议 3–4 天

  1. D1：KVCache 数据结构 + 单测思维
  2. D2：DecoderLayer 无 cache 跑通
  3. D3：接入 cache + RoPE offset 意识
  4. D4：Block + tied + 统计 API + 回归
你的仓库提示

本地 a5 快照印象

  • attention / pos_emb 可能已有实现痕迹
  • norm / mlp / transformer 仍可能是 stub
  • 优先：把 A2–A4 通过版拷入 a5 再拼装
拼装前先 import 每个子模块，确保不再 NotImplemented。
收官心态

先零件正确，再总装

A5 的 bug 80% 来自旧模块接口漂移
总装时一次只接一层

pytest 同步子模块
链接

A5 资料

  • Blog：A5 Transformer Block
  • 本地：~/Documents/nju-llm/a5
  • HF Llama DecoderLayer / Cache
  • 返回总目录：hub.html
1 / 1

快捷键