examlab .net 用最有效率的方法,考取最有價值的證照
本篇導覽 約 20 分鐘

SageMaker 分散式訓練

4,000 字 · 約 20 分鐘閱讀 ·

MLA-C01 Domain 2 Task 2.2 SageMaker 分散式訓練:data vs model vs tensor parallelism、SMDP/SMP 函式庫、ml.p3/p4d/p5 選擇、EFA 網路、Training Compiler XLA、分散式 checkpointing、Horovod,以及是否使用分散式訓練的決策樹。

立即做 20 題練習 → 免費 · 不用註冊 · MLA-C01

SageMaker 上的分散式訓練是將單一訓練任務拆分到多個 GPU 和多個執行個體的技術,讓原本在一台機器需要數天才能完成的模型在數小時內完成,或讓無法裝入單一 GPU 記憶體的模型得以訓練。在 MLA-C01 考試上,分散式訓練是 Domain 2(ML 模型開發,佔 26%)Task 2.2(訓練與精煉模型)的核心考點。這是工程考試,不是資料科學考試——題目很少問「解釋跨 pipeline-parallel mesh 的反向傳播」,幾乎都是「給定模型大小、資料集大小、執行個體預算和訓練時間約束,選擇正確的分散式訓練策略和執行個體型別」。

本指南從 ML Engineer 視角出發。涵蓋三種平行策略(data、model、tensor/pipeline)、兩個 SageMaker 函式庫(SMDP 用於 data parallelism、SMP 用於 model parallelism)、分散式訓練實際跑的執行個體型別(ml.p3.16xlarge、ml.p4d.24xlarge、ml.p5.48xlarge)、讓多節點訓練可行的網路層(Elastic Fabric Adapter)、疊加在上面的 XLA 圖優化層 SageMaker Training Compiler,以及讓分散式訓練對 Spot 中斷和節點故障具備韌性的 checkpointing 模式。同時涵蓋考試喜歡考的同等重要的問題——何時不使用分散式訓練,因為同步開銷可能讓小模型在八個 GPU 上的訓練速度比一個 GPU 還慢。

什麼是分散式訓練及 MLA-C01 為何在意

Definition — SageMaker 分散式訓練 — MLA-C01 ML Engineer 學習筆記. This MLA-C01 topic covers a domain-specific AWS service or pattern. Confirm the canonical definition from official AWS documentation before relying on third-party summaries — service names and feature scoping have shifted over time.

分散式訓練是在多個計算單元上跑單一訓練任務的技術——一個執行個體內的多個 GPU,或每個執行個體各有多個 GPU 的多個執行個體——以減少掛鐘時間或適應更大的模型。MLA-C01 考分散式訓練是因為現代 ML 工作負載已跨越讓單 GPU 訓練不切實際的門檻:基礎模型微調、大型資料集上的電腦視覺、零售規模的時間序列預測,都需要分散式訓練才能具備商業可行性。社群訊號一致——分散式訓練的深度比考生預期的難,考題會出選錯平行策略卻不減少訓練時間反而倍增算力成本的情境。

分散式訓練解決兩個不同的瓶頸

分散式訓練解決兩個截然不同的問題,你選擇的平行策略取決於你遇到哪一個。訓練時間是吞吐量問題——你的模型可以裝進一個 GPU,但訓練需要 50 小時而業務需要 5 小時的結果。模型大小是記憶體問題——你的模型有 700 億個參數,訓練時的權重加上激活值加上 optimizer state 就是裝不進單一 80 GB H100。Data parallelism 解決訓練時間;model parallelism 解決模型大小;tensor parallelism 解決模型大小中更窄的一塊——個別層超過單一 GPU 記憶體的情況。錯誤策略配錯瓶頸會燒錢——100M 參數模型用 model parallelism 跑在八個 GPU 上,訓練速度更慢且成本八倍,而正確的 data-parallel 設定不需要這些。

為何用 SageMaker 而非自管 EC2

你可以用裸機 EC2 加 NCCL、MPI 和自己的調度跑分散式訓練。SageMaker 分散式訓練加上三件考試要你認識的東西。第一,SageMaker Data Parallel 和 Model Parallel 函式庫是 AWS 調優的實作,比原生開源預設值更充分利用 AWS 網路拓樸(placement group、EFA)。第二,SageMaker 管理叢集生命週期——在同一個 placement group 啟動執行個體、分發訓練映像、設定節點間網路、完成後拆解——讓工程師專注在訓練腳本。第三,SageMaker 把分散式訓練與 Spot、checkpoint、warm pool、Experiments 和 Model Registry 整合,這些都在考試範圍內。自管 EC2 分散式訓練是有效的架構選擇,但對 MLA-C01 答案,SageMaker 幾乎永遠是預期服務,除非題幹明確排除它。

