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）

　?。?/p>

　　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實現(xiàn)使用的Index-select算子，Index-select會涉及索引展開、內(nèi)存復(fù)用等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數(shù)據(jù)

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

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

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

　　。.. # 省略若干計算代碼

　　comm_handle.wait（）

　　對應(yīng)中間的計算就能夠跟通信進行overlap，只要我們提前梳理好網(wǎng)絡(luò)拓撲，完全是沒問題的。

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

　　PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

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

　　6. 在適當時機清空緩存可以大幅降低內(nèi)存占用

　　torch.cuda.empty_cache（）

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

　　訓(xùn)練過程中，禁止使用torch.cuda.empty_cache（），除非切換不同任務(wù)（如train/eval切換），因為cache blocks釋放會觸發(fā)Stream Synchronize，開銷較大。

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

　　data = data.cuda（non_blocking=True）

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

　　8. 在CPU負載比較空的時候，還是要充分利用的

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

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

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

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

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

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

　　11. 對于連續(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）

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

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

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

　　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）

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

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

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

　　審核編輯：黃飛