A3 Mega 仮想マシンで Megatron-LM を使用して Llama2 をトレーニングする

A3 Mega 仮想マシンで Megatron-LM を使用して Llama2 をトレーニングする

Standard

概要

このクイックスタートでは、A3 Mega でコンテナベースの Megatron-LM PyTorch ワークロードを実行する方法について説明します。コードは、GitHub リポジトリ megatron-gke で入手できます。

始める前に

次の手順で Google Kubernetes Engine（GKE）API を有効にします。

アカウントにログインします。 Google Cloud を初めて使用する場合は、 アカウントを作成して、実際のシナリオで Google プロダクトのパフォーマンスを評価してください。 Google Cloud新規のお客様には、ワークロードの実行、テスト、デプロイができる無料クレジット $300 分を差し上げます。

In the Google Cloud console, on the project selector page, select or create a Google Cloud 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.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

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

In the Google Cloud console, on the project selector page, select or create a Google Cloud 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.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

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

1. プロジェクトで次のロール（複数の場合あり）が割り当てられていることを確認します。 roles/container.admin、roles/compute.networkAdmin、roles/iam.serviceAccountUser

   ロールを確認する

   1. コンソールで、[IAM] ページに移動します。 Google Cloud

      IAM に移動
   2. プロジェクトを選択します。

   3. [Principal] 列で、自分または自分が所属するグループの行をすべて確認します。所属するグループについては、管理者にお問い合わせください。

   4. 自分のメールアドレスを含む行の [**ロール**] 列で、ロールのリストに必要なロールが含まれているかどうか確認します。

   ロールを付与する

   1. コンソールで、[IAM] ページに移動します。 Google Cloud

      IAM に移動
   2. プロジェクトを選択します。
   3. [**アクセスを許可**] をクリックします。

   4. [新しいプリンシパル] フィールドに、ユーザー ID を入力します。 これは通常、Google アカウントのメールアドレスです。

   5. [**ロールを選択**] をクリックして、ロールを検索します。
   6. 追加のロールを付与するには、 [Add another role] をクリックして各ロールを追加します。
   7. [保存] をクリックします。

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: GPUDirect-TCPXO とマルチネットワーキングが有効になっている A3 Mega GKE クラスタの名前。
   • CONTROL_PLANE_LOCATION: クラスタのコントロール プレーンの Compute Engine のロケーション。リージョン クラスタの場合はリージョン、ゾーンクラスタの場合はゾーンを指定します。
   • PROJECT_ID: 実際の Google Cloud プロジェクト ID。

2. 認証に Google Cloud 認証情報を使用するように Google Cloud CLI を構成します。

   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 をデプロイする

トポロジ対応スケジューラを使用すると、指定された GPU トポロジを持つノードに GKE Pod をデプロイできます。

次の 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 をビルドして Google Cloud Artifact Registry に push する

1. Cloud Storage バケットと Docker リポジトリを作成します。scripts/setup-and-configure-resources.sh script で、バケット名とリポジトリ名を作成した名前に置き換えて、スクリプトを実行します。

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

2. pytorch-megatron:23.11-py3 イメージをビルドしてリポジトリに push します。scripts/build-and-push-docker-image.sh ファイル内の Docker リポジトリ名が、scripts/setup-and-configure-resources.sh スクリプトで使用したリポジトリ名と一致していることを確認します。push する前に Docker イメージのタグ名を編集することもできます。

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

   注: このイメージは nvcr.io/nvidia/pytorch:23.11-py3 をベースにしており、変更は最小限です。

Megatron-LM Llama2 ベンチマークを起動する

1. helm/values.yaml ファイルを編集して、Cloud Storage バケットと、前のセクションで作成した Docker イメージを指定します。構成例については、サンプル構成をご覧ください。

2. 省略可: selected-configuration.sh ファイルを編集して、デフォルトの Helm 構成に加えた変更を指定することもできます。

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

   HELM_EXPERIMENT_NAME は、テストの任意の名前に置き換えます。

   注: Helm テストを複数回実行する場合は、helm uninstall コマンドを使用して既存のテストを消去するか、別の名前で新しいテストを作成します。

このテストでは、Nsight Systems プロファイリング ツールからの指標を megatron-experiments ディレクトリで指定された Cloud Storage バケットに書き込みます。

クリーンアップ

このページで使用したリソースについて、 Google Cloud アカウントに課金されないようにするには、 次の手順を実施します。

GKE クラスタを削除します。

[クラスタ] ページに移動します。

[クラスタ] に移動

1. CLUSTER_NAME のチェックボックスをオンにします。
2. [削除] をクリックします。
3. 削除を確定するには、「CLUSTER_NAME」と入力して [削除] をクリックします。

Cloud Storage バケットを削除する

[バケット] ページに移動します。

[バケット] に移動

1. このクイックスタート用に作成した Cloud Storage バケットのチェックボックスをオンにします。

2. [削除] をクリックします。

3. 削除を確定するには、「DELETE」と入力して [削除] をクリックします。

次のステップ

• GKE で GPU の使用方法の詳細を確認する