PyTorch高效编程实战指南

　　1. 能用_all_gather_base的，不用all_gather

　　output = torch.empty（input.numel（） * world_size， dtype=input.dtype， device=input.device）

　　torch.distributed._all_gather_base（output， input， group=xxx）

　　vs.

　　output_list = ［

　　torch.empty（input.numel（）， dtype=input.dtype， device=input.device）

　　for _ in range（world_size）

　　］

　　torch.distributed.all_gather（output_list， input， group=xxx）

　　output = torch.cat（output_list， dim=0）

　　内存碎片更少，操作更少，性能/内存均有收益！

　　2. 能用专有算子的，不用通用算子

　　如 F.embedding vs. Index-select

　　Megatron-LM master实现使用的Index-select算子，Index-select会涉及索引展开、内存复用等HostCPU逻辑，效率较低

　　3. 对于生命周期较长的Tensors，可以共用contiguous buffer

　　data = torch.zeros（global_size， dtype=xx， device=xx）

　　start_idx = 0

　　for i in range（len（item_list））：

　　item_list［i］ = data［start_idx:start_idx+item_list［i］.numel（）］.view（item_list［i］.shape）

　　torch.cuda.empty_cache（） # 清空原始已释放的item list数据

　　CUDA内存池是对齐分配的，使用分散的block会带来内存碎片，同时对于相同操作，可以直接对contiguous buffer进行操作，减少了更多的算子下发，大块计算效率也会更高。

　　4. 尽可能使用异步通信，提高计算/通信overlap

　　comm_handle = torch.distributed.all_reduce（data， group=xxx， async_op=True）

　　。.. # 省略若干计算代码

　　comm_handle.wait（）

　　对应中间的计算就能够跟通信进行overlap，只要我们提前梳理好网络拓扑，完全是没问题的。

　　5. 对于输入数据size频繁变化的场景，使用Expandable Segments

　　PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

　　跟cudaMalloc直接分配Kernel可访问的内存地址不同，该机制操作的是虚拟内存空间（对应的物理内存地址不具备访问权限），可以通过驱动map更多的物理内存在已分配的block的后面，从而使得segments可向上扩展，一定程度上提高了cache match的效率，减少内存碎片。

　　6. 在适当时机清空缓存可以大幅降低内存占用

　　torch.cuda.empty_cache（）

　　在训练任务初始化时，经常会创建一些临时的设备Tensors，如果在训练任务开始时不及时清理，会造成内存池碎片化，最终导致内存占用增加。

　　训练过程中，禁止使用torch.cuda.empty_cache（），除非切换不同任务（如train/eval切换），因为cache blocks释放会触发Stream Synchronize，开销较大。

　　7. non-blocking H2D拷贝是安全的，可以无脑使用

　　data = data.cuda（non_blocking=True）

　　在后续对当前数据有依赖的地方会主动插入sync point，保证数据安全；在没有立即对数据产生依赖的场景，可以使得数据H2D拷贝和计算并行。

　　8. 在CPU负载比较空的时候，还是要充分利用的

　　如数据加载的时候可以尽量将部分操作放在CPU负载。当前Megatron master主干在这一块还是很有优化空间的。

　　https://zhuanlan.zhihu.com/p/670569490

　　但是尽量不要在网络中间插入to cpu操作，会触发同步，反而弄巧成拙。

　　9. 加速通信算子内存释放，可以无脑使用

　　10. 训练/推理过程中不要触及内存上限

　　如果内存观测是在持续上下跳动，那就是触及了内存上限，虽然整体程序能正常run起来，这时候已经频繁触发了内存池回收，每一次block回收都会触发一次Stream Synchronize，虽然平均利用率看起来可能超过90%，但是整体性能会降低的非常多。

　　11. 对于连续的ElementWise算子，可以使用NvFuser加速

　　@torch.jit.script

　　def bias_dropout_add（x_with_bias， residual， prob， training）：

　　x， bias = x_with_bias # unpack

　　x = x + bias

　　out = torch.nn.functional.dropout（x， p=prob， training=training）

　　out = residual + out

　　return out

　　torch._C._jit_set_nvfuser_enabled（True）

　　前反向过程可以通过NvFuser实时生成高效的融合Kernel，但是注意torch.jit.script装饰器下的所有操作必须能被TorchScript语法解释，不然还是不能work的（具体可以去看PyTorch官方文档的TorchScript语法介绍）。

　　12. 模型运行过程中不要流同步阻塞算子下发

　　

　　D2H操作、内存回收、以及主动调用流同步（torch.cuda.synchronize（））等都会阻塞算子下发（保证对应Stream清空），那么后续算子如果执行过快（比下发快），那就会造成GPU间隙，所以说这个下发越快越好、越多越好，上图这个曲线是越缓越好，下发即执行那就是性能随时都可能坑。

　　13. 尽量使用TensorCore，避免使用CUDACore

　　# 直接使用cumsum

　　b = a.cumsum（dim=-1）

　　# 使用矩阵计算替代

　　a = torch.matmul（a.view（x， b， s）， triu_matrix）

　　c = a［：， ：-1， -1］.cumsum（-1）

　　a［：， 1：， ：］.add_（c.unsqueeze（-1））

　　a = a.view（x， b*s）

　　上图的计算如果替换成矩阵计算，加速数十倍，在cumsum维度过高的情况下，开销是异常大的。所以在遇到类似场景，都尽量转换成矩阵计算，即使计算量增加很多，速度还是有巨大收益的。

　　14. 集群通信需要寻找合适的bucket size

　　对于分桶通信，最优bucket size往往跟集群规模相关，需要自适应修改，并不一定是越小越好，不然训练性能损失惨重。

　　审核编辑：黄飞