提升大规模神经网络训练效率:计算优化策略解析

探索提升大规模神经网络训练效率的计算优化策略,重点分析GPU架构、GEMM优化及共享内存访问等关键技术。

原文标题:原创|训练大规模神经网络:计算(一)

原文作者:数据派THU

冷月清谈:

本文从计算的角度探讨如何提高大规模神经网络的训练效率。首先回顾了GPU的架构,特别是SM(流式处理器)的内部结构,强调理解硬件结构对于编写高效CUDA代码的重要性。然后讨论了GEMM(通用矩阵乘法)在神经网络中的核心作用,并展示了基础的PyTorch和CUDA代码实现。文章着重分析了GEMM的优化方法,包括通过分布式共享内存的访问架构来减少全局内存访问,以及通过分区计算来充分利用SM。此外,还提到了操作数重用缓存如何减少寄存器组冲突。总的来说,文章旨在为读者提供关于如何通过计算优化来加速神经网络训练的实用指导。

怜星夜思:

1、文章提到了GEMM在神经网络中的重要性,但实际应用中,针对不同规模的矩阵,选择合适的GEMM实现方案(比如cuBLAS、MKL等)会有差异。大家在实际项目中是如何根据矩阵规模和硬件环境选择的?有没有一些通用的benchmark工具或者经验分享?
2、文章中提到通过共享内存(SMEM)来减少全局内存访问,对于减少访存延迟有很大的帮助。但是SMEM的大小是有限制的,那么在实际应用中,如何合理地划分block size,以及如何设计数据在SMEM中的排布,才能最大程度地提升性能?
3、文章提到了操作数重用缓存(Operand Reuse Cache)可以减少寄存器组(Register Bank)冲突。除了文中提到的寄存器排布方法,还有没有其他避免寄存器组冲突的技巧?在实际CUDA编程中,大家通常如何避免和调试这类问题?

原文内容

作者:陈陟原
本文约3300字,建议阅读6分钟
本系列文章将从计算和通信的角度分别讨论如何提高网络的训练效率。


随着神经网络的规模逐年递增,如何高效的训练神经网络在业内得到了日益增长的关注。


本系列文章将从计算和通信的角度分别讨论如何提高网络的训练效率。


背景


在开始之前,让我们先回顾一下两个核心概念:


1.什么是GPU


  GPU的全称为图形处理器


下图展示了英伟达GH100图形处理器的架构:



GH100图形处理器包括144SMsStreaming Processor, 流式处理器),分布在8GPCsGraphics processing cluster,图形处理集群)里。


SM是英伟达图形处理器的基本构建单元。


下图展示了一个SM当中的结构:



在现代英伟达图形处理器中,每个SM包括4SMSPsStreaming Processor Sub-Partitions,流式处理器子分区)。这4SMSPs共享一个一级缓存(L1 Cache,包括一级指令缓存以及一级数据缓存)。从Volta架构开始,共享内存被合并进入一级数据缓存中,来提供高效的跨SMSPs通信。


由于篇幅所限,我们无法细致的介绍每个组件。但是理解这些基础的硬件结构很重要,因为针对硬件结构编写CUDA核心是高效代码的核心。


2.什么是GEMM

 

GEMM的全称为通用矩阵乘法


几乎所有神经网络的操作都和乘法(当然还有加法)有关。比如,在卷积神经网络当中,我们常常通过im2col算法来将一个卷积运算转化成一个乘法运算。而在Transformer当中,不管是获得querykeyvaluelinear运算还是三者之间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


因此,我们的CGMACompute to Global Memory Access, 计算-全局内存访问)比仅仅只有1/3。这是完全不可接受的。为此提出了分布式共享内存的访问架构。


 下图展示了一个GH100 GPU的内存结构:




共享内存的访问大概花费32 cycles(时钟周期)。


Hopper架构Distributed Shared Memory每个时钟周期可以传输128 bytes的数据。


全局内存访问则需要188 cycles如果二级缓存命中,296 cycles如果TLBTranslation 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等比赛,原创手把手教系列文章;

调研分析组:通过专访等方式调研大数据的应用,探索数据产品之美;

系统平台组:追踪大数据&人工智能系统平台技术前沿,对话专家;

