PyTorch用户侧性能优化

 原文：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里的句子个数可以不一样！其他模型，可以按序列长度分桶，每轮所有进程都拿相等长度的桶；或者把数据集按长度排好序，依顺序去拿；