[fix] float16 下 attention mask 的 -1e9 溢出成 -inf,导致整条序列输出 NaN

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)。
pull/863/head
Linxiushen 2026-09-20 11:50:28 +08:00
parent cc312c1cc6
commit ba10c4d17b
1 changed files with 6 additions and 2 deletions

View File

@ -125,10 +125,14 @@ class Attention(nn.Module):
if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)): if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)):
output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal) output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal)
else: else:
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim) # 先转 float32 再加 mask:-1e9 超出 float16 的表示范围(min≈-65504),
# 在 float16 张量上累加会饱和成 -inf。若某个 query 在因果范围内可见的
# key 全被 mask 掉(左 padding 时必然出现),softmax 整行就是 NaN,
# 并在下一层通过 K/V 投影扩散到同序列里有效的 token 上。
scores = (xq @ xk.transpose(-2, -1)).float() / math.sqrt(self.head_dim)
if self.is_causal: scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1) if self.is_causal: scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9 if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv output = self.attn_dropout(F.softmax(scores, dim=-1).type_as(xq)) @ xv
output = output.transpose(1, 2).reshape(bsz, seq_len, -1) output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
output = self.resid_dropout(self.o_proj(output)) output = self.resid_dropout(self.o_proj(output))
return output, past_kv return output, past_kv