適合大規模分散式培訓的專用基礎架構 Amazon SageMaker HyperPod 上市

今天，我們推出  ，它透過為大規模分散式培訓提供專用基礎設施，幫助縮短基礎模型的培訓時間。現在，您可以使用 SageMaker HyperPod 對基礎模型進行數週甚至數月的訓練，同時 SageMaker 會主動監控叢集的健康狀況，並透過替換故障節點和從檢查點恢復模型訓練來提供自動化的節點和作業彈性。

叢集預配置了 SageMaker 的分散式訓練庫，可幫助您在所有節點上分割訓練資料和模型，以便並行處理，並幫助您充分利用叢集的計算和網路基礎設施。您還可以安裝額外的框架、除錯工具和最佳化庫，進一步自定義訓練環境。

下面為您介紹如何開始使用 SageMaker HyperPod。在下面的演示中，我建立了一個 SageMaker HyperPod，展示如何使用  GitHub 儲存庫中共享的示例訓練   模型。

建立和管理叢集

作為 SageMaker HyperPod 管理員，您可以使用   或  建立和管理叢集。在 中，導航至  Amazon SageMaker，在左側選單中選擇  HyperPod 叢集下的 叢集管理，然後選擇 建立叢集。

在接下來的設定中，提供一個叢集名稱，並根據所選例項型別和要分配給每個例項組的例項數量配置例項組。

此外，需要準備一個或多個生命週期指令碼並上傳到   桶，以便在叢集建立期間在每個例項組中執行。有了生命週期指令碼，就可以自定義叢集環境並安裝所需的庫和軟體包。您可以在 GitHub 儲存庫中找到 SageMaker HyperPod 的 。

使用 Amazon CLI

您還可以使用 Amazon CLI 建立和管理叢集。在演示中，我在 JSON 檔案中指定了叢集配置。我選擇建立兩個例項組，一個名為 “ controller-group ” ，用於叢集控制器節點，另一個名為 “ worker-group ” ，用於叢集工作節點。對於將執行模型訓練的工作節點，我指定使用由   晶片提供支援的   例項。

// demo-cluster.json[
    {
        "InstanceGroupName": "controller-group",
        "InstanceType": "ml.m5.xlarge",
        "InstanceCount": 1,
        "LifeCycleConfig": {
            "SourceS3Uri": "s3://<your-s3-bucket>/<lifecycle-script-directory>/",
            "OnCreate": "on_create.sh"
            },
        "ExecutionRole": "arn:aws:iam::111122223333:role/my-role-for-cluster",
        "ThreadsPerCore": 1
    },
    {
        "InstanceGroupName": "worker-group",
        "InstanceType": "ml.trn1.32xlarge",
        "InstanceCount": 4,
        "LifeCycleConfig": {
            "SourceS3Uri": "s3://<your-s3-bucket>/<lifecycle-script-directory>/",
            "OnCreate": "on_create.sh"
            },
        "ExecutionRole": "arn:aws:iam::111122223333:role/my-role-for-cluster",
        "ThreadsPerCore": 1
    }]

為建立叢集，我執行了以下 Amazon CLI 命令：

aws sagemaker create-cluster \
    --cluster-name antje-demo-cluster \
    --instance-groups file://demo-cluster.json

建立後，您可以使用  amazon sagemaker describe-cluster 和  amazon sagemaker list-cluster-nodes 檢視叢集和節點的詳細資訊。請記下控制器節點的叢集 ID 和例項 ID 。需要這些資訊才能連線到叢集。

您還可以選擇附加   等共享檔案系統。欲使用 FSx for Lustre，需要使用   配置來設定叢集。這是一個  模板，展示瞭如何  以及如何 。

連線您的叢集

作為叢集使用者，您需要擁有叢集管理員提供的叢集訪問許可權。獲得訪問許可權後，您可以使用 SSH 連線到叢集，並計劃和執行作業。您可以使用預安裝的 來連線叢集的控制器節點。

在演示中，我執行了以下命令，將控制器節點的叢集 ID 和例項 ID 指定為目標。

aws ssm start-session \
    --target sagemaker-cluster:ntg44z9os8pn_controller-group-i-05a854e0d4358b59c \
    --region us-west-2

