TAS で GKE ワークロードをスケジュールする

クラスタ化された GPU 一般的な GPU

AI / ML ワークロードでは、Pod 間の通信が大量に発生します。この要件により、Pod 間のネットワーク帯域幅はワークロードの実行時間と費用に直接影響します。この帯域幅は、クラスタ内の仮想マシン（VM）インスタンスの配置によって異なります。

このドキュメントでは、Google Kubernetes Engine（GKE）クラスタで大規模な AI または ML ワークロードのスケジューリングを最適化し、パフォーマンスと信頼性の両方を向上させる方法について説明します。具体的には、低レイテンシ通信にトポロジ対応スケジューリング（TAS）を使用するようにクラスタを構成します。このアプローチにより、通信オーバーヘッドが最小限に抑えられ、ワークロードのパフォーマンスを最大化できます。

トポロジ対応スケジューリング（TAS）とは

TAS を使用すると、大規模言語モデル（LLM）トレーニングの効率を大幅に向上させることができます。TAS は、ネットワーク トポロジにワーカーを戦略的に配置して、勾配集約中の通信オーバーヘッドを最小限に抑えます。勾配集約では、ワーカーが特定のランク順で通信する必要があります。TAS は、順次通信するワーカー間のネットワーク ホップを最小限に抑えることで、ネットワーク競合を軽減し、帯域幅の使用率を最適化します。これにより、収束が高速化され、トレーニング時間が短縮されます。LLM モデルが大規模化するにつれて、分散トレーニングのパフォーマンスとスケーラビリティを最大化するために TAS が不可欠になります。

TAS は、予約にバインドされた容量を介して取得できる、高密度に配置された容量で最適に動作します。Spot VM やその他の消費オプションを使用すると、容量が近くに割り当てられる可能性が低くなるため、このシナリオでは TAS が適切に機能しない可能性があります。

始める前に

作業を始める前に、次のタスクが完了していることを確認してください。

• Google Kubernetes Engine API を有効にする。
Google Kubernetes Engine API を有効化
• このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化します。gcloud CLI をインストール済みの場合は、gcloud components update コマンドを実行して最新のバージョンを取得します。以前のバージョンの gcloud CLI では、このドキュメントのコマンドを実行できない場合があります。
  注: gcloud CLI がすでにインストールされている場合には、必ず compute/region プロパティを設定してください。主にゾーンクラスタを使用する場合は、代わりに compute/zone を設定します。デフォルトのロケーションを設定することで、gcloud CLI のエラー（One of [--zone, --region] must be supplied: Please specify location など）を防止できます。クラスタのロケーションが設定したデフォルトと異なる場合は、特定のコマンドでロケーションの指定が必要になることがあります。

• クラスタに接続するには、次のコマンドを実行します。

  gcloud container clusters get-credentials CLUSTER_NAME
  

  CLUSTER_NAME は、使用するクラスタの名前に置き換えます。

GKE クラスタを準備する

TAS を使用してワークロードを実行するように GKE クラスタを準備する手順は次のとおりです。

1. TAS を有効にして Kueue をインストールする

2. GKE クラスタのトポロジを表示する

3. Kueue を構成する

TAS を有効にして Kueue をインストールする

TAS は、割り当てとジョブによる割り当ての使用方法を管理する Kubernetes ネイティブ システムである Kueue と組み合わせて使用することをおすすめします。TAS には Kueue バージョン 0.10.0 以降が必要であり、明示的に有効にする必要があります。

Kueue をインストールして TAS を有効にするには、次のいずれかのオプションを選択します。

Kueue マニフェスト

1. Kueue をインストールします。

   kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/v0.10.0/manifests.yaml
   

2. Kueue で TAS を有効にします。

   kubectl -n kueue-system patch deployment kueue-controller-manager \
       --type json -p='[{"op": "add", "path": "/spec/template/spec/containers/0/args/-", "value": "--feature-gates=TopologyAwareScheduling=true"}]'
   

Helm チャート

Helm チャートを使用して、TAS を有効にして Kueue をインストールします。

