Qwen模型推理时attention_mask首位为0导致输出NaN问题分析

问题现象

在使用Qwen-7B-Chat模型进行推理时，当输入的attention_mask张量中第一个位置被置为0（即序列起始位置被mask掉），模型forward函数的输出logits会全部变为NaN值。这种现象在以下场景中尤为明显：

1. 进行批量推理时（batch_size > 1）
2. 使用不同长度的输入序列进行左填充（left-padding）处理时
3. 未启用flash-attn的情况下

技术背景

attention_mask在Transformer模型中用于控制注意力机制的计算范围，通常：

• 1表示该位置参与注意力计算
• 0表示该位置被mask掉，不参与计算

在批量处理不同长度序列时，通常会对较短序列进行填充（padding），并使用attention_mask来屏蔽这些填充位置的影响。

问题分析

经过深入分析，这个问题与以下几个因素相关：

1. 模型架构差异：Qwen1.0系列模型存在此问题，而Qwen1.5及后续版本已修复
2. PyTorch版本影响：在PyTorch 2.0.1版本下表现正常，但在PyTorch 2.4及以上版本会出现问题
3. 注意力实现方式：使用flash-attn可以规避此问题，但原生实现会出现NaN

根本原因可能是：

• 当序列起始位置被mask时，模型的自注意力计算可能产生了数值不稳定
• 在特定PyTorch版本下，矩阵运算的实现方式变化导致了数值溢出的问题

解决方案

针对此问题，推荐以下几种解决方案：

1. 升级模型版本：迁移到Qwen1.5或Qwen2.0等后续版本，这些版本已经修复了此问题

2. 启用flash-attn：

# 安装flash-attn
pip install flash-attn

# 在代码中启用
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-7B-Chat", use_flash_attention_2=True)

3. 调整PyTorch版本：降级到PyTorch 2.0.1版本可以临时解决问题

4. 修改padding策略：避免使用左填充，或确保序列起始位置不被mask

最佳实践建议

对于生产环境中的使用，建议：

1. 优先使用最新的Qwen2.0系列模型
2. 保持PyTorch等依赖库的版本更新
3. 在批量推理时，尽量保持输入序列长度一致，减少padding需求
4. 对于必须使用padding的场景，建议启用flash-attn以获得更好的性能和稳定性

总结

这个问题展示了深度学习模型在实际应用中的一些潜在陷阱，特别是在处理变长序列和注意力掩码时。理解模型的行为、保持框架和模型的更新，以及合理使用优化技术（如flash-attn），都是确保模型稳定运行的重要因素。