為何要分散式訓練 — 三個驅動力

三個力量把工作負載從單 GPU 推向分散式。

驅動力 1 — 訓練時間約束

300 億參數 LLM 在 100 GB 文字上的微調,在單一 ml.p4d.24xlarge(8 個 A100 GPU)上大約需要 80 小時。用 data parallelism 跨四個 ml.p4d.24xlarge 節點(32 個 A100 GPU),同樣的微調大約 22 小時完成。成本不是 32 倍算力——次線性縮放意味著接近 5 倍——但掛鐘時間 4 倍的縮短是業務付費的理由。訓練時間在考試上的呈現是「團隊需要在週五 Demo 前出結果」或「競爭分析需要每日重訓」。當掛鐘時間是約束時,橫向擴展。

驅動力 2 — 超過單一 GPU 記憶體的模型大小

1750 億參數模型用 FP16 光是權重就需要 350 GB,加上訓練期間的激活值和 optimizer state 大約 2 倍——遠超 A100 或 H100 的 80 GB。無論多有耐心,都無法在一個 GPU 上訓練;模型根本載不進去。Model parallelism 把模型拆分到多個 GPU,每個 GPU 持有一部分參數。Tensor parallelism 把個別矩陣乘法切分到多個 GPU。Pipeline parallelism 把模型切成序列階段在不同 GPU 上以流水線排程執行。當模型裝不下時,跨 GPU 拆分。

驅動力 3 — 讓單 GPU IO 成為瓶頸的資料集大小

10 TB 影像資料集太大,無法在單一執行個體的本地 NVMe 上暫存,每個 epoch 從 S3 以 File mode 讀取浪費數小時。Data parallelism 讓每個節點並行讀取資料集的一個分片,結合 FSx for Lustre 作為共用的高吞吐層,消除 IO 瓶頸。考試會出「訓練很慢但 GPU 使用率 30%」——答案很少是「更多 GPU」,通常是「先修 IO 瓶頸」,然後再擴展。

白話文解釋 分散式訓練

分散式訓練是那種抽象概念(「data parallel」、「tensor parallel」)在你把它們對應到實際工作之前感覺很學術的主題。三個類比讓它具體化。

類比 1 — 醫院中央廚房擴大供膳規模

想像醫院中央廚房要為一場大型院慶準備 2000 份便當(訓練模型)。食譜是模型,食材是資料集,距離活動開始的小時數是訓練時間預算。

Data parallelism 是請八位一模一樣的廚師,每人拿到同一份食譜和八分之一的食材。每位廚師並行準備自己那份;結束時大家互相對照心得(「你鹽放多少?」)並平均調整——那個對照步驟就是 AllReduce 梯度同步。食譜(模型)裝進每位廚師的腦袋;唯一的約束是吞吐量。這是 SageMaker Data Parallel Library(SMDP)模式。

Model parallelism 是食譜超龐大——200 頁的料理全書——沒有任何一位廚師能全部背起來。你把食譜本切成章節:廚師 1 負責前菜章、廚師 2 負責主菜章、廚師 3 負責甜點章。他們在流水線中工作——廚師 1 完成 B 組的前菜,廚師 2 還在做 A 組的主菜時廚師 1 已經開始 C 組的前菜。這是 pipeline parallelism,SageMaker Model Parallel Library 的預設。

Tensor parallelism 是連一個章節都太長了——假設主菜章有 50 頁。你把那個單一章節拆給四位廚師,每人同時處理不同的段落,不斷互相協調。Tensor parallelism 對廚師間溝通要求很高,所以廚師必須緊鄰站(同一執行個體、NVLink)——把他們放在不同廚房(不同執行個體)會毀掉吞吐量。

Elastic Fabric Adapter(EFA) 是廚房之間的高速氣動管道系統,讓不同大樓的廚師傳遞餐點的速度如同並肩站立。沒有 EFA,跨廚房溝通靠飛鴿傳書(標準乙太網路),每個 batch 結束時的梯度同步毀掉了加速。Training Compiler 是事先閱讀食譜並為效率重寫它的副廚師——合併步驟、消除多餘的準備工作、挑最佳順序——讓每位廚師做出同樣料理所需的工作少了 30%。

checkpoint 到 S3 是主廚每 15 分鐘影印廚房進度並鎖在保險箱裡;如果廚房失火(Spot 中斷),訓練從最後一份影印繼續,而不是從生食材重新開始。

類比 2 — 工程師備考期間的筆記系統

想像一位 ML Engineer 備考 MLA-C01(訓練一個複雜的心智模型)。所有官方文件是訓練資料,每天剩下的備考時間是訓練時間預算。

