RTX 3080 在 pytorch 中使用 CUDA 11.3 时速度不会比 GTX 1650 Ti 快很多吗？

Question

所以我在不同的机器上有两个 GPU：

• GTX 1650 Ti：4 GB，1024 个 CUDA 核心，在笔记本电脑中
• RTX 3080：10 GB，8960 个 CUDA 核心，台式机

我正在使用它们来训练 PyTorch 深度学习模型。对于这两者，我都设置了一个带有 CUDA 11.3 工具包的 conda PyTorch 环境。最初，我预计 RTX 会比 GTX 提供多倍的加速。然而，我很快注意到 RTX 的性能似乎“仅”是 GTX 的两倍。我真的不确定在这里应该期待多少，但考虑到它是 RTX 而不是板载笔记本显卡，只有两倍的加速会有点令人失望 ^^。根据 gpu.userbenchmark.com（以及一般情况），RTX 应该与 GTX 完全不同。请参阅此摘录比较（其中左边是 GTX，右边是 RTX）： https://i.stack.imgur.com/XXaag.png

根据 pytorch 文档中的示例，这里有一些简短的基准代码：

import torch
print(torch.__version__)
print(torch.version.cuda)

def batched_dot_mul_sum(a, b):
    '''Computes batched dot by multiplying and summing'''
    return a.mul(b).sum(-1)

def batched_dot_bmm(a, b):
    '''Computes batched dot by reducing to bmm'''
    a = a.reshape(-1, 1, a.shape[-1])
    b = b.reshape(-1, b.shape[-1], 1)
    return torch.bmm(a, b).flatten(-3)

# Input for benchmarking
x = torch.randn(10000, 64, device='cuda')

# Ensure that both functions compute the same output
assert batched_dot_mul_sum(x, x).allclose(batched_dot_bmm(x, x))

import torch.utils.benchmark as benchmark

t0 = benchmark.Timer(
    stmt='batched_dot_mul_sum(x, x)',
    setup='from __main__ import batched_dot_mul_sum',
    globals={'x': x})

t1 = benchmark.Timer(
    stmt='batched_dot_bmm(x, x)',
    setup='from __main__ import batched_dot_bmm',
    globals={'x': x})

print(t0.timeit(100000))
print(t1.timeit(100000))

输出如下：

• GTX：

1.10.2
11.3
<torch.utils.benchmark.utils.common.Measurement object at 0x000001F7518E8C40>
batched_dot_mul_sum(x, x)
setup: from __main__ import batched_dot_mul_sum
  72.80 us
  1 measurement, 100000 runs , 1 thread
<torch.utils.benchmark.utils.common.Measurement object at 0x000001F7518E8D90>
batched_dot_bmm(x, x)
setup: from __main__ import batched_dot_bmm
  70.01 us
  1 measurement, 100000 runs , 1 thread

• RTX：

1.10.2
11.3
<torch.utils.benchmark.utils.common.Measurement object at 0x000001AF2F63F5E0>
batched_dot_mul_sum(x, x)
setup: from __main__ import batched_dot_mul_sum
  29.88 us
  1 measurement, 100000 runs , 1 thread
<torch.utils.benchmark.utils.common.Measurement object at 0x000001AF26F99B20>
batched_dot_bmm(x, x)
setup: from __main__ import batched_dot_bmm
  27.02 us
  1 measurement, 100000 runs , 1 thread

我们可以看到，使用 RTX 时，速度只提高了两倍多一点，这似乎并不多（？）。因此，我可能做错了什么，导致 RTX 的性能没有得到充分利用。

RTX 附带安装了 CUDA 11.6。最新的 conda PyTorch 发行版只支持 11.3，但可以使用夜间二进制文件安装 pytorch CUDA 11.6 版本：

pip install torch --pre --extra-index-URL https://download.pytorch.org/whl/nightly/cu116

使用该环境，基准测试脚本（带有 Cuda 11.6 的 RTX）的相应输出为：

1.13.0.dev20220528+cu116
11.6
<torch.utils.benchmark.utils.common.Measurement object at 0x0000011A24BC6100>
batched_dot_mul_sum(x, x)
setup: from __main__ import batched_dot_mul_sum
  22.96 us
  1 measurement, 100000 runs , 1 thread
<torch.utils.benchmark.utils.common.Measurement object at 0x0000011A15803B50>
batched_dot_bmm(x, x)
setup: from __main__ import batched_dot_bmm
  21.88 us
  1 measurement, 100000 runs , 1 thread

因此，根据 CUDA 版本，额外加速效果很小。总体而言，RTX 相对于 GTX 的优势似乎不大，尤其是考虑到其中一个是板载笔记本电脑显卡，另一个是实际设备。

另一个选择是从源代码构建带有 cuda 11.6 的 PyTorch。也许只有这样，RTX 的性能才能得到充分利用？到目前为止，我一直避免使用此选项，因为可以使用夜间二进制文件选项。从源代码构建会有所不同吗？也许这就是问题的原因。否则，与 GTX 相比，我真的不知道 RTX 性能缓慢的原因是什么 :(。

谢谢，JZ