原文:Performance Tuning Guide — PyTorch Tutorials 1.10.0+cu102 documentation
1. DataLoad: num_workers开启多进程读取;pin_memory允许异步H2D拷贝。
2. torch.no_grad():不需要backward的地方,用这个可以不创建反向图,不缓存激活值;
3. 如果BatchNorm后面接Conv,则该Conv不需要再指定bias做相加了;
4. param.grad=None取代model.zero_grad() 和 optimizer.zero_grad()
5. point-wise operator,在函数头部加@torch.compile来自动将该函数进行operator fusion,减少了每步的中间结果在显存和GPU之间的传输延迟;利用了PyTorch2的TorchInductor这个compiler:
@torch.compile
def gelu(x):
return x * 0.5 * (1.0 + torch.erf(x / 1.41421))
6. Chanels Last: 图像的NCHW格式(torch.continuous_format),改成NHWC格式(torch.channels_last),可以充分利用TensorCore的FP16计算,提升22%的速度;
7. torch.utils.checkpoint可以让一些forward function的输出激活值不进行保留,backward时当场计算;时间换空间;(叫做Gradient Checkpointing,实际是activation的checkpointing,在backward时最多再跑一整遍forward激活计算)
8. 关掉一些影响性能的操作:
a. torch.autograd.detect_anomaly可以把发生异常或者nan时,该backward所对应的forward op名称显示出来;
b. profiler
c. torch.autograd.gradcheck:用来检查梯度是否算的正确(通过对input的微小调整和梯度值的变化,来检测)
9. CPU加速:
a. NUMA:尽量让进程不要访问其他socket上的内存,只访问自己的;DDP多进程跑时,每个进程占用1个CPU。
numactl --cpunodebind=0 --membind=0 python train.py
b. 配置OpenMP的并行度、CPU affinity(把线程绑到固定的CPU上能减少cache miss)等参数;
c. Intel的libiomp在Intel芯片上的表现比GNU的libgomp,性能要好;
d. Memory allocator: Jemalloc优化了内存碎片和并发访问;TCMalloc使得频繁被访问的小对象可以保持在cache里;
e. oneDNN;和torch.jit.trace以及amp的autocast配合使用,可以对计算图进行融合;
f. torch-ccl使用了Intel的oneCCL,加速单机多socket CPU之间的collective communication;
10. GPU加速:
a. 打开cuDNN的auto tuner,一上来先尝试多种不同的kernel,后面就用尝试的几种kernel里最快的那个来跑。"torch.backends.cudnn.benchmark=True"
b. 尽量避免CPU-GPU之间的同步。例如:print某GPU上的tensor, tensor.item(), H2D或者D2H拷贝,tensor.nonzero(),拿到tensor的值用来影响if语句结果;
c. 直接在GPU上创建tensor,而不是copy过去:
慢:torch.rand(size).cuda()
快:torch.rand(size, device='cuda')
d. 使用AMP(Automatic Mixed Precision)来利用TensorCores做FP16计算; (input tensor size要对齐、padding到8的倍数)
11. 预先分配足够大的内存/显存,以后每个batch就可以复用之 ,避免出现遇到longer sequence的batch时再次开辟;(用户自己一上来用longest的input来跑1轮forward&backward,不要让梯度更新模型,之后再跑真实数据,就OK)
12. DDP多进程,比DP多线程,避免了GIL锁,DDP有更优性能;特别是Intel的oneCCL通信库,torch-ccl,CPU之间组播通信加速库;
13. 多次batch叠加到一起,进行1次权重更新;
14. DDP默认会用模型正向声明的顺序,来做backward和all-reduce相overlap的buckets填充;也可关闭一个开关来让backward时遇到的op来作为填充bucket的顺序;(DDP已经实现了AllReduce和partial梯度计算的overlap,即AllReduce传输上一个bucket的参数时,当前GPU正在计算下一个bucket的参数)
15. Load balance的解决:Transformer类别的模型,可以让各个进程的batch里的token总数保持相似,batch里的句子个数可以不一样!其他模型,可以按序列长度分桶,每轮所有进程都拿相等长度的桶;或者把数据集按长度排好序,依顺序去拿;