Data parallelism 是組成 8 人讀書會,每人各讀官方文件的不同章節,然後定期聚會分享筆記——每個人的心智模型因此更新。文件(模型)可以存在每個人的腦袋,限制是吞吐量(消化速度)。

Model parallelism 是一份文件量龐大到任何一個人都背不起來。你把它切成服務域(Domain 1 由 A 負責、Domain 2 由 B 負責),每個人精通自己的域,在模擬考時接力解題。Pipeline overhead 是每次模擬考開始要等第一個人解完第一題再接棒的等待時間。

Tensor parallelism 是連單一 Domain 都太複雜,需要兩人同時看同一頁文件互相補充理解。他們必須坐在同一張桌子(同一執行個體 NVLink);一個人在圖書館一個人在咖啡廳就無法同步(不能跨執行個體)。

EFA 是讓不同房間的讀書小組能透過高速視訊即時同步筆記,接近面對面的效果。Training Compiler 是把文件整理成高效學習路徑的學習顧問。S3 checkpointing 是每 15 分鐘把筆記上傳雲端;如果電腦當機,從最後一份雲端筆記繼續,不用重看所有文件。

類比 3 — 汽車工廠的組裝流水線

想像汽車工廠生產一批車隊(訓練大型模型)。組裝流程是模型,零件庫存是訓練資料,生產截止日是訓練時間。

Data parallelism 是跑八條相同的組裝線並行生產,每條線從零件庫存中取自己那份。每天結束時工廠管理員比較各線的品質指標並在每條線更新標準流程。每條線跑完整的組裝流程——約束是吞吐量,不是流程複雜度。

Model parallelism 是當組裝流程有太多工站(5000 個)以至於沒有任何一條線夠長。你把流程拆分:1 號線負責工站 1-1500、2 號線負責 1501-3000,以此類推。每條線把半成品車傳給下一條。Pipeline 重疊讓 1 號線開始 B 號車時,2 號線還在做 A 號車的中間階段。

Tensor parallelism 是連一個工站都太複雜,需要四位工人同時在同一個引擎上工作,各自負責不同的螺絲。他們必須並肩站在同一個工站(同一執行個體、NVLink 連線);分散在工廠不同地方會毀掉同步。

EFA 是工廠不同大樓之間的高速輸送系統——等效於物理分隔時仍能並肩站立。Training Compiler 是審查組裝流程並為效率重寫它的工業工程師——消除多餘步驟、合併工站、最佳化操作順序。checkpoint 到 S3 是工廠班長每 15 分鐘拍下每個工站每輛車的狀態並存在防火保險箱;如果颱風摧毀工廠,生產從最後的快照繼續,而不是從原材料重新開始。

Data Parallelism — 大多數分散式訓練的預設

Data parallelism 是最常見且最簡單的分散式訓練策略。先掌握這個,因為它涵蓋大多數 MLA-C01 分散式訓練情境。

Data Parallelism 的運作方式

每個 GPU 取得完整的模型副本和資料的不同 mini-batch 切片。所有 GPU 在各自的切片上獨立執行前向和反向傳播。反向傳播後,梯度在所有 GPU 間同步——AllReduce 操作——讓每個 GPU 最終得到相同的平均梯度。Optimizer 接著以相同方式更新每個副本的參數。模型在每一步都在所有 GPU 間保持同步。Data parallelism 在網路頻寬飽和前幾乎線性縮放。

Data Parallelism 是正確選擇的時機

三個訊號指向 data parallelism。第一,模型舒適地裝進一個 GPU 記憶體——通常對 100 億參數以下的模型,或含 optimizer state 訓練時 10 億參數以下。第二,資料集夠大以至於單 GPU 訓練受掛鐘時間而非算力限制。第三,團隊需要最少的程式碼變更——SageMaker Data Parallel 的 data parallelism 只需要在現有 PyTorch DataLoader 和 optimizer 外面加幾行包裝。

SageMaker Data Parallel Library(SMDP)

SMDP 是針對 ml.p4d、ml.p4de 和 ml.p5 執行個體使用的 EC2 網路拓樸調優的 AWS 優化 data-parallel 函式庫。SMDP 用自訂 AllReduce 取代開源 NCCL AllReduce,充分利用 EFA fabric 和 placement-group 拓樸。函式庫處理梯度壓縮、通訊與計算重疊,以及自適應 batch 大小。透過 Estimator 設定中的 distribution={"smdistributed": {"dataparallel": {"enabled": True}}} 啟用。訓練腳本變更很少——用 smdistributed.dataparallel.torch.distributed 包裹模型、用分散式 sampler 取代標準 DataLoader,然後繼續。

開源替代方案 — PyTorch DDP 和 Horovod

