在 A3 Mega 虛擬機器上使用 Megatron-LM 訓練 Llama2

總覽

本快速入門導覽課程將說明如何在 A3 Mega 上，執行以容器為基礎的 Megatron-LM PyTorch 工作負載。您可以在這個 GitHub 存放區取得程式碼：megatron-gke。

事前準備

請按照下列步驟啟用 Google Kubernetes Engine (GKE) API：

  1. In the Cloud de Confiance console, on the project selector page, select or create a Cloud de Confiance project.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  2. Verify that billing is enabled for your Cloud de Confiance project.

  3. Enable the GKE API.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the API

  4. 確認您在專案中具備下列角色： roles/container.admin、roles/compute.networkAdmin、roles/iam.serviceAccountUser

    檢查角色

    1. 前往 Cloud de Confiance 控制台的「IAM」頁面。

      前往「IAM」頁面
    2. 選取專案。

    3. 在「主體」欄中，找出所有識別您或您所屬群組的資料列。如要瞭解自己所屬的群組，請與管理員聯絡。

    4. 針對指定或包含您的所有列，請檢查「角色」欄，確認角色清單是否包含必要角色。

    授予角色

    1. 前往 Cloud de Confiance 控制台的「IAM」頁面。

      前往「IAM」頁面
    2. 選取專案。
    3. 按一下「Grant access」(授予存取權)。

    4. 在「New principals」(新增主體) 欄位中，輸入您的使用者 ID。 這通常是指員工身分集區中使用者的 ID。詳情請參閱「在 IAM 政策中代表工作團隊集區使用者」，或聯絡管理員。

    5. 按一下「Select a role」(選取角色)，然後搜尋角色。
    6. 如要授予其他角色，請按一下「Add another role」(新增其他角色)，然後新增其他角色。
    7. 按一下「Save」(儲存)。

建立 A3 Mega 叢集

建立啟用 GPUDirect-TCPXO 和多網路功能的 A3 Mega GKE 叢集。詳情請參閱「使用 GPUDirect 和多重網路，盡量提高 GPU 網路頻寬」。

設定環境

  1. 為一些常見參數建立環境變數

    export CLUSTER_NAME=CLUSTER_NAME
export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
export PROJECT_ID=PROJECT_ID

    更改下列內容：

    • CLUSTER_NAME：A3 Mega GKE 叢集的名稱，該叢集已啟用 GPUDirect-TCPXO 和多網路功能。
    • CONTROL_PLANE_LOCATION：叢集控制層的 Compute Engine 位置。為地區叢集提供區域，或為區域叢集提供可用區。
    • PROJECT_ID：您的 Cloud de Confiance 專案 ID。

  2. 將 Google Cloud CLI 設為使用您的 Cloud de Confiance by S3NS 憑證進行驗證：

    gcloud auth login

    詳情請參閱「Google Cloud CLI 驗證」。

  3. 安裝 kubectl 和 GKE gcloud CLI 外掛程式：

    sudo apt-get install kubectl
sudo apt-get install google-cloud-sdk-gke-gcloud-auth-plugin

  4. 擷取 GKE 叢集的憑證：

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --location=${CONTROL_PLANE_LOCATION} \
  --project=${PROJECT_ID}

  5. 如果尚未安裝，請安裝 Helm：

    curl -fsSL -o get_helm.sh https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3
chmod 700 get_helm.sh
./get_helm.sh && rm get_helm.sh
sudo chmod +x /usr/local/bin/helm

使用拓撲感知排程器部署 Pod

您可以使用拓撲感知排程器，將 GKE Pod 部署至具有指定 GPU 拓撲的節點。

在下列 kubectl 指令中，您將直接使用存放區中的檔案。或者，您也可以在本機複製存放區，然後 kubectl 指令可以改為參照本機檔案。

詳情請參閱「拓撲排程器」。

  1. 設定服務帳戶：

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/service-account.yaml

  2. 在 configmap 中安裝拓撲排程器指令碼：

    curl -OL  https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/schedule-daemon.py
curl -OL  https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/label-nodes-daemon.py

kubectl -n kube-system create configmap topology-scheduler-scripts \
    --from-file=schedule-daemon.py=schedule-daemon.py \
    --from-file=label-nodes-daemon.py=label-nodes-daemon.py

  3. 安裝拓撲標籤 DaemonSet 和拓撲排程器 Pod：

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/label-nodes-daemon.yaml
$ kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/schedule-daemon.yaml

  4. 觀察拓撲排程器的動作：

    kubectl -n kube-system logs topology-scheduler-pod

執行工作負載

建構 Dockerfile 並推送至 Cloud de Confiance by S3NS Artifact Registry

  1. 建立 Cloud Storage bucket 和 Docker 存放區。 在 scripts/setup-and-configure-resources.sh script 中，將值區和存放區名稱替換為您建立的名稱，然後執行指令碼：

    bash scripts/setup-and-configure-resources.sh

  2. 建構 pytorch-megatron:23.11-py3 映像檔並推送至存放區。 確認 scripts/build-and-push-docker-image.sh 檔案中的 Docker 存放區名稱，與 scripts/setup-and-configure-resources.sh 指令碼中使用的存放區名稱相符。您也可以在推送前編輯 Docker 映像檔標記名稱。

    bash scripts/build-and-push-docker-image.sh

啟動 Megatron-LM Llama2 基準測試

  1. 編輯 helm/values.yaml 檔案，指定 Cloud Storage bucket 和先前章節中建立的 Docker 映像檔。如需設定範例，請參閱「sample-configurations」。

  2. 選用：您也可以編輯 selected-configuration.sh 檔案，指定對預設 Helm 設定所做的任何變更。

    helm install HELM_EXPERIMENT_NAME helm/ --values helm/values.yaml

    將 HELM_EXPERIMENT_NAME 替換成實驗的任意名稱。

實驗會將 Nsight Systems 分析工具的指標寫入 megatron-experiments 目錄中指定的 Cloud Storage bucket。

清除所用資源

為了避免系統向您的 Cloud de Confiance 帳戶收取本頁面所用資源的費用，請按照下列步驟操作。

刪除 GKE 叢集：

前往「Clusters」(叢集) 頁面：

前往「Clusters」(叢集)

  1. 勾選 CLUSTER_NAME 的核取方塊。
  2. 按一下「刪除」圖示 。
  3. 如要確認刪除，請輸入 CLUSTER_NAME，然後按一下「刪除」。

刪除 Cloud Storage bucket

前往「Buckets」(值區) 頁面：

前往值區

  1. 選取您為本快速入門建立的 Cloud Storage bucket 核取方塊。

  2. 按一下「刪除」圖示 。

  3. 如要確認刪除，請輸入 DELETE，然後按一下「刪除」。

後續步驟