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）

　　內(nèi)存碎片更少，操作更少，性能/內(nèi)存均有收益！

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

　　如 F.embedding vs. Index-select

　　Megatron-LM master實(shí)現(xiàn)使用的Index-select算子，Index-select會(huì)涉及索引展開(kāi)、內(nèi)存復(fù)用等HostCPU邏輯，效率較低

　　3. 對(duì)于生命周期較長(zhǎng)的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數(shù)據(jù)

　　CUDA內(nèi)存池是對(duì)齊分配的，使用分散的block會(huì)帶來(lái)內(nèi)存碎片，同時(shí)對(duì)于相同操作，可以直接對(duì)contiguous buffer進(jìn)行操作，減少了更多的算子下發(fā)，大塊計(jì)算效率也會(huì)更高。

　　4. 盡可能使用異步通信，提高計(jì)算/通信overlap

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

　　。.. # 省略若干計(jì)算代碼

　　comm_handle.wait（）

　　對(duì)應(yīng)中間的計(jì)算就能夠跟通信進(jìn)行overlap，只要我們提前梳理好網(wǎng)絡(luò)拓?fù)?，完全是沒(méi)問(wèn)題的。

　　5. 對(duì)于輸入數(shù)據(jù)size頻繁變化的場(chǎng)景，使用Expandable Segments

　　PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

　　跟cudaMalloc直接分配Kernel可訪問(wèn)的內(nèi)存地址不同，該機(jī)制操作的是虛擬內(nèi)存空間（對(duì)應(yīng)的物理內(nèi)存地址不具備訪問(wèn)權(quán)限），可以通過(guò)驅(qū)動(dòng)map更多的物理內(nèi)存在已分配的block的后面，從而使得segments可向上擴(kuò)展，一定程度上提高了cache match的效率，減少內(nèi)存碎片。

　　6. 在適當(dāng)時(shí)機(jī)清空緩存可以大幅降低內(nèi)存占用

　　torch.cuda.empty_cache（）

　　在訓(xùn)練任務(wù)初始化時(shí)，經(jīng)常會(huì)創(chuàng)建一些臨時(shí)的設(shè)備Tensors，如果在訓(xùn)練任務(wù)開(kāi)始時(shí)不及時(shí)清理，會(huì)造成內(nèi)存池碎片化，最終導(dǎo)致內(nèi)存占用增加。

　　訓(xùn)練過(guò)程中，禁止使用torch.cuda.empty_cache（），除非切換不同任務(wù)（如train/eval切換），因?yàn)閏ache blocks釋放會(huì)觸發(fā)Stream Synchronize，開(kāi)銷較大。

　　7. non-blocking H2D拷貝是安全的，可以無(wú)腦使用

　　data = data.cuda（non_blocking=True）

　　在后續(xù)對(duì)當(dāng)前數(shù)據(jù)有依賴的地方會(huì)主動(dòng)插入sync point，保證數(shù)據(jù)安全；在沒(méi)有立即對(duì)數(shù)據(jù)產(chǎn)生依賴的場(chǎng)景，可以使得數(shù)據(jù)H2D拷貝和計(jì)算并行。

　　8. 在CPU負(fù)載比較空的時(shí)候，還是要充分利用的

　　如數(shù)據(jù)加載的時(shí)候可以盡量將部分操作放在CPU負(fù)載。當(dāng)前Megatron master主干在這一塊還是很有優(yōu)化空間的。

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

　　但是盡量不要在網(wǎng)絡(luò)中間插入to cpu操作，會(huì)觸發(fā)同步，反而弄巧成拙。

　　9. 加速通信算子內(nèi)存釋放，可以無(wú)腦使用

　　10. 訓(xùn)練/推理過(guò)程中不要觸及內(nèi)存上限

　　如果內(nèi)存觀測(cè)是在持續(xù)上下跳動(dòng)，那就是觸及了內(nèi)存上限，雖然整體程序能正常run起來(lái)，這時(shí)候已經(jīng)頻繁觸發(fā)了內(nèi)存池回收，每一次block回收都會(huì)觸發(fā)一次Stream Synchronize，雖然平均利用率看起來(lái)可能超過(guò)90%，但是整體性能會(huì)降低的非常多。

　　11. 對(duì)于連續(xù)的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）

　　前反向過(guò)程可以通過(guò)NvFuser實(shí)時(shí)生成高效的融合Kernel，但是注意torch.jit.script裝飾器下的所有操作必須能被TorchScript語(yǔ)法解釋，不然還是不能work的（具體可以去看PyTorch官方文檔的TorchScript語(yǔ)法介紹）。

　　12. 模型運(yùn)行過(guò)程中不要流同步阻塞算子下發(fā)

　　D2H操作、內(nèi)存回收、以及主動(dòng)調(diào)用流同步（torch.cuda.synchronize（））等都會(huì)阻塞算子下發(fā)（保證對(duì)應(yīng)Stream清空），那么后續(xù)算子如果執(zhí)行過(guò)快（比下發(fā)快），那就會(huì)造成GPU間隙，所以說(shuō)這個(gè)下發(fā)越快越好、越多越好，上圖這個(gè)曲線是越緩越好，下發(fā)即執(zhí)行那就是性能隨時(shí)都可能坑。

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

　　# 直接使用cumsum

　　b = a.cumsum（dim=-1）

　　# 使用矩陣計(jì)算替代

　　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）

　　上圖的計(jì)算如果替換成矩陣計(jì)算，加速數(shù)十倍，在cumsum維度過(guò)高的情況下，開(kāi)銷是異常大的。所以在遇到類似場(chǎng)景，都盡量轉(zhuǎn)換成矩陣計(jì)算，即使計(jì)算量增加很多，速度還是有巨大收益的。

　　14. 集群通信需要尋找合適的bucket size

　　對(duì)于分桶通信，最優(yōu)bucket size往往跟集群規(guī)模相關(guān)，需要自適應(yīng)修改，并不一定是越小越好，不然訓(xùn)練性能損失慘重。

　　審核編輯：黃飛