如果你的團隊已經使用 PyTorch DistributedDataParallel(DDP)或 Horovod,兩者都可以在 SageMaker 上不修改地執行——SageMaker 提供啟動腳本配置叢集。SMDP 在 AWS 特定執行個體型別上通常優於 DDP,因為有 EFA 優化,但在較小叢集上效能差距縮小。Horovod 是跨 PyTorch、TensorFlow 和 MXNet 的廠商中立選擇——可攜性好但在 AWS 上比 SMDP 慢。

SageMaker Data Parallel Library 專門針對帶 EFA 網路的 ml.p3.16xlarge、ml.p4d.24xlarge、ml.p4de.24xlarge 和 ml.p5.48xlarge 執行個體優化。 在這個執行個體集合之外——ml.g4dn、ml.g5 或較小的 p 家族——SMDP 要嘛退回標準 NCCL 要嘛拒絕啟用。SMDP 的效能論點只在工作負載大到足以使用 EFA 配備執行個體時成立;對較小的工作,ml.g5 或 ml.p3.8xlarge 上的 PyTorch DDP 更具成本效益。在 MLA-C01 考試上,引用「ml.p4d.24xlarge 叢集、500 億參數 LLM 微調、最少程式碼變更」的題幹指向 SMDP;引用「因預算選 ml.g5.12xlarge、較小模型」的題幹指向不帶 SageMaker 函式庫的 PyTorch DDP。

Model Parallelism — 當模型裝不下時

Model parallelism 是模型大到無法裝進任何單一 GPU 時的答案。MLA-C01 考試測試何時從 data parallel 切換到 model parallel。

Pipeline Parallelism

Pipeline parallelism 按層把模型拆分到 GPU。100 層 transformer 的第 1-25 層可能在 GPU 0 跑,第 26-50 層在 GPU 1,以此類推。每個 micro-batch 流過 pipeline;GPU 0 處理 micro-batch 2 時,GPU 1 在處理 micro-batch 1。Pipeline 讓 GPU 保持忙碌,除了開始(填充 pipeline)和結束(清空它)——「pipeline bubble」開銷。Pipeline parallelism 按階段數降低每 GPU 記憶體,讓需要 200 GB 的模型可以在 4x80GB GPU 上跑。

Tensor Parallelism

Tensor parallelism 跨 GPU 切割個別操作。需要 100 GB 記憶體的大型矩陣乘法,變成在四個 GPU 上並行跑的四個 25 GB 部分矩陣乘法,結果透過 AllReduce 合併。Tensor parallelism 有很高的 GPU 間通訊——它應該只跨 NVLink 連接的 GPU(同一執行個體內),不跨執行個體。為 tensor parallelism 跨越網路會殺死吞吐量。

SageMaker Model Parallel Library(SMP)

SMP 是為 PyTorch 和 TensorFlow 自動化 pipeline 和 tensor parallelism 的 AWS 函式庫。SMP 自動劃分模型圖、把分區放在 GPU 上、排程 micro-batch 通過 pipeline,並插入必要的集體操作。工程師在 Estimator distribution 字典中設定平行度(pipeline_parallel_degreetensor_parallel_degree),SMP 處理其餘部分。SMP 支援自動混合精度(FP16/BF16)、activation checkpointing(反向傳播時重新計算激活值以降低記憶體),以及與 FSDP(Fully Sharded Data Parallel)整合以進一步節省記憶體。

結合 Data 和 Model Parallelism

對非常大的模型,你同時使用兩者。跨節點,用 data parallelism 複製模型(每個節點在自己的資料分片上訓練)。在節點內,用 tensor 和 pipeline parallelism 拆分模型(模型太大無法裝進一個 GPU,但可以裝進一個 ml.p4d.24xlarge 的 8 個 GPU)。這有時稱為「3D parallelism」——data、pipeline、tensor——這是訓練 1000 億參數以上模型的方式。SMP 開箱支援這個混合模式。

跨執行個體的 tensor parallelism 摧毀吞吐量,因為每一層的 AllReduce 都要跨越網路。 Tensor parallelism 應限制在由 NVLink 連接的單一執行個體內的 GPU(ml.p4d.24xlarge 內部提供 600 GB/s GPU 對 GPU 頻寬)。Pipeline parallelism 對網路的容忍度更高,因為激活值只在階段邊界跨 GPU,不是每個操作都要。考試會出「設定 tensor parallelism 跨 4 個節點,看到單節點訓練 30% 吞吐量」——正確修法是把 tensor parallelism 保持在一個節點內,跨節點使用 pipeline parallelism 加 data parallelism。混淆平行策略型別是社群引用最多的分散式訓練痛點之一。

分散式訓練的執行個體選擇