helm install kueue oci://us-central1-docker.pkg.dev/k8s-staging-images/charts/kueue \
    --version="v0.10.0" \
    --create-namespace \
    --namespace=kueue-system \
    --set="controllerManager.featureGates[0].name=TopologyAwareScheduling,controllerManager.featureGates[0].enabled=true"

Kueue をインストールしたら、次のセクションで説明するように、管理対象のインフラストラクチャを認識するように Kueue を構成する必要があります。

GKE クラスタのトポロジを表示する

Spot VM としてプロビジョニングされた A4X、A4、A3 Ultra、A3 Mega、A3 High（8 GPU）ノードのトポロジを表示する前に、GKE ノードにコンパクト プレースメントを定義して、TAS の物理トポロジを公開する必要があります。それ以外の場合はエラーが発生します。

特定のノードプール内の GKE クラスタノードのトポロジを表示するには、次のコマンドを実行します。

kubectl get nodes -l cloud.google.com/gke-nodepool=NODE_POOL_NAME \
    -ocustom-columns='NAME:.metadata.name,BLOCK:.metadata.labels.cloud\.google\.com/gce-topology-block,SUBBLOCK:.metadata.labels.cloud\.google\.com/gce-topology-subblock,HOST:.metadata.labels.cloud\.google\.com/gce-topology-host' | sort -k2,4

NODE_POOL_NAME は、ノードプールの名前に置き換えます。

出力で VM 上の GKE ノードの物理トポロジを理解するには、次のノードラベルを参照してください。

• cloud.google.com/gce-topology-block: VM が配置されている予約済みブロックの組織固有の ID。

• cloud.google.com/gce-topology-subblock: VM が配置されているサブブロックの組織固有の ID。

• cloud.google.com/gce-topology-host: VM が配置されているホストの ID。

• kubernetes.io/hostname: Kubernetes ノードのホスト名。通常、このホスト名は GKE ノード名でもあります。

2 つの VM が共有するラベル値が多いほど、VM は物理的に近い場所に配置されます。これらの用語の詳細については、用語をご覧ください。

Kueue を構成する

Kueue をインストールしたら、管理するインフラストラクチャを指定するように Kueue を構成する必要があります。通常、Kueue には、静的インフラストラクチャまたはクラスタの自動スケーリングが有効になっている動的インフラストラクチャの ClusterQueue リソース割り当ての定義が必要です。ClusterQueue は、ワークロードがリクエストするリソースが ClusterQueue で定義されたリソースプール以下の場合にのみ、ワークロードを承認します。このセクションの説明に沿って Kueue を構成すると、Kueue は次のように TAS を使用してワークロードを許可します。

• TAS ワークロード: Kueue は、物理インフラストラクチャのトポロジと現在の使用状況の両方を確認します。

• TAS 以外のワークロード: Kueue は物理インフラストラクチャのトポロジをチェックしません。Kueue は、構成で定義された割り当て全体を管理し、ノードの割り当てを kube-scheduler に任せます。

Kueue に ClusterQueue リソース割り当て定義を指定する方法については、次の例をご覧ください。

• 非常に高い割り当て: Kueue は、リクエストされたリソースに基づいてワークロードの許可を停止することはほとんどありません。TAS の定義に基づいて、Kueue はインフラストラクチャ トポロジに基づいてワークロードを許可する場合と許可しない場合があります。詳細については、非常に高いリソース割り当てをご覧ください。

• 現実的な割り当て: Kueue は、ワークロードがリクエストするリソースがこれらのリソース割り当て上限内にある場合にのみ、ワークロードを許可します。TAS の定義に基づいて、Kueue はワークロードを許可する前にインフラストラクチャ トポロジを確認します。詳細については、現実的なリソース割り当てをご覧ください。

以降のセクションでリソース割り当てに言及している箇所は、すべて ClusterQueue リソース割り当てを指しています。

非常に高いリソース割り当て

次の例では、非常に高いリソース割り当てを使用しているため、Kueue は使用可能なリソース割り当てに基づいてワークロードを停止することはありません。代わりに、Kueue は使用可能なノードのトポロジ情報を使用して、トポロジをワークロードの要件と一致させようとします。

