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往往跟集群規模相關，需要自適應修改，并不一定是越小越好，不然訓練性能損失慘重。

　　審核編輯：黃飛