運算與通訊如何重疊？Gradient Bucket、CUDA Stream與Critical Path

運算與通訊重疊，是在Backward尚未完全結束前，把已經準備好的Gradient Bucket送進All-Reduce，同時讓GPU繼續計算其他Layer。目標不是讓網路傳得更快，而是讓部分通訊離開Critical Path，減少GPU純等待時間。

重疊不一定改善效能。Compute與Collective可能同時搶用SM、HBM、PCIe、NVLink和NIC；Bucket太小會產生大量啟動成本，Bucket太大又會太晚開始傳輸。正確做法是用Timeline找出裸露通訊時間，再調整Bucket、Stream、Topology與Parallelism。

重點快讀

• Backward通常從後層往前層產生Gradient。
• DDP把多個Gradient整理成Bucket，Bucket Ready後即可All-Reduce。
• 非同步Collective會排進Communication CUDA Stream。
• 重疊成立的前提是後續Compute不依賴該Collective結果。
• Bucket太小增加Launch與Protocol開銷；太大延後通訊起點。
• 通訊Kernel可能占用SM和記憶體頻寬，反而拖慢Compute。
• 真正指標是Exposed Communication與Step Time，不是兩條Timeline看起來重疊。
• Profiler要覆蓋Warm-up、Stable Iteration與P99 Slow Rank。

沒有重疊的時間線

Forward  ───────────────
Backward ─────────────────────
AllReduce                     ───────────
Optimizer                                ─────

Step time = compute + communication + optimizer

所有Gradient算完才開始All-Reduce，整段通訊直接落在Critical Path。GPU可能仍有Collective Kernel在跑，但模型計算已停止。

理想重疊時間線

Backward bucket 0 ───── ready → AllReduce 0 ─────
Backward bucket 1      ───── ready → AllReduce 1 ─────
Backward bucket 2           ───── ready → AllReduce 2 ───

Exposed communication =
最後一個Backward結束後仍未完成的Collective

前面Bucket的通訊被後續Backward覆蓋，最後只剩無法隱藏的尾端。改善幅度取決於可覆蓋Compute長度、通訊速度和資源競爭。

Gradient Bucket如何運作？

PyTorch DistributedDataParallel會把Parameter Gradient按Bucket Size整理。某個Bucket中所有Gradient Ready後，Reducer便能啟動Collective，不必等待完整Backward。

PyTorch DDP的bucket_cap_mb預設通常為25 MiB；這只是通用起點。模型層大小、Network與GPU不同時，最佳Bucket會改變。

ddp_model = DistributedDataParallel(
    model,
    device_ids=[rank],
    bucket_cap_mb=25,
    gradient_as_bucket_view=True,
)

DDP可能在初始Iteration後重建Bucket，Communication Hook不應永久依賴第一輪Bucket Index或順序。

CUDA Stream如何提供非同步？

Compute Kernel通常排進Default或Compute Stream；NCCL Collective可排進獨立Communication Stream。不同Stream具備獨立依賴時，GPU有機會同時執行。

• Gradient Ready Event讓Comm Stream知道資料可讀。
• Optimizer使用Gradient前必須等待All-Reduce完成。
• 多Process Group需要所有Rank保持一致Collective順序。
• 非同步API回傳Work Handle，不代表資料已可使用。
• 跨Stream要用Event、Wait或框架提供的同步。

錯誤同步可能造成Race、Deadlock或讀取未完成資料。刪除所有wait()不等於更高效，只是把依賴變成未定義行為。

為什麼重疊後反而變慢？

• Collective Kernel占用大量SM，壓縮Backward Kernel。
• Compute與Comm同時讀寫HBM，造成Bandwidth Contention。
• GPU到NIC共用PCIe路徑。
• 小Bucket使Protocol與Launch占比過高。
• 多個Process Group的Collective競爭。
• 某個Slow Rank讓所有Rank在尾端等待。
• CPU Scheduler或Python發起不及。
• Torch Compile Graph Break改變Bucket與Overlap。

