トポロジ対応スケジューリング（TAS）で GKE ワークロードをスケジュールする

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

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

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

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

TAS は、予約によって取得できる密度の高い容量で最適に動作します。Flex Start VM または Spot VM では、容量が近くに割り当てられる可能性が低いため、このシナリオでは TAS がうまく機能しない可能性があります。

始める前に

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

  • Google Kubernetes Engine API を有効にする。
    • Google Kubernetes Engine API を有効化
  • このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。gcloud CLI をインストール済みの場合は、gcloud components update コマンドを実行して最新のバージョンを取得します。以前のバージョンの gcloud CLI では、このドキュメントのコマンドを実行できない場合があります。

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

    gcloud container clusters get-credentials CLUSTER_NAME

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

GKE クラスタを準備する

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

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

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

  3. Kueue を構成する

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

割り当てとジョブによる割り当ての使用方法を管理する Kubernetes ネイティブ システムである Kueue で TAS を使用することをおすすめします。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"

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

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

スポット VM としてプロビジョニングされた A4X、A4、A3 Ultra、A3 Mega、A3 High（GPU 数 8）ノードのトポロジを表示する前に、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 に任せます。

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

  • 非常に高い割り当て: 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 があります。

次のリソース割り当て定義を使用する手順は次のとおりです。

  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 つのフィールドを追加で指定する必要があります。

使用する 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"

次のステップ