次のリソース割り当て定義を使用するには、次の手順を完了します。

1. 任意のファイル エディタを開きます。次に、kueue-tas-config-very-high-quota.yaml という名前の YAML ファイルに次の割り当て定義を追加します。

     apiVersion: kueue.x-k8s.io/v1alpha1
     kind: Topology
     metadata:
       name: "gke-default"
     spec:
       levels:
       - nodeLabel: "cloud.google.com/gce-topology-block"
       - nodeLabel: "cloud.google.com/gce-topology-subblock"
       - nodeLabel: "cloud.google.com/gce-topology-host"
       - nodeLabel: "kubernetes.io/hostname"
   ---
     kind: ResourceFlavor
     apiVersion: kueue.x-k8s.io/v1beta1
     metadata:
       name: "tas-flavor"
     spec:
       nodeLabels:
         cloud.google.com/gke-nodepool: "NODE_POOL_NAME"
       topologyName: "gke-default"
       tolerations:
       - key: "nvidia.com/gpu"
         operator: "Exists"
         effect: NoSchedule
   ---
     apiVersion: kueue.x-k8s.io/v1beta1
     kind: ClusterQueue
     metadata:
       name: "tas-cluster-queue"
     spec:
       namespaceSelector: {}
       resourceGroups:
       - coveredResources: ["nvidia.com/gpu"]
         flavors:
         - name: "tas-flavor"
           resources:
           - name: "nvidia.com/gpu"
             nominalQuota: 10000000
   ---
     apiVersion: kueue.x-k8s.io/v1beta1
     kind: LocalQueue
     metadata:
       namespace: "default"
       name: "tas-user-queue"
     spec:
       clusterQueue: "tas-cluster-queue"
   

   NODE_POOL_NAME は、ノードプールの名前に置き換えます。

2. Kueue ジョブ キューイング システムのリソース割り当て構成を作成して適用します。

   kubectl create -f kueue-tas-config-very-high-quota.yaml
   

現実的なリソース割り当て

前の例では、GPU リソースのみを構成しました。ただし、Kueue は Kubernetes 互換のすべてのリソースを管理できます。

次の例では、CPU、メモリ、GPU を含む、より現実的なリソース割り当てを定義しています。これは 100 台の a3-ultragpu-8g マシンを対象としています。1 台のマシンに 224 個の vCPU、2,944 GB のメモリ、8 個の GPU が搭載されています。同じ原則が、それぞれの CPU とメモリの仕様を使用する一般 GPU（A2、A3 Edge、G2、G4、N1）マシンタイプにも適用されます。

次のリソース割り当て定義を使用するには、次の手順を完了します。

1. 任意のファイル エディタを開きます。次に、kueue-tas-config-real-quota.yaml という名前の YAML ファイルに次の割り当て定義を追加します。

     apiVersion: kueue.x-k8s.io/v1alpha1
     kind: Topology
     metadata:
       name: "gke-default"
     spec:
       levels:
       - nodeLabel: "cloud.google.com/gce-topology-block"
       - nodeLabel: "cloud.google.com/gce-topology-subblock"
       - nodeLabel: "cloud.google.com/gce-topology-host"
       - nodeLabel: "kubernetes.io/hostname"
   ---
     kind: ResourceFlavor
     apiVersion: kueue.x-k8s.io/v1beta1
     metadata:
       name: "tas-flavor"
     spec:
       nodeLabels:
         cloud.google.com/gke-nodepool: "NODE_POOL_NAME"
       topologyName: "gke-default"
       tolerations:
       - key: "nvidia.com/gpu"
         operator: "Exists"
         effect: NoSchedule
   ---
     apiVersion: kueue.x-k8s.io/v1beta1
     kind: ClusterQueue
     metadata:
       name: "tas-cluster-queue"
     spec:
       namespaceSelector: {} # match all
       resourceGroups:
       - coveredResources: ["cpu", "memory", "nvidia.com/gpu"]
         flavors:
         - name: "tas-flavor"
           resources:
           # numbers below represent quota of 100 a3-ultragpu-8g machines
           - name: "cpu"
             nominalQuota: 22400
           - name: "memory"
             nominalQuota: 294400Gi
           - name: "nvidia.com/gpu"
             nominalQuota: 800
   ---
     apiVersion: kueue.x-k8s.io/v1beta1
     kind: LocalQueue
     metadata:
       namespace: "default"
       name: "tas-user-queue"
     spec:
       clusterQueue: "tas-cluster-queue"
   

   NODE_POOL_NAME は、ノードプールの名前に置き換えます。

