Attention.forward 的手写 softmax 分支(flash SDPA 关闭时走这里):
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
...
if attention_mask is not None: scores += (1.0 - mask) * -1e9
output = ... F.softmax(scores.float(), dim=-1) ...
scores 的 dtype 跟随 xq/xk。float16 模型下它是 float16,而
torch.finfo(torch.float16).min ≈ -65504,把 -1e9 累加进去会饱和成 -inf
(不是一个很大的有限负数)。当某个 query 在因果范围内可见的 key 全被 mask
掉时(左 padding 必然出现),softmax 整行 = NaN,并在下一层通过 K/V 投影
扩散到同序列里 mask=1 的有效 token 上。
.float() 加在 softmax 那一刻已经太晚——溢出发生在它之前。
修复:把 float32 提升提前到加 mask 之前,两处共 2 行。
fp32/bf16 路径数值等价(fp32 下原来的 .float() 本就是 no-op)。
|
||
|---|---|---|
| .. | ||
| __init__.py | ||
| model_lora.py | ||
| model_minimind.py | ||
| tokenizer.json | ||
| tokenizer_config.json | ||