電子發燒友App
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）
內存碎片更少，操作更少，性能/內存均有收益！
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往往跟集群規模相關，需要自適應修改，并不一定是越小越好，不然訓練性能損失慘重。
審核編輯：黃飛
工商網監
評論