PyTorch中如何选择float32矩阵乘法的内部精度？核心问题解析

PyTorch 1.12版本将默认fp32数学精度改为「highest precision」，同时引入了torch.set_float32_matmul_precision API，允许用户指定medium、high、highest三者之一作为float32矩阵乘法的内部计算精度。官方文档提到，选择较低精度「可显著提升性能，且部分程序中精度损失的影响可忽略不计」。

1. 如何判断自身程序是否能从低精度设置中获益？是否完全依赖经验判断？ #

不需要完全依赖经验，可通过以下方法判断：

• 看核心计算类型：如果程序包含大量矩阵乘法操作（比如深度学习中的全连接层、注意力机制层），低精度设置几乎肯定能带来性能提升——因为Tensor Core对低精度计算的硬件优化非常显著。
• 对比测试验证：分别用默认的highest和目标低精度（high或medium）跑完整的训练/推理流程，记录两个关键指标：
  • 性能指标：比如每秒迭代次数、总训练时长、推理延迟
  • 任务指标：比如模型准确率、loss值、回归任务的误差等
    如果性能提升明显，且任务指标的下降在业务可接受范围内，就适合用低精度。
• 结合任务场景：对精度敏感度较低的任务（如普通图像分类、文本分类）更容易从低精度中获益；而对精度要求极高的任务（如医疗影像高精度分割、科学计算中的数值模拟）则需要谨慎测试，避免精度损失影响结果有效性。
• 检查数值稳定性：如果切换低精度后，训练过程中出现loss突然变为NaN、指标剧烈波动等情况，说明低精度导致了数值不稳定，这种场景就不适合使用低精度设置。

2. 如何在high和medium这两种低精度间选择？是否完全依赖经验？有何推荐方法？ #

两者核心区别在于内部计算精度：

• high模式下，矩阵乘法内部使用TensorFloat32（tf32），精度接近标准fp32，但计算速度大幅提升
• medium模式下，内部使用fp16，性能提升比high更显著，但精度损失也更大

选择时不用全靠经验，推荐按以下步骤来：

1. 优先尝试high模式：它是平衡精度与性能的最优选择，tf32的精度损失极小，绝大多数深度学习任务（分类、检测、NLP等）都能兼容，同时能充分利用Tensor Core的硬件加速能力。
2. 在性能需求更高时考虑medium：如果high带来的性能提升仍无法满足需求（比如大规模批量训练、低延迟推理场景），且任务对精度的容忍度较高，可以尝试medium。但必须做严格的对比测试，确保任务指标的下降在可接受范围内。
3. 关注数值稳定性：使用medium时，要额外注意训练过程中的数值问题，比如是否需要调整学习率、添加梯度裁剪、使用混合精度训练的配套策略（如动态损失缩放），避免出现NaN或模型不收敛的情况。
4. 结合硬件特性：对于支持tf32的GPU（如A100、RTX 30/40系列），high模式的性能提升已经很明显；而对于较早的Tensor Core GPU（如V100），medium可能是更直接的性能提升选项，但同样要结合任务精度要求判断。

内容的提问来源于stack exchange，提问作者thesofakillers