2. Kueue ジョブ キューイング システムのリソース割り当て構成を作成して適用します。

   kubectl create -f kueue-tas-config-real-quota.yaml
   

   出力は次のようになります。

   topology.kueue.x-k8s.io/gke-default created
   resourceflavor.kueue.x-k8s.io/tas-flavor created
   clusterqueue.kueue.x-k8s.io/tas-cluster-queue created
   localqueue.kueue.x-k8s.io/tas-user-queue created
   

Kueue を使用して TAS でワークロードをスケジュールする

次のシナリオでは、トポロジ リクエスト タイプとトポロジ リクエスト レベルを使用して、一般的なワークロードとインフラストラクチャの組み合わせを管理するように Kueue と TAS に指示する方法について説明します。

• 使用可能なトポロジ リクエスト タイプ（優先または必須）は次のとおりです。

  • kueue.x-k8s.io/podset-preferred-topology: Kueue は、特定のトポロジレベル内でワークロード全体をスケジュールすることを優先しますが、このトポロジレベルに収まらないワークロードも許可します。単一のトポロジレベルに収まるワークロードの場合、Kueue はそのワークロードをそのトポロジレベルの複数のインスタンスにスケジュールすることがあります。

  • kueue.x-k8s.io/podset-required-topology: Kueue は、ワークロード全体が選択したトポロジレベルに収まるまで、このワークロードの許可を試行し続けます。

• 使用可能なトポロジ リクエスト レベルは次のとおりです。これにより、ジョブの実行を優先または必要とする物理インフラストラクチャをより具体的に指定できます。

  • cloud.google.com/gce-topology-block

  • cloud.google.com/gce-topology-subblock

  • cloud.google.com/gce-topology-host

  • kubernetes.io/hostname

これらの値を使用してワークロードをスケジュールするには、次の Job YAML ファイルを使用します。

apiVersion: batch/v1
kind: Job
metadata:
  generateName: JOB_NAME
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
spec:
  parallelism: NUMBER_OF_REPLICAS
  completions: NUMBER_OF_REPLICAS
  completionMode: Indexed
  template:
    metadata:
      annotations:
        ANNOTATIONS_STRING
    spec:
      containers:
      - name: dummy-job
        image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
        args: ["60s"]
        resources:
          requests:
            nvidia.com/gpu: "1"
          limits:
            nvidia.com/gpu: "1"
      restartPolicy: Never

次の変数を置き換えます。

• JOB_NAME: Job の名前。

• NUMBER_OF_REPLICAS: 並列で実行されている Pod の数。