自然语言处理组:重于实践,积极参加比赛及策划各类文本分析项目;

制造业大数据组:秉工业强国之梦,产学研政结合,挖掘数据价值;

数据可视化组:将信息与艺术融合,探索数据之美,学用可视化讲故事;

网络爬虫组:爬取网络信息,配合其他各组开发创意项目。


点击文末“阅读原文”,报名数据派研究部志愿者,总有一组适合你~



转载须知


如需转载,请在开篇显著位置注明作者和出处(转自:数据派THUID:DatapiTHU),并在文章结尾放置数据派醒目二维码。有原创标识文章,请发送【文章名称-待授权公众号名称及ID】至联系邮箱,申请白名单授权并按要求编辑。

未经许可的转载以及改编者,我们将依法追究其法律责任。






关于我们

数据派THU作为数据科学类公众号,背靠清华大学大数据研究中心,分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识,努力建设数据人才聚集平台、打造中国大数据最强集团军。




新浪微博:@数据派THU

微信视频号:数据派THU

今日头条:数据派THU

点击“阅读原文”加入组织~



SMEM的合理使用确实是个难题。我的经验是,block size的确定要综合考虑SMEM的大小、寄存器的使用情况、以及warp的调度效率。数据在SMEM中的排布,要尽量避免bank conflict,这通常需要仔细设计内存访问模式。可以使用一些工具,比如NVIDIA Visual Profiler,来分析SMEM的利用率和bank conflict的情况,然后不断调整block size和数据排布。

楼上说的benchmark方法很扎实!我补充一点,在选择GEMM库时,要考虑的因素还有精度问题。有些加速库为了追求速度可能会牺牲一些精度,在对精度要求高的场景下需要谨慎选择。另外,一些深度学习框架(比如TensorFlow、PyTorch)底层已经对GEMM进行了优化,可以直接利用框架提供的接口,通常也能获得不错的性能。

SMEM这玩意儿,调好了是神技,调不好就是鸡肋。我曾经遇到过一个case,为了用SMEM,把block size设的特别大,结果导致warp调度效率很低,反而性能下降了。所以,还是要具体问题具体分析,不能一概而论。

选择恐惧症表示:benchmark一时爽,选起来火葬场。其实我觉得可以先用框架默认的GEMM实现,然后profile一下看看是不是性能瓶颈。如果是,再考虑换其他库。毕竟优化是无止境的,还是要抓住主要矛盾。

楼上说的NVIDIA Visual Profiler是神器!我再补充一点,SMEM的使用也要和计算强度相匹配。如果计算量太小,SMEM带来的收益可能还不如访存带来的开销。所以,要尽量提高计算强度,才能充分发挥SMEM的作用。此外,双缓冲(double buffering)也是一个常用的技巧,可以隐藏访存延迟。

针对GEMM方案的选择,我通常先用一些benchmark工具跑一下,比如Google Benchmark,分别测试不同矩阵规模下cuBLAS、MKL以及一些开源的GEMM库(比如OpenBLAS)的性能。然后根据实际硬件环境(CPU、GPU型号)和矩阵规模的分布情况,选择一个综合性能最优的方案。有时候小矩阵用MKL,大矩阵用cuBLAS也是一种策略。

寄存器这块儿,水太深了… 我一般都是先保证代码逻辑正确,然后profile一下看看是不是有性能瓶颈。如果是,再考虑优化寄存器使用。感觉这是一个trade-off,优化寄存器会增加代码复杂度,降低可读性,需要谨慎对待。

我补充一个比较tricky的方法:padding。有时候,可以通过在数组中插入一些Padding元素,来改变内存访问模式,从而避免bank conflict和寄存器组冲突。当然,这种方法需要在存储空间和性能之间进行权衡。

避免寄存器组冲突,除了合理的寄存器排布,还可以尝试展开循环、使用局部变量等方法来减少寄存器的使用。此外,还可以使用一些编译器优化选项,比如-Xptxas -dlcm=cg,来让编译器自动进行寄存器分配优化。调试这类问题,可以使用NVIDIA Compute Sanitizer,它能够检测到寄存器组冲突,并给出详细的报告。