時間線重疊只是視覺現象。若兩個Kernel同時執行後各自變慢，總Step Time可能不降反升。

Critical Path怎麼算？

exposed_comm =
max(0, communication_finish - backward_finish)

hidden_comm =
total_communication_time - exposed_comm

step_time =
forward + backward_with_contention
+ exposed_comm + optimizer + bubbles

除了Exposed Comm，也要比較重疊前後的Backward Duration。若通訊被隱藏50毫秒，Backward卻因競爭增加60毫秒，整體仍然退化。

Communication Hook能做什麼？

PyTorch DDP Communication Hook允許團隊接管每個GradBucket的通訊策略，例如標準All-Reduce、FP16／BF16壓縮、PowerSGD或自訂Future Pipeline。

• Hook接收GradBucket和State。
• 回傳Future，在完成後更新Bucket。
• 壓縮可能降低Network，也可能引入品質誤差。
• 自訂All-Gather通常比All-Reduce昂貴。
• 不同Rank必須執行一致邏輯。
• Hook和Optimizer Overlap需要額外Warm-up與驗證。

通訊Hook屬實驗與底層能力，模型Loss、Convergence和Checkpoint相容性都要重跑。

FSDP與Prefetch

Fully Sharded Data Parallel還會在Forward前All-Gather下一個Unit的Parameter，Backward時Reduce-Scatter Gradient。Forward Prefetch和Backward Prefetch都嘗試讓通訊與相鄰Module Compute重疊。

• FSDP Unit太小會增加Collective次數。
• 太大會提高Peak Memory和等待。
• Prefetch過多可能同時占用記憶體。
• Activation Checkpoint會改變Compute窗口。
• HSDP會改變節點內與跨節點流量。

Profiler要看什麼？

1. 排除初始化、JIT與第一輪Bucket Rebuild。
2. 擷取多個穩定Iteration。
3. 標記Forward、Backward、Bucket與Optimizer。
4. 查看NCCL Kernel和CUDA Stream。
5. 找出Compute空洞與Collective尾端。
6. 比較每個Rank的Timeline。
7. 查看GPU SM、HBM、NVLink與NIC。
8. 對照Step Time、Tokens／s和Loss。

Nsight Systems適合跨CPU、CUDA與NCCL時間線；PyTorch Profiler適合Operator、Module與Distributed事件。兩者應配合使用。

調整順序

1. 先建立不重疊或框架預設基線。
2. 確認Topology、Rank Mapping和Network正常。
3. 測小、中、大Bucket。
4. 觀察Backward是否被Contention拖慢。
5. 再測壓縮、Optimizer Overlap或FSDP Prefetch。
6. 每次只改一個主要參數。
7. 使用相同Seed、Batch和Data比較。
8. 保存最佳與Rollback配置。

拓撲和NCCL演算法可閱讀GPU拓撲感知通訊怎麼做？；推論的TP、EP與KV傳輸可閱讀LLM推論為什麼卡在網路？。

驗收指標

常見問題

Bucket越小，重疊一定越好嗎？

不一定。小Bucket更早啟動，但Message與Launch更多，Network Bandwidth可能無法充分利用。

使用async_op就會自動重疊嗎？

不會。還需要獨立Stream、正確依賴，以及後續可平行Compute；同步太早或資源競爭都會失效。

GPU Utilization提高就代表加速嗎？

不代表。Collective與Busy Wait也可能提高Utilization，仍要看Step Time、Tokens／s和Critical Path。

官方資料

• PyTorch DDP Communication Hooks
• DistributedDataParallel
• torch.distributed

運算與通訊重疊的本質，是把可獨立資料交換移出Critical Path。Gradient Bucket和CUDA Stream提供手段，Profiler則負責證明等待真的被縮短，而不是被藏在更慢的Compute裡。