CUDA 融合核函数

GPU（图形处理器）通常有许多计算单元，可以在快速低延迟的情况下完成多个任务。在使用 GPU 进行计算时，我们需要最大限度地利用其强大的算力，使其较少地空闲。 核函数是一组用于实现 PyTorch 操作的指令。例如，当你使用 torch.add 函数时，它将通过一个 PyTorch 调度器来决定需要执行哪些操作。 当我们使用 GPU 进行计算时，在执行一个或多个 PyTorch 操作时，GPU 会从其显存中读取数据，然后执行操作并将结果写入显存。这个过程可能会造成一些开销。如果我们能合并多个 PyTorch 操作并将它们组成一个单一的融合核函数，那么 GPU 可以更好地利用其算力，从而降低开销。 在融合核函数中，多个离散的计算和进出显存的数据移动被替换为少量数据移动的融合计算。在一些情况下，一些数学变换也被用来提高计算速度。 我们需要使用许多自定义 CUDA 融合核函数来快速高效地训练。例如，LayerNorm 的融合核函数，以及用于融合缩放、掩码和 softmax 的核函数。所有这些核函数都可以在 Megatron-LM 代码库中找到，并通过 PyTorch 的 JIT 功能与 GeLU 融合。融合这些操作能够最大限度地利用 GPU 的算力和减少内存瓶颈。因此，当执行瓶颈在内存的操作时，可以同时执行 Bias Add 操作，而不会增加运行时间。