题目重点：从shape和因果性检查定位问题，再用独立实现对照数值。
题意与假设
这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。
- 训练示例是单层pre-LN causal decoder，无dropout、无训练器，token/position参数显式随机初始化，multihead attention后接GELU FFN、residual和tied vocabulary head。
- debug任务需要实际starter code；本文给可验证的练习实现，不声称复原某份未知notebook。训练loss可调用独立cross entropy实现，但输入/target必须自行shift。
解题路线
- 先画输入B×T→embedding B×T×D→Q/K/V B×H×T×Dh→score B×H×T×T→attended B×T×D→logits B×T×Vocab。确认D%H=0，转置与reshape顺序不能混淆，head拼接回原顺序。
- score=QKᵀ/√Dh，未来位置用−∞后softmax；把禁止位置置0不是屏蔽，softmax(0)仍有权重。算加权值应weights@V，不应weights@y/输入标签；如果错误张量恰巧同shape，单shape断言抓不到。
- 位置参数需要初始化且在计算路径里；np.empty/未初始化Parameter可能带垃圾值。完整训练路径要算loss并backward，next-token训练logits[:, :-1]对应ids[:,1:]，不能预测同一token后报告语言模型loss。
- 按最小case调试：T1注意力必须为1；改变未来token不改前面logits；独立Torch forward核对layernorm、GELU和residual；训练另查梯度非零与少量数据能否过拟合。数值一致仅验证当前前向合同，不证明训练器或真实数据管道无误。
手动走一遍
T3的mask允许矩阵为[[1,0,0],[1,1,0],[1,1,1]]；将第二/三token改掉，位置0 logits不应变化。若D4、H2，每head Dh2，scale是√2而不是√4。attention输出与V错误分支同shape时，应检查数值参考和单元行为，不能只打印shape。
Python 参考实现
下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。
代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。
import numpy as np
def layer_norm(x):
return (x-x.mean(axis=-1,keepdims=True))/np.sqrt(x.var(axis=-1,keepdims=True)+1e-5)
def gelu(x):return .5*x*(1+np.tanh(np.sqrt(2/np.pi)*(x+.044715*x**3)))
def initialize(vocab,width,heads,max_length,seed=0):
if min(vocab,width,heads,max_length)<1 or width%heads:raise ValueError('dimensions/heads')
rng=np.random.default_rng(seed);scale=width**-.5
return dict(token=rng.normal(size=(vocab,width))*scale,position=rng.normal(size=(max_length,width))*scale,
q=rng.normal(size=(width,width))*scale,k=rng.normal(size=(width,width))*scale,
v=rng.normal(size=(width,width))*scale,o=rng.normal(size=(width,width))*scale,
ff1=rng.normal(size=(width,4*width))*scale,ff2=rng.normal(size=(4*width,width))*scale,heads=heads)
def forward(tokens,p,cache=None):
ids=np.asarray(tokens)
if ids.ndim!=2 or ids.shape[1]<1 or not np.issubdtype(ids.dtype,np.integer):raise ValueError('B,T integer tokens')
if np.any(ids<0) or np.any(ids>=len(p['token'])):raise ValueError('token ID')
b,t=ids.shape;d=p['token'].shape[1];h=p['heads'];dh=d//h
offset=0 if cache is None else cache['k'].shape[2]
if offset+t>len(p['position']):raise ValueError('position length')
x=p['token'][ids]+p['position'][offset:offset+t];normalized=layer_norm(x)
split=lambda a:a.reshape(b,t,h,dh).transpose(0,2,1,3)
q=split(normalized@p['q']);k=split(normalized@p['k']);v=split(normalized@p['v'])
if cache is not None:
if cache['k'].shape!=(b,h,offset,dh) or cache['v'].shape!=cache['k'].shape:raise ValueError('cache shape')
k=np.concatenate([cache['k'],k],axis=2);v=np.concatenate([cache['v'],v],axis=2)
scores=q@k.swapaxes(-1,-2)/np.sqrt(dh)
permitted=np.arange(k.shape[2])[None,:]<=offset+np.arange(t)[:,None]
scores=np.where(permitted[None,None,:,:],scores,-np.inf)
shifted=scores-scores.max(axis=-1,keepdims=True);weights=np.exp(shifted);weights/=weights.sum(axis=-1,keepdims=True)
attended=(weights@v).transpose(0,2,1,3).reshape(b,t,d)
hidden=x+attended@p['o'];hidden=hidden+gelu(layer_norm(hidden)@p['ff1'])@p['ff2']
logits=layer_norm(hidden)@p['token'].T
return logits,hidden,dict(k=k,v=v)
def classify(tokens,p,pooling_mask,head,bias):
# Encoder-like use is a different contract; this classifier retains causal hidden states.
_,hidden,_=forward(tokens,p);mask=np.asarray(pooling_mask,dtype=bool)
if mask.shape!=hidden.shape[:2] or np.any(mask.sum(axis=1)==0):raise ValueError('nonempty pooling mask')
# Mask must be right padding; masked suffix cannot influence earlier causal states.
if np.any(mask[:,1:] & ~mask[:,:-1]):raise ValueError('right padding')
pooled=(hidden*mask[:,:,None]).sum(axis=1)/mask.sum(axis=1,keepdims=True)
return pooled@head+bias
复杂度
长度T、宽D，一层attention O(BT²D)，线性投影/FFN O(BTD²)，score内存O(BHT²)。示例显式score，不是FlashAttention；Torch独立参考只检验small tensors。
易错点与追问
容易失分的边界
- softmax沿head维或time查询维而非key维会改变归一化。
- 忘记使用计算出的loss而直接对logits.sum backward，不是目标训练。
- starter具体第四个bug未给全时不能编造“必考四个错误”。
继续深挖
- 长序列内存不足时，block attention怎样保持正确归一化？
- dropout、autocast和不同LayerNorm epsilon加入后，参考误差预算怎样调整？