• ANNOTATIONS_STRING: 次の表をご覧ください。

  リクエストされたトポロジのタイプとレベル	説明	ANNOTATIONS_STRING
  ホスト名内で実行することが推奨されます。	この構成では、容量が断片化されていても、ワークロードのリソース要件を満たすのに十分なリソースが利用可能であれば、ワークロードが許可されます。Kueue は、Pod をできるだけコンパクトにスケジュールします。	kueue.x-k8s.io/podset-preferred-topology: "kubernetes.io/hostname"
  ホスト内で実行する必要がある	この構成では、ワークロードのリソース要件を満たす十分なリソースを備えたホストが利用可能な場合にのみ、ワークロードが許可されます。 これは、ホストごとに複数の VM がある場合（たとえば、より小さいマシンタイプ）や、複数の Pod を単一のノードで実行できる場合に便利です。このような場合、ワークロードが承認されると、単一のホストで実行されます。	kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-host"
  ホスト内で実行することが推奨されます。	この構成では、容量が断片化されていても、ワークロードのリソース要件を満たすのに十分なリソースが利用可能であれば、ワークロードが許可されます。Kueue は、ホスト内で Pod をスケジュールしようとします。必要に応じて、追加のホストを使用します。	kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-host"
  サブブロック内で実行するには必須	この構成では、ワークロードのリソース要件を満たすのに十分なリソースを持つサブブロックが利用可能な場合にのみ、ワークロードが許可されます。	kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-subblock"
  サブブロック内で実行することが推奨されます。	この構成では、容量が断片化されていても、ワークロードのリソース要件を満たすのに十分なリソースが利用可能であれば、ワークロードが許可されます。Kueue は、サブブロック内で Pod をスケジュールしようとします。必要に応じて、追加のサブブロックを使用します。この場合、Kueue は、要件を満たすのに十分な容量を持つサブブロックよりも、断片化されていても利用可能な容量が多いサブブロックを優先します。	kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-subblock"
  ブロック内で実行する必須	この構成では、ブロック内で使用可能なリソースがワークロードのリソース要件を満たしている場合にのみ、ワークロードが許可されます。承認された場合、Kueue はワークロードをスケジュールするサブブロックとホストの数を最小限に抑えます。これにより、使用可能な容量が断片化する可能性があります。	kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
  ブロック内で実行することが推奨	この構成では、容量が断片化されている場合でも、ワークロードのリソース要件を満たすのに十分なリソースが利用可能であれば、ワークロードが許可されます。Kueue は、ブロック内で Pod をスケジュールしようとします。必要に応じて、追加のブロックを使用します。	kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-block"

Kueue を使用して TAS で PodGroup を使用してワークロードをスケジュールする

PodGroup を使用する場合は、PodGroup 内のすべての Pod に 3 つの追加フィールドを指定する必要があります。

• ラベル:

  • kueue.x-k8s.io/pod-group-name: 集計に使用される PodGroup の名前。

  • kueue.x-k8s.io/pod-group-pod-index: PodGroup 内の各 Pod のインデックス。

• アノテーション:

  • kueue.x-k8s.io/pod-group-total-count: PodGroup 内の Pod の合計数。

使用する ML フレームワークによっては、PodGroup のリーダーが GPU を必要とする場合と必要としない場合があります。Kueue の制限により、これらのケースは異なる方法で処理する必要があります。次の例は、1 つのリーダーと 2 つのワーカーを持つ 3 つの Pod の PodGroup を作成する方法を示しています。

ケース 1: リーダーがワーカーでもあり、GPU が必要

リーダーがワーカーの 1 つであり、GPU も必要な場合、リーダーは PodGroup 内で任意の数を持つことができます。わかりやすくするために、次の例ではリーダーのインデックスは 0 です。

apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-leader-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "0"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  containers:
  - name: leader
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
  restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-1-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "1"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-2-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "2"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

ケース 2: リーダーがワーカーではなく、GPU を必要としない

Kueue の制限によりリーダーがワーカーの 1 つでない場合、Kueue が PodSet を作成する方法により、リーダーは PodGroup の最後のインデックスを持つ必要があります。リーダーに最後のインデックスがなく、最初のワーカーが最初のインデックスを使用しない場合、Kueue はランク割り当てを適用しません。

次の例をご覧ください。

---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-leader-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "2"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  containers:
  - name: leader
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        cpu: "1"
      limits:
        cpu: "1"
  restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-0-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "0"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-1-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "1"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

次のステップ

• GKE クラスタでノードの健全性予測を有効にする方法については、ノードの健全性予測を有効にするをご覧ください。

• GKE クラスタと AI ワークロードに関連する一般的なイベントの管理については、AI に最適化された GKE クラスタを管理するをご覧ください。

• Kueue を使用して GKE でジョブをスケジュールする方法については、Kueue を使用してバッチシステムをデプロイするをご覧ください。

• 事前構成された Kueue と TAS のブループリントを使用して GKE クラスタをデプロイするには、Cluster Toolkit の GKE デプロイの概要をご覧ください。