神经网络如何映射到GPU？

Question

我想了解当GPU执行神经网络时，操作如何映射到GPU的硬件资源。我熟悉 GPU（尤其是 NVIDIA）的架构，并且通常知道它们如何执行神经网络，但我不知道如何对硬件资源的操作进行详细和细粒度的调度以及内核如何执行他们。我想知道是否有任何工具或一组工具可以实现这一点。

更具体地说，假设我在 pytorch 中有一个预先训练的神经网络，并希望在 NVIDIA 3090 GPU 上运行它。如何通过 SM 或线程将操作（无论是在 MAC 操作还是在 NN 的神经元/通道/层）详细调度到相应的硬件资源？

Answer 1

这个问题非常广泛，但我可以提供一些信息。

当您加载神经网络时，您正在加载权重张量。这些权重通常加载到 CPU 上，然后传递到 GPU 内存 (HBM)。

除了权重之外，您还有模型逻辑（即 pytorch 模型的

forward

方法）。请注意，模型逻辑与权重本身是分开的。

模型逻辑决定何时执行什么权重。

假设我们有模型：

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()

        self.layer1 = nn.Linear(32, 8)
        self.layer2 = nn.Linear(8, 1)

    def forward(self, x):
        x = self.layer1(x)
        x = torch.relu(x)
        x = self.layer2(x)
        return x

模型状态字典中的权重是

layer1

和

layer2

的权重/偏差张量。我们的模型执行逻辑是

layer1/relu/layer2

方法中的

forward

代码。

当我们对模型进行推理时，

forward

方法决定操作的顺序。

每一层都有对应的GPU内核。内核决定如何将输入权重/激活分解为网格/块并在 GPU SM 之间分配。

通常 GPU 一次执行一层，并为该层使用尽可能多的计算。

使用上面的模型，它看起来像这样：

1. 给定一些输入

   x

   和 HBM 中的模型权重
2. 将

   x

   和

   layer1

   的权重移至 SRAM
3. 使用

   torch.nn.Linear

   和

   x

   权重

   执行 

   layer1
的 GPU 内核
4. 将结果返回至HBM
5. 将

   x

   （现在是

   layer1

   的结果）移入SRAM
6. 在

   torch.relu

    上执行 

   x
的 GPU 内核
7. 将结果移回 HBM
8. 将

   x

   （现在是relu操作的结果）和

   layer2

   的权重移动到SRAM
9. 使用

   torch.nn.Linear

   和

   x

   权重

   执行 

   layer2
的 GPU 内核
10. 将结果移回 HBM

对于上述情况，每个内核执行都会将输入/权重分配到不同的网格/块来执行。具体的分布取决于内核本身的逻辑。

这会因为内核融合（即在 relu 上）等的可能性而变得更加复杂，但这是基本思想。

您可以使用 pytorch profiler 来查看在前向传递过程中正在执行哪些内核。