使用 Slurm 在叢集上排程和執行作業

SageMaker HyperPod 在釋出之初便支援使用   進行工作負載編排。Slurm 是熱門的開源叢集管理和作業排程系統。在叢集建立過程中，可以透過生命週期指令碼安裝和設定 Slurm 。 展示了操作方法。然後，可以使用標準 Slurm 命令來計劃和啟動作業。有關架構詳細資訊和有用命令，請檢視　 。

在本演示中，我使用了  ，該示例展示瞭如何使用 Trn1 例項在 Slurm 上訓練 Llama 2 7B。我的叢集已經用 Slurm 設定好了，並且已經掛載了 FSx for Lustre 檔案系統。

注意

Llama 2 模型由   管理。您可以透過 請求訪問許可權。

設定叢集環境

SageMaker HyperPod 支援在各種環境中進行訓練，其中包括   、  、  和   。我按照   中的說明，構建了虛擬環境  amazon_neuron_venv_pytorch，並設定了   和   庫，以便在 Trn1 例項上訓練模型。

準備模型、分詞器和資料集

我按照說明，下載了 Llama 2 模型和分詞器，然後將模型轉換為 Hugging Face 格式。接著又下載了　 並對其進行了分詞。最後一個準備步驟是使用預先（AOT）編譯來預編譯 Llama 2 模型，以加快模型訓練。

在叢集上啟動作業

現在，我準備就緒，可以使用　 sbatch 命令開始模型訓練作業。

sbatch --nodes 4 --auto-resume=1 run.slurm ./llama_7b.sh

您可以使用  squeue 命令檢視作業佇列。一旦訓練作業執行起來，SageMaker HyperPod 的彈性功能就會自動啟用。SageMaker HyperPod 會自動檢測硬體故障，根據需要替換節點，如果設定了  auto-resume 引數，還會從檢查點恢復訓練，如上一條命令所示。

您可以在以下檔案中檢視模型訓練作業的輸出：

tail -f slurm-run.slurm-<JOB_ID>.out

表明模型訓練已開始的示例輸出如下所示：

Epoch 0:  22%|██▏       | 4499/20101 [22:26:14<77:48:37, 17.95s/it, loss=2.43, v_num=5563, reduced_train_loss=2.470, gradient_norm=0.121, parameter_norm=1864.0, global_step=4512.0, consumed_samples=1.16e+6, iteration_time=16.40]
Epoch 0:  22%|██▏       | 4500/20101 [22:26:32<77:48:18, 17.95s/it, loss=2.43, v_num=5563, reduced_train_loss=2.470, gradient_norm=0.121, parameter_norm=1864.0, global_step=4512.0, consumed_samples=1.16e+6, iteration_time=16.40]
Epoch 0:  22%|██▏       | 4500/20101 [22:26:32<77:48:18, 17.95s/it, loss=2.44, v_num=5563, reduced_train_loss=2.450, gradient_norm=0.120, parameter_norm=1864.0, global_step=4512.0, consumed_samples=1.16e+6, iteration_time=16.50]

欲進一步監控和剖析模型訓練作業，可以使用  或任何其他工具。

現已推出

SageMaker HyperPod 現已在以下亞馬遜雲科技商業服務區域推出：美國東部（俄亥俄州）、美國東部（北弗吉尼亞州）、美國西部（俄勒岡州）、亞太地區（新加坡）、亞太地區（悉尼）、亞太地區（東京）、歐洲（法蘭克福）、歐洲（愛爾蘭）和歐洲（斯德哥爾摩）地區。

瞭解更多資訊：

● 有關定價資訊和支援的叢集例項型別列表，請參閱 

● 檢視

● 訪問 

文章來源：  https://aws.amazon.com/blogs/aws/introducing-amazon-sagemaker-hyperpod-a-purpose-built-infrastructure-for-distributed-training-at-scale?trk=cndc-detail/

— 

附言：即使文章標題下只有一個署名，亞馬遜雲科技的部落格文章也是團隊共同努力的成果。因此，我要感謝   、  、  、  、 　和　 　在示例程式碼方面給予的慷慨幫助以及對管理大規模模型訓練基礎設施、Slurm 和 SageMaker HyperPod 方面專業知識的不吝分享。

作者 Antje Barth