From ba10c4d17ba41cfed4bab5cbea7b6561c1d1cc7e Mon Sep 17 00:00:00 2001 From: Linxiushen <1787979356@qq.com> Date: Sun, 20 Sep 2026 11:50:28 +0800 Subject: [PATCH] =?UTF-8?q?[fix]=20float16=20=E4=B8=8B=20attention=20mask?= =?UTF-8?q?=20=E7=9A=84=20-1e9=20=E6=BA=A2=E5=87=BA=E6=88=90=20-inf?= =?UTF-8?q?=EF=BC=8C=E5=AF=BC=E8=87=B4=E6=95=B4=E6=9D=A1=E5=BA=8F=E5=88=97?= =?UTF-8?q?=E8=BE=93=E5=87=BA=20NaN?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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）。 --- model/model_minimind.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/model/model_minimind.py b/model/model_minimind.py index 0f0c3ef..1b34cee 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -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)): 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: - 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 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 = self.resid_dropout(self.o_proj(output)) return output, past_kv