執行個體選擇是分散式訓練決策的一半。考試測試哪些執行個體支援哪些功能。

GPU 執行個體家族

對 MLA-C01,四個 GPU 執行個體家族主導分散式訓練題目。

  • ml.p3.16xlarge — 8 個 V100 GPU(16 GB 或 32 GB),25 Gbps 網路。舊款但對中等規模訓練具成本效益。16xlarge 有 EFA,較小的 p3 沒有。
  • ml.p4d.24xlarge — 8 個 A100 GPU(40 GB),400 Gbps EFA 網路。2023-2024 年代大規模 data-parallel 和 model-parallel 訓練的主力機型。
  • ml.p4de.24xlarge — 8 個 A100 GPU(80 GB),400 Gbps EFA 網路。80GB A100 讓每 GPU 可裝更大的模型。
  • ml.p5.48xlarge — 8 個 H100 GPU(80 GB),3200 Gbps EFA 網路。目前 AWS 基礎模型訓練的最先進機型。
  • ml.g5.48xlarge — 8 個 A10G GPU(24 GB),100 Gbps 網路。比 p4d 便宜,適合小規模 data parallelism,但沒有 EFA。
  • ml.g4dn.12xlarge — 4 個 T4 GPU(16 GB),適中的網路。適合推論和小型訓練;不建議用於大型分散式訓練。

為何即使有 Model Parallelism,GPU 記憶體仍然重要

更大的 GPU 記憶體移動了需要 model parallelism 的門檻。70B 模型 FP16 重約 140 GB;在 80GB H100 上,權重只需要 2 個 GPU 的 model parallelism,留下空間給跨多節點的 data parallelism。在 40GB A100 上需要 4 個 GPU 的 model parallelism,減半了 data-parallel 縮放。考試問「給定 300 億參數模型,哪個執行個體讓我們用最簡單的設定訓練」——選你負擔得起的最大 GPU 記憶體。

網路 — 為何 EFA 是分散式訓練的乘數

Elastic Fabric Adapter(EFA)是繞過 OS 核心的基於 Libfabric 的網路介面,在執行個體之間提供次微秒延遲。EFA 是分散式訓練擴展超過幾個節點的必要條件——沒有它,AllReduce 集體操作在 8 個以上節點時成為瓶頸。EFA 在 ml.p3.16xlarge、ml.p4d、ml.p4de、ml.p5 和少數其他大型執行個體上可用。SageMaker 在你設定支援執行個體上的分散式訓練並請求 placement group 時自動啟用 EFA。

Placement Group 和共置

SageMaker 自動把分散式訓練執行個體放在 cluster placement group 中,確保所有執行個體在同一個實體 AZ 分區,以達到低延遲節點間通訊。你不需要手動設定——SageMaker 處理它。但你需要知道分散式訓練工作可能在繁忙區域等待 placement-group capacity。

在超過兩個 ml.p3.16xlarge、ml.p4d、ml.p4de 或 ml.p5 執行個體上跑分散式訓練時,永遠啟用 EFA。SageMaker 在你選擇這些執行個體型別並設定 SMDP 或 SMP 時自動啟用 EFA。 EFA 在相同執行個體上提供比標準 ENI 網路 4-10 倍的 AllReduce 吞吐量,直接轉化為線性 vs 次線性縮放。成本差異是零——EFA 包含在這些執行個體中,不額外收費——所以跳過 EFA 的唯一理由是你的工作負載確實裝在一個執行個體上,這樣根本不需要分散式訓練。在考試上,任何提議 ml.p4d 用於分散式訓練但省略 EFA 網路的答案都是次佳的。

SageMaker Training Compiler — 優化層

Training Compiler 是基於 XLA 的圖優化層,為目標加速器硬體編譯 PyTorch 和 TensorFlow 模型,通常在相同硬體上提供 20-50% 的加速。

Training Compiler 做什麼

Training Compiler 分析訓練圖並應用 kernel fusion、記憶體佈局優化、混合精度轉換和算子排程。它產生一個在目標硬體上比未優化的 PyTorch/TF 版本跑更快的硬體優化編譯圖。加速幅度因模型而異——transformer 架構通常看到 30% 的改善,卷積模型可能較少。

何時使用 Training Compiler vs 跳過

Training Compiler 在每個訓練工作開始時增加編譯時間——通常 5-15 分鐘。對短訓練跑,編譯開銷可能超過加速收益。Training Compiler 對 4 小時以上的跑最有益。它也有模型架構限制——不是每個 PyTorch 算子都支援,不支援的算子強制退回急切執行,侵蝕加速效果。在提交完整多天分散式訓練跑並啟用 compiler 之前,先在一個小型代表性訓練工作上測試相容性。

結合 Training Compiler 與 SMDP 和 SMP

