在 HyperPod 上使用 Slurm 執行分散式訓練工作負載 - Amazon SageMaker AI

本文為英文版的機器翻譯版本,如內容有任何歧義或不一致之處,概以英文版為準。

在 HyperPod 上使用 Slurm 執行分散式訓練工作負載

SageMaker HyperPod 專門用於訓練大型語言模型 (LLM) 和基礎模型 (FM) 的工作負載。這些工作負載通常需要針對 ML 基礎設施和資源使用多個平行化技術和最佳化的操作。使用 SageMaker HyperPod,您可以使用下列 SageMaker AI 分散式訓練架構:

在 SageMaker HyperPod 上使用 SMDDP

SMDDP 程式庫是集體通訊程式庫,可改善分散式資料平行訓練的運算效能。SMDDP 程式庫使用下列開放原始碼分散式訓練架構:

SMDDP 程式庫透過為 SageMaker HyperPod 提供下列項目,解決金鑰集體通訊操作的通訊負荷。

  • 程式庫提供AllGather針對 最佳化AWS。 AllGather 是一種用於碎片資料平行訓練的關鍵操作,這是一種由熱門程式庫提供的記憶體效率資料平行處理技術。其中包括 SageMaker AI 模型平行化 (SMP) 程式庫、DeepSpeed 零冗餘最佳化工具 (ZeRO) 和 PyTorch 全碎片資料平行化 (FSDP)。

  • 程式庫會透過充分利用AWS網路基礎設施和 SageMaker AI ML 執行個體拓撲來執行最佳化node-to-node通訊。

執行範例資料平行訓練任務

探索下列使用 SMDDP 程式庫實作資料平行化技術的分散式訓練範例。

設定在 SageMaker HyperPod 上使用 SMDDP 程式庫的環境

以下是在 SageMaker HyperPod 上使用 SMDDP 程式庫的訓練環境要求。

  • PyTorch v2.0.1 及更新版本

  • CUDA v11.8 及更新版本

  • 大於 3 的 libstdc++ 執行時期版本

  • Python v3.10.x 及更新版本

  • ml.p4d.24xlargeml.p4de.24xlarge,這是 SMDDP 程式庫支援的執行個體類型

  • 在訓練主機上啟用 imdsv2

根據您想要執行分散式訓練任務的方式,有兩種安裝 SMDDP 程式庫的選項:

  • 使用 SMDDP 二進位檔案的直接安裝。

  • 使用預先安裝了 SMDDP 程式庫的 SageMaker AI 深度學習容器 (DLC)。

預先安裝了 SMDDP 程式庫的 Docker 映像檔或 SMDDP 二進位檔案的 URL 會列示在 SMDDP 程式庫文件中支援的架構

在 SageMaker HyperPod DLAMI 上安裝 SMDDP 程式庫
  • pip install --no-cache-dir https://smdataparallel.s3.amazonaws.com/binary/pytorch/<pytorch-version>/cuXYZ/YYYY-MM-DD/smdistributed_dataparallel-X.Y.Z-cp310-cp310-linux_x86_64.whl

    注意

    如果您在 Conda 環境中工作,請確定您使用 conda install 而非 pip 安裝 PyTorch。

    conda install pytorch==X.Y.Z torchvision==X.Y.Z torchaudio==X.Y.Z pytorch-cuda=X.Y.Z -c pytorch -c nvidia
在 Docker 容器上使用 SMDDP 程式庫
  • SMDDP 程式庫會預先安裝在 SageMaker AI 深度學習容器 (DLC) 上。若要使用 SMDDP 程式庫尋找 PyTorch 的 SageMaker AI 架構 DLC 清單,請參閱 SMDDP 程式庫文件中支援的架構。您也可以自帶已安裝必要相依性的 Docker 容器來使用 SMDDP 程式庫。若要進一步了解如何設定自訂 Docker 容器以使用 SMDDP 程式庫,另請參閱使用 SageMaker AI 分散式資料平行程式庫建立您自己的 Docker 容器

    重要

    若要在 Docker 容器中使用 SMDDP 程式庫,請將 /var/log 目錄從主機電腦掛載到容器中的 /var/log。您可以在執行容器時新增下列選項來完成此操作。

    docker run <OTHER_OPTIONS> -v /var/log:/var/log ...

若要了解通常如何使用 SMDDP 執行資料平行訓練任務,請參閱使用 SageMaker AI 分散式資料平行化程式庫進行分散式訓練

在 SageMaker HyperPod 叢集上使用 SMP HyperPod

SageMaker AI 模型平行化 (SMP) 程式庫提供各種最先進的模型平行化技術,包括:

  • 全碎片資料平行化

  • 專家平行化

  • 使用 FP16/BF16 和 FP8 資料類型的混合精確度訓練

  • 張量平行化

SMP 程式庫也與 PyTorch FSDP、NVIDIA Megatron 和 NVIDIA 轉換器引擎等開放原始碼架構相容。

執行範例模型平行化訓練工作負載

SageMaker AI 服務團隊在 awsome-distributed-training/3.test_cases/17.SM-modelparallelv2 提供使用 SMP 程式庫實作模型平行化的範例訓練任務。