探索提升大规模神经网络训练效率的计算优化策略，重点分析GPU架构、GEMM优化及共享内存访问等关键技术。
原文标题：原创｜训练大规模神经网络：计算（一）
原文作者：数据派THU
冷月清谈：
怜星夜思：
2、文章中提到通过共享内存(SMEM)来减少全局内存访问，对于减少访存延迟有很大的帮助。但是SMEM的大小是有限制的，那么在实际应用中，如何合理地划分block size，以及如何设计数据在SMEM中的排布，才能最大程度地提升性能？
3、文章提到了操作数重用缓存(Operand Reuse Cache)可以减少寄存器组(Register Bank)冲突。除了文中提到的寄存器排布方法，还有没有其他避免寄存器组冲突的技巧？在实际CUDA编程中，大家通常如何避免和调试这类问题？
原文内容
作者：陈陟原本文约3300字，建议阅读6分钟本系列文章将从计算和通信的角度分别讨论如何提高网络的训练效率。
随着神经网络的规模逐年递增，如何高效的训练神经网络在业内得到了日益增长的关注。
本系列文章将从计算和通信的角度分别讨论如何提高网络的训练效率。
背景
在开始之前，让我们先回顾一下两个核心概念：
1.什么是GPU？
GPU的全称为图形处理器
下图展示了英伟达GH100图形处理器的架构：
GH100图形处理器包括144个SMs（Streaming Processor, 流式处理器），分布在8个GPCs（Graphics processing cluster，图形处理集群）里。
SM是英伟达图形处理器的基本构建单元。
下图展示了一个SM当中的结构：
在现代英伟达图形处理器中，每个SM包括4个SMSPs（Streaming Processor Sub-Partitions，流式处理器子分区）。这4个SMSPs共享一个一级缓存（L1 Cache，包括一级指令缓存以及一级数据缓存）。从Volta架构开始，共享内存被合并进入一级数据缓存中，来提供高效的跨SMSPs通信。
由于篇幅所限，我们无法细致的介绍每个组件。但是理解这些基础的硬件结构很重要，因为针对硬件结构编写CUDA核心是高效代码的核心。
2.什么是GEMM？
GEMM的全称为通用矩阵乘法
几乎所有神经网络的操作都和乘法（当然还有加法）有关。比如，在卷积神经网络当中，我们常常通过im2col算法来将一个卷积运算转化成一个乘法运算。而在Transformer当中，不管是获得query、key、value的linear运算还是三者之间attention操作都是乘法运算。因此，高效的GEMM算子非常重要。
数学上，我们定义一个GEMM的操作为：
以下是一段基础的PyTorch代码来实现一个GEMM运算：
```
import torch
def naive_gemm_torch(A, B):
M, K = A.shape
K2, N = B.shape
B = B.to(A)
C = torch.zeros(M, N, dtype=A.dtype, device=A.device)
for i in range(M):
for j in range(N):
row_i = A[i, :]
col_j = B[:, j]
C[i, j] = torch.sum(row_i * col_j)
return C
这个代码的效率很低，因此不要实际去执行他。重要的是他展现了我们该如何实现一个GEMM操作。
下面，让我们实现一段等价的CUDA代码。通过并行执行矩阵乘法来提高运行速度。
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = '''
__global__ void naive_gemm_kernel(const float* A, const float* B, float* C, int M, int N, int K) {
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;
if (row < M && col < N) {
float sum = 0.0;
for (int i = 0; i < K; i++) {
sum += A[row * K + i] * B[i * N + col];
}
C[row * N + col] = sum;
}
}
torch::Tensor naive_gemm_cuda(torch::Tensor left, torch::Tensor right) {
const auto height = left.size(0);
const auto hidden = left.size(1);
const auto width = right.size(1);
auto result = torch::empty({height, width}, left.options());
dim3 threads_per_block(16, 16);
dim3 number_of_blocks((width + threads_per_block.x - 1) / threads_per_block.x,
(height + threads_per_block.y - 1) / threads_per_block.y);
naive_gemm_kernel<<
>>(
left.data_ptr
(), right.data_ptr
(), result.data_ptr
(), height, width, hidden);
return result;
}
'''
cpp_source = "torch::Tensor naive_gemm_cuda(torch::Tensor left, torch::Tensor right);"
cuda_extensions = load_inline(
name="cuda_extensions",
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=["naive_gemm_cuda"],
with_cuda=True,
extra_cuda_cflags=["-O2"],
)
随后，让我们来测试一下这段CUDA代码和PyTorch原生算子（当然，实际上则是英伟达cuBLAS算子）之间的效率差距。
可以看出来，该算法只在小规模运算上有些许的性能优势。这是因为我们没有做优化，因此这些优化产生的额外overhead对性能产生了负面影响。但当矩阵足够大时，这些优化可以带来超过十倍的性能提升。那么，如何优化呢？
3.如何提升GEMM
回顾GEMM的核心是：
```
__global__ void naive_matmul_kernel(const float* A, const float* B, float* C, int M, int N, int K) {
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;
if (row < M && col < N) {
float sum = 0.0;
for (int i = 0; i < K; i++) {
sum += A[row * K + i] * B[i * N + col];
}
C[row * N + col] = sum;
```
注意到在这个运算当中，需要执行3次访存操作：
(1)加载 A
(2)加载 B
(3)存储 C
因此，我们的CGMA（Compute to Global Memory Access, 计算-全局内存访问）比仅仅只有1/3。这是完全不可接受的。为此提出了分布式共享内存的访问架构。
下图展示了一个GH100 GPU的内存结构：
共享内存的访问大概花费32 cycles（时钟周期）。
Hopper架构Distributed Shared Memory每个时钟周期可以传输128 bytes的数据。
全局内存访问则需要188 cycles如果二级缓存命中，296 cycles如果TLB（Translation Lookaside Buffer，翻译后备缓冲）命中，616 cycles如果TLB丢失。
也就是说，我们可以节省150+ cycles如果我们能够将数据提前预读到SMEM当中。
4. 通过分区计算来充分利用SM
回顾在矩阵乘法的数学定义是，因此。
所以，我们可以依据这个公式，将每一个块的计算分配给不同的SM来实现并行计算。如下图所示：
图源：https://siboehm.com，原始地址：https://siboehm.com/assets/img/MMM/Basic_MMM.png
由于篇幅所限，本文无法详细展开具体的分块流程--实际上要复杂很多，因为具体的策略是出于平衡访存和计算而设计的。李少侠和Pzzzzz写过很好的文章来讲述如何进行分区计算，包括双缓冲加载，此处我们不再赘述。需要注意的是，采用多少线程进行计算和GPU核心高度相关。在过去几代当中，每一代内存带宽都有大幅提升。因此，需要通过实验来确定最佳的线程/块值。
通过操作数重用缓存（Operand Reuse Cache）来减少寄存器组（Register Bank）冲突。
众所周知，SMEM其实分成了很多个大小为4 bytes的组。其实寄存器也包括了4个组。和SMEM（以及其他存储结构一样），一次只能访问组内的一个寄存器。因此同时访问会导致Stall。
为了缓解这一问题，Maxwell架构引入了操作数重用缓存。操作数重用缓存可以为每一个指令的指令槽存储8 bytes的数据。当下一条指令的同一个指令槽要访问同一个寄存器中的数据时，可以直接从操作数重用缓存中读取来避免缓存访问。
下图展示了一个寄存器排布的方法，红、蓝、黄、绿分别对应了一个寄存器组。通过正确的排布寄存器，我们可以消除C寄存器和阻塞寄存器的所有组冲突。
图源：https://github.com/NervanaSystems/maxas/wiki/SGEMM，原始地址：https://github.com/xldrx/maxas.wiki.fixed/raw/master/img/RegisterBanks.png
An Improved MAGMA GEMM for Fermi GPUs
Performance Upper Bound Analysis and Optimization of SGEMM on Fermi and Kepler GPUs
[NVIDIA H100 Tensor Core GPU Architecture Overview](https://resources.nvidia.com/en-us-tensor-core)
[TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings](https://arxiv.org/abs/2304.01433)
[Dissecting the NVIDIA Volta GPU Architecture via Microbenchmarking](https://arxiv.org/abs/1804.06826)
[Dissecting the NVIDIA Turing T4 GPU via Microbenchmarking](https://arxiv.org/abs/1903.07486)
[The Ultra-Scale Playbook: Training LLMs on GPU Clusters](https://huggingface.co/spaces/nanotron/ultrascale-playbook)
[CUDA Refresher](https://developer.nvidia.com/blog/tag/cuda-refresher/)
[CUDA C++ Programming Guide](https://docs.nvidia.com/cuda/cuda-c-programming-guide/)
[GPU MODE](https://www.youtube.com/@GPUMODE)
[86.41 / 66.17 - Sistemas Digitales - University of Buenos Aires](https://campusgrado.fi.uba.ar/course/view.php?id=383)
[Accelerating Matrix Multiplication with Block Sparse Format and NVIDIA Tensor Cores](https://developer.nvidia.com/blog/accelerating-matrix-multiplication-with-block-sparse-format-and-nvidia-tensor-cores/)
[Fast Multidimensional Matrix Multiplication on CPU from Scratch](https://siboehm.com/articles/22/Fast-MMM-on-CPU)
作者简介
陈陟原，目前在香港大学多媒体实验室（HKU-MMLab）就读硕士研究生，师从刘希慧教授。他的研究兴趣是AI4Bio，集中于分子生物学。
数据派研究部介绍
数据派研究部成立于2017年初，以兴趣为核心划分多个组别，各组既遵循研究部整体的知识分享和实践项目规划，又各具特色：
算法模型组：积极组队参加kaggle等比赛，原创手把手教系列文章；
调研分析组：通过专访等方式调研大数据的应用，探索数据产品之美；
系统平台组：追踪大数据&人工智能系统平台技术前沿，对话专家；
自然语言处理组：重于实践，积极参加比赛及策划各类文本分析项目；
制造业大数据组：秉工业强国之梦，产学研政结合，挖掘数据价值；
数据可视化组：将信息与艺术融合，探索数据之美，学用可视化讲故事；
网络爬虫组：爬取网络信息，配合其他各组开发创意项目。
点击文末“阅读原文”，报名数据派研究部志愿者，总有一组适合你~
