Merge pull request #863 from Linxiushen/fix-fp16-attention-mask-nan

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

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