Training Compiler 可與 SMDP(data parallelism)和 SMP(model parallelism)疊加。設定是分層的:SMDP/SMP 定義工作如何跨 GPU 分配;Training Compiler 優化在每個 GPU 上跑的工作。組合加速是相乘而非相加——在 80% data-parallel 縮放效率上疊加 30% compiler 加速,在規模上有意義。

韌性分散式訓練的 Checkpointing

分散式訓練工作跑數小時到數天。沒有 checkpointing,任何中斷——Spot 回收、執行個體故障、網路分區——都毀掉所有進度。

S3 Checkpoint 設定

SageMaker 訓練工作在 Estimator 中接受 checkpoint_s3_uricheckpoint_local_path。訓練期間寫到本地 checkpoint 路徑的檔案自動同步到 S3,工作重新啟動時 S3 內容恢復到本地路徑。訓練腳本必須(1)定期將模型狀態儲存到本地路徑,(2)開始時檢查是否有現有 checkpoint 並從中繼續。

分散式 Checkpoint 模式

對 data-parallel 訓練,只有一個 GPU 寫 checkpoint——通常是 rank 0。其他 rank 有相同的模型狀態,重複寫入浪費 S3 吞吐量。SMP 提供分散式 checkpoint 工具,從所有 rank 收集分片的模型狀態成單一連貫 checkpoint,然後從 rank 0 寫入。

Checkpoint 頻率取捨

頻繁 checkpointing 防止中斷,但需要 IO 時間。典型模式:短工作每個 epoch(每 epoch 1 小時以下),長工作每 N 步(一個 epoch 可能需要 8 小時以上)。Spot 訓練工作從更頻繁的 checkpointing 中受益——目標是在 Spot 回收時損失不超過 30 分鐘的工作。

沒有 S3 checkpointing 的分散式訓練不符合生產標準——每一次 Spot 中斷、每一次節點故障、每一次工作重新啟動都毀掉所有進度。 SageMaker 的 checkpoint_s3_uri 參數自動化本地執行個體儲存和 S3 之間的同步;重新啟動時,先前的 checkpoint 在訓練繼續前恢復。對 Managed Spot Training(節省高達 90% 的成本),checkpointing 不是選項——多小時訓練工作的 Spot 回收概率高到最終一定會發生無法復原的工作。考試測試這個組合:任何把「Managed Spot Training」和分散式訓練結合的答案必須同時設定 checkpoint 路徑;提議 Spot 但沒有 checkpoint 的答案是錯的。

何時使用分散式訓練

分散式訓練有開銷。GPU 間和節點間的同步成本增加延遲,可能對小型工作負載超過加速收益。

同步稅

Data parallelism 中的每個 batch 都需要跨所有 GPU 的 AllReduce。AllReduce 時間隨叢集大小增長,受網路頻寬限制。對每 GPU batch 小的小模型,同步時間可能超過每 GPU 計算時間——意味著更多 GPU 讓工作更慢,而不是更快。收支平衡點取決於模型大小、batch 大小和網路。經驗法則:100M 參數以下且 batch 小於 64,分散式訓練很少划算。

通訊計算比

Data-parallel 縮放效率取決於每 batch 計算時間與通訊時間的比值。大 batch 的大模型計算時間長,攤薄通訊成本——有效縮放到數十個節點。小 batch 的小模型計算時間短,通訊主導——超過 2-4 個節點縮放效果差。

先在單節點多 GPU 上縮放

加執行個體前,先在一個執行個體內縮放。ml.p4d.24xlarge 有 8 個由 NVLink 以 600 GB/s 連接的 GPU——比執行個體之間的 EFA 快得多。一個 ml.p4d.24xlarge 上跨 8 個內部 GPU 的 data parallelism 訓練幾乎線性縮放,沒有任何跨執行個體開銷。只有當一個執行個體不夠用時才應該擴展到多個執行個體。

改增加 Batch Size

如果 GPU 記憶體還有餘裕,增加 batch size 在不分散的情況下提高吞吐量。batch size 32 的 ml.p4d.24xlarge 可能 GPU 記憶體使用不足;增到 batch size 256 保持相同執行個體數量但每秒處理 8 倍更多資料。結合 gradient accumulation 以維持有效 batch size 同時適應記憶體。

MLA-C01 分散式訓練的常見考試陷阱

考試會埋設特定的錯誤觀念。認識它們。

陷阱 1 — 對小模型選擇 Model Parallelism

剛看完基礎模型文章的工程師預設選 model parallelism。對舒適裝進一個 GPU 的 1 億參數模型,model parallelism 增加 pipeline 開銷卻零收益。Data parallelism(或根本不用分散式)才是正確的。

