关于torch.nn.functional.linear在输入为float16或bfloat16时计算类型的疑问

嘿，我完全理解翻PyTorch源码找这种底层细节有多头疼，我来帮你搞清楚这个问题！

你的问题核心是：当输入张量都是torch.float16时，linear操作的计算精度是FP16还是FP32？答案其实取决于你运行代码的设备（CPU/CUDA），具体如下：

1. 在支持FP16的CUDA GPU上运行（比如Turing、Ampere架构及以上） #

当你把a和b（这里b作为linear的权重参数）都定义为torch.float16且放在CUDA设备上时，整个计算过程会直接在FP16精度下执行。

这是因为现代GPU有专门的Tensor Core硬件来加速FP16的矩阵乘法（GEMM）操作，PyTorch会默认利用这个硬件特性，计算全程保持FP16精度，输出结果的 dtype 也是torch.float16。

2. 在CPU上运行 #

大多数CPU没有原生的FP16计算单元，所以PyTorch会先把FP16的输入张量转换为FP32精度来执行矩阵乘法运算，计算完成后再把结果转换回FP16输出。

也就是说，虽然最终输出是FP16，但中间的核心计算过程是用FP32完成的。

如何验证？ #

你可以通过简单的代码来快速确认：

import torch
from torch.nn.functional import linear

# 自动选择可用设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
a = torch.ones(2, 3, dtype=torch.float16, device=device)
b = torch.ones(2, 3, dtype=torch.float16, device=device)

output = linear(a, b)
print(f"输出 dtype: {output.dtype}")  # 无论CPU/CUDA，输出都是torch.float16

# 用性能分析器看计算细节（仅CUDA有效）
if device.type == "cuda":
    with torch.autograd.profiler.profile(use_cuda=True) as prof:
        linear(a, b)
    # 查看GEMM操作的精度，会显示为half（即FP16）
    print(prof.key_averages().table(sort_by="cuda_time_total"))

确实PyTorch的底层实现分散在C++和CUDA代码里，上层Python接口看不到这些细节，所以通过实际测试或者官方文档的精度指南来确认会更高效～

备注：内容来源于stack exchange，提问作者YSF