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

環境を設定する

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: 実際の 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 をデプロイする

トポロジ対応スケジューラを使用すると、指定された 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 をビルドして Cloud de Confiance by S3NS 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
   

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 は、テストの任意の名前に置き換えます。

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