陷阱 2 — 跨執行個體使用 Tensor Parallelism

Tensor parallelism 通訊量大,必須留在一個執行個體內(NVLink)。跨執行個體擴展會毀掉吞吐量。Pipeline parallelism 是跨執行個體的 model-parallel 選擇;tensor parallelism 是執行個體內的選擇。

陷阱 3 — 多節點分散式訓練略過 EFA

沒有 EFA,多節點訓練受 ENI 網路瓶頸。考試把沒有 EFA 的 ml.p4d 列為錯誤答案;正確選擇永遠包含 EFA。

陷阱 4 — 分散式訓練沒有 Checkpointing

特別是使用 Managed Spot Training 時。Spot 回收是機率性的;多小時工作最終一定會被回收。沒有 checkpoint 意味著從頭開始。

陷阱 5 — 假設線性縮放

GPU 加倍不等於訓練時間減半。AllReduce 開銷、IO 和 pipeline bubble 造成次線性縮放。32 個 GPU 的實際縮放效率通常是 60-80%;256 個 GPU 如果不仔細調優可能降到 40-60%。

陷阱 6 — Training Compiler 總是加速

錯誤。對短工作,編譯時間可能超過加速效果。先在代表性工作上測試。

陷阱 7 — 忽略 IO 才是真正的瓶頸

如果 GPU 使用率低於 80%,瓶頸是 IO 或 CPU 預處理,不是 GPU 計算。加更多 GPU 沒幫助;修資料 pipeline 才有。

陷阱 8 — Spot 分散式訓練不考慮超參數

分散式 Spot 可能在訓練中途中斷。依賴一致步數的超參數(warmup schedule、LR decay schedule)需要意識到重新啟動。

陷阱 9 — 混淆 SMDP 與 SMP

SMDP 是 data parallel——每個 GPU 有完整模型。SMP 是 model parallel——每個 GPU 有一部分。它們是解決不同問題的不同函式庫。考試反覆出這個區分。

陷阱 10 — 混合 Horovod 與 SMDP

選一個分散式框架。混合使用產生不可預測的行為。SageMaker Estimator 的 distribution 參數只能選一個策略。

MLA-C01 exam priority — SageMaker 分散式訓練 — MLA-C01 ML Engineer 學習筆記. This topic carries weight on the MLA-C01 exam. Master the trade-offs, decision boundaries, and the cost/performance triggers each AWS service exposes — the exam will test scenarios that hinge on knowing which service is the wrong answer, not just which is right.

FAQ — SageMaker 上的分散式訓練

Q1 — 什麼時候用 SageMaker Data Parallel(SMDP)vs PyTorch DDP?

在 EFA 配備執行個體(ml.p3.16xlarge、ml.p4d、ml.p4de、ml.p5)上跑,且想要 AWS 特定拓樸上最大 AllReduce 吞吐量時,用 SMDP。SMDP 在多節點設定上通常比 NCCL 好 20-40%。在較小的沒有 EFA 的執行個體(ml.g5、ml.p3.8xlarge)上、你的團隊已有現有 DDP 程式碼且 AWS 特定調優預算有限、或需要在 AWS 和其他雲之間的可攜性時,用 PyTorch DDP。MLA-C01 考試對任何大型 p4d/p5 分散式訓練情境偏好 SMDP;PyTorch DDP 作為較小成本受限訓練的答案出現。

Q2 — 我的 700 億參數模型裝不進單一 A100 80GB。最簡單的分散式訓練設定是什麼?

跨一個 ml.p4de.24xlarge 或 ml.p5.48xlarge 的 8 個 GPU 的 pipeline parallelism。70B 模型 FP16 權重約 140 GB;分散到 8 個 GPU 每個 GPU 17.5 GB 用於權重,留下激活值和 optimizer state 的空間。設定 SMP pipeline_parallel_degree=8,啟用 activation checkpointing 降低激活值記憶體,用 BF16 混合精度減半權重記憶體。如果單節點不夠(例如,訓練也需要 data parallelism 以提高吞吐量),擴展到多節點,hybrid pipeline + data parallelism——節點內 pipeline parallel,跨節點 data parallel。

Q3 — 我設定了 8 個 ml.p4d.24xlarge 節點,但訓練比 1 個節點快不了多少。我錯過了什麼?

按診斷樹排查。(1)EFA 是否啟用——沒有 EFA,ENI 頻寬瓶頸 AllReduce。SageMaker 在支援執行個體上自動啟用 EFA;在 CloudWatch metrics 中驗證。(2)資料載入器是否是瓶頸——查 GPU 使用率。如果低於 80%,資料集 IO 或預處理使 GPU 挨餓。切換到 FastFile mode 或 FSx for Lustre。(3)Batch size 是否太小——小的每 GPU batch 讓通訊成本主導。增加 batch size 或使用 gradient accumulation。(4)AllReduce 策略是否正確——SMDP 的優化 AllReduce 通常優於原生 NCCL。(5)節點是否共置——SageMaker 應該自動把它們放在 placement group;如果你在裸機 EC2 上跑,需要手動設定。

Q4 — 我可以對分散式工作使用 Managed Spot Training 嗎?

可以,但你必須設定 checkpointing。Managed Spot 節省高達 90% 的訓練成本,但 Spot capacity 可能在 2 分鐘警告後被回收。對分散式訓練,在 Estimator 中設定 checkpoint_s3_uricheckpoint_local_path。訓練腳本必須定期儲存狀態到本地路徑,並在啟動時從 S3 恢復。SageMaker 自動處理 S3 同步。預期超過幾小時的工作會有中斷;調整 checkpoint 頻率讓每次中斷損失不超過 30 分鐘的掛鐘時間。

Q5 — Pipeline parallelism 和 tensor parallelism 的區別是什麼,各自何時使用?

Pipeline parallelism 沿層軸拆分模型——不同 GPU 持有不同層。通訊發生在層邊界(激活值在階段之間傳遞)。Pipeline parallelism 通訊量輕,可容忍執行個體之間的 EFA 延遲,是跨執行個體 model parallelism 的正確選擇。Tensor parallelism 跨 GPU 拆分個別操作(矩陣乘法)。通訊發生在每個操作內,要求 NVLink 級別的頻寬——tensor parallelism 只能在一個執行個體內工作。組合使用:節點內 8 個 GPU 之間的 tensor parallelism(NVLink),跨節點的 pipeline parallelism(EFA),更多組群節點複製整個結構用於 data parallelism。

Q6 — 我想在分散式 PyTorch 工作上用 SageMaker Training Compiler。應該預期什麼?

預期工作開始時有 5-15 分鐘的編譯開銷,然後每 GPU 計算速度提高 20-50%。Compiler 與 SMDP 乾淨疊加——設定 distribution={"smdistributed": {"dataparallel": {"enabled": True}}}compiler_config=TrainingCompilerConfig()。先跑一個短的代表性測試(例如 30 分鐘)以驗證(1)你的所有算子都支援(不支援的算子退回急切執行並侵蝕加速)、(2)你的 batch size 仍然裝得下——compiler 可能改變記憶體佈局,編譯前裝得下的 batch 可能編譯後 OOM。Compiler 對 4 小時以上的訓練跑最值得,此時編譯成本得以攤薄。

Q7 — 我的叢集有時啟動失敗,出現「容量不足」錯誤。如何處理?

分散式訓練需要同一 AZ 中的 placement-group capacity,對繁忙區域的 ml.p4d 和 ml.p5 可能稀缺。三個緩解方案。(1)使用 SageMaker 的訓練工作重試邏輯——設定 max_retry_attempts 讓工作自動重新排隊。(2)選擇較不繁忙的訓練區域(us-west-2 通常比 us-east-1 有更多 p4d/p5 capacity)。(3)透過 Capacity Reservations 或 Savings Plans 為可預測的訓練排程保留 capacity。(4)使用 warm pool——一旦成功申請到叢集 capacity,warm pool 以適中的成本維持它被分配給後續訓練工作。對一次性大型微調,簡單重試通常夠用;對生產重訓排程,保留 capacity。

進一步閱讀 — 分散式訓練的官方 AWS 文件

對超出 MLA-C01 範圍的深度,權威 AWS 資源包括:SageMaker 分散式訓練登陸頁(概覽和決策指南)、SageMaker Data Parallel Library 文件(SMDP API 參考和調優)、SageMaker Model Parallel Library 文件(SMP 設定和平行策略)、SageMaker Training Compiler 文件(支援的算子和基準測試結果)、Elastic Fabric Adapter 文件(網路概念和支援執行個體),以及 SageMaker 訓練執行個體型別矩陣(哪些執行個體支援哪些功能)。

AWS Machine Learning Blog 有在 SMP 訓練大型模型的實際案例研究——包括微調 Llama 級別模型、BERT 預訓練和大規模電腦視覺的操作說明。AWS Well-Architected Machine Learning Lens 涵蓋分散式訓練架構的效能支柱指引。AWS re:Invent 和 re:MARS 的分散式訓練議題包含 SMP 和 SMDP 的實作示範。GitHub 倉庫 awslabs/sagemaker-distributed-training-workshop 和 SageMaker examples aws/amazon-sagemaker-examples 下包含 data parallelism 和 model parallelism 場景的端到端 notebook 範例。

官方資料來源

更多 MLA-C01 主題