GKE Autopilot에서 TPU 워크로드 배포

이 페이지에서는 Google Kubernetes Engine(GKE) Autopilot 클러스터에서 Cloud TPU 가속기(TPU)를 사용하여 배포하여 머신러닝(ML) 워크로드를 가속화하는 방법을 설명합니다. 다음 개념에 익숙해야 합니다.

  1. Cloud TPU 소개
  2. Cloud TPU 시스템 아키텍처
  3. GKE의 TPU 정보

Autopilot에서 TPU 작동 방식

Autopilot 워크로드에서 TPU를 사용하려면 워크로드 매니페스트에서 TPU 버전과 해당 TPU 버전에 지원되는 토폴로지를 요청합니다. 그런 다음 Kubernetes resources.requests 및 resources.limits 필드를 사용하여 사용할 워크로드의 TPU 칩 수를 지정합니다. 워크로드를 배포하면 GKE에서 요청된 TPU 구성이 있는 노드를 프로비저닝하고 노드에서 포드를 예약합니다. GKE는 각 워크로드를 자체 노드에 배치하므로 중단 위험이 최소화되면서 각 포드에서 노드의 전체 리소스에 액세스할 수 있습니다.

Autopilot의 TPU는 다음 기능과 호환됩니다.

  1. 스팟 포드
  2. 특정 용량 예약
  3. 연장된 실행 시간 포드

TPU 구성 계획

TPU를 요청하기 전에 워크로드 CPU 및 메모리 요구사항을 기준으로 원하는 구성을 결정합니다. 다음 사항을 결정해야 합니다.

  • TPU 버전: 특정 Cloud TPU 버전(예: v5e)
  • 선택한 TPU 버전의 토폴로지: TPU의 배열 및 수

선택한 TPU 버전과 토폴로지에 따라 GKE에서 노드를 단일 호스트 슬라이스 또는 멀티 호스트 슬라이스로 프로비저닝할지 여부가 결정됩니다. 단일 호스트 슬라이스에서 각 노드는 다른 TPU 슬라이스 노드와 독립적입니다. 멀티 호스트 슬라이스에서 GKE는 TPU 슬라이스에 상호 연결된 VM이 있는 노드 그룹을 만듭니다. 멀티 호스트 슬라이스는 원자적입니다. 즉, GKE는 상호 연결된 전체 노드 그룹을 하나의 단위로 수직 확장 또는 축소합니다.

사용 가능한 TPU 버전, 해당 토폴로지, CPU 및 메모리 용량, 결과 슬라이스 유형에 대한 자세한 내용은 Autopilot TPU 구성 선택을 참조하세요.

가격 책정

가격 책정 정보는 Autopilot 가격 책정을 참조하세요.

시작하기 전에

시작하기 전에 다음 태스크를 수행했는지 확인합니다.

  • Google Kubernetes Engine API를 사용 설정합니다.
    • Google Kubernetes Engine API 사용 설정
  • 이 태스크에 Google Cloud CLI를 사용하려면 gcloud CLI를 설치한 후 초기화합니다. 이전에 gcloud CLI를 설치한 경우 gcloud components update를 실행하여 최신 버전을 가져옵니다.
  • GKE 버전 1.29.2-gke.1521000 이상이 실행되는 Autopilot 클러스터가 있는지 확인합니다.
  • 예약된 TPU를 사용하려면 기존 특정 용량 예약이 있는지 확인합니다. 자세한 내용은 예약된 영역별 리소스 사용을 참조하세요.

TPU 할당량이 있는지 확인

TPU 슬라이스 노드를 만들려면 기존 용량 예약을 사용하지 않는 한 TPU 할당량을 사용할 수 있어야 합니다. 예약된 TPU를 사용하는 경우 이 섹션을 건너뛰세요.

GKE에서 TPU 슬라이스 노드를 만들려면 Cloud TPU API 할당량(tpu.googleapis.com)이 아닌 Compute Engine API 할당량(compute.googleapis.com)이 필요합니다. 할당량 이름은 일반 Autopilot 포드와 스팟 포드에서 서로 다릅니다.

TPU용 Compute Engine API 할당량의 한도 및 현재 사용량을 확인하려면 다음 단계를 따르세요.

  1. Google Cloud 콘솔의 할당량 페이지로 이동합니다.

    할당량으로 이동

  2. 필터 상자에서 다음을 수행합니다,

    1. 서비스 속성을 선택하고 Compute Engine API를 입력한 다음 Enter 키를 누릅니다.

    2. 유형 속성을 선택하고 할당량을 선택합니다.

    3. 이름 속성을 선택하고 다음과 같이 원하는 TPU 유형에 따라 할당량 이름을 입력합니다.

      • TPU v5p(tpu-v5p-slice): TPU v5p 칩
      • TPU v5e(tpu-v5-lite-podslice): TPU v5 Lite PodSlice 칩
      • TPU v5e(tpu-v5-lite-device): TPU v5 Lite Device 칩
      • TPU v4(tpu-v4-podslice): TPU v4 PodSlice 칩

      스팟 포드의 경우 해당 '선점형' 할당량을 선택합니다.

    4. 측정기준(예: 위치) 속성을 선택하고 region:을 입력한 다음 GKE에서 TPU를 생성할 리전의 이름을 입력합니다. 예를 들어 us-west4-a 영역에 TPU 슬라이스 노드를 만들려면 region:us-west4를 입력합니다. TPU 할당량은 리전 기준이므로 동일한 리전 내의 모든 영역에서 동일한 TPU 할당량을 사용합니다.

입력한 필터와 일치하는 할당량이 없으면 원하는 리전에 대해 지정된 할당량이 프로젝트에 부여되지 않으므로 TPU 할당량 상향을 요청해야 합니다.

TPU 애플리케이션 준비

TPU 워크로드 준비 요구사항은 다음과 같습니다.

  1. JAX, PyTorch, TensorFlow와 같은 프레임워크는 libtpu 공유 라이브러리를 사용하여 TPU VM에 액세스합니다. libtpu에는 XLA 컴파일러, TPU 런타임 소프트웨어, TPU 드라이버가 포함됩니다. PyTorch 및 JAX의 각 출시 버전에는 특정 libtpu.so 버전이 필요합니다. GKE에서 TPU를 사용하려면 다음 버전을 사용해야 합니다.
    TPU 유형 libtpu.so 버전
    TPU v5e
    tpu-v5-lite-podslice
    tpu-v5-lite-device
    TPU v5p
    tpu-v5p-slice
    • 권장 jax[tpu] version: 0.4.19 이상
    • 권장 torchxla[tpuvm] 버전: 2023년 10월 23일 나이틀리 버전 빌드를 사용하는 것이 좋습니다
    TPU v4
    tpu-v4-podslice
  2. TPU 리소스를 요청하는 컨테이너에 다음 환경 변수를 설정합니다.
    • TPU_WORKER_ID: 포드마다 고유한 정수입니다. 이 ID는 TPU 슬라이스의 고유한 worker-id를 나타냅니다. 이 필드에 지원되는 값 범위는 0부터 포드 수에서 1을 뺀 값까지입니다.
    • TPU_WORKER_HOSTNAMES: 슬라이스 내에서 서로 통신해야 하는 TPU VM 호스트 이름이나 IP 주소의 쉼표로 구분된 목록입니다. 슬라이스의 TPU VM마다 호스트 이름이나 IP 주소가 있어야 합니다. IP 주소 또는 호스트 이름 목록은 TPU_WORKER_ID에 의해 정렬되고 색인이 0으로 지정됩니다.

      • completionMode: Indexed, subdomain, parallelism > 1, 요청 google.com/tpu 속성을 사용하여 작업을 만든 경우 GKE에서 변형 웹훅을 사용하여 이러한 환경 변수를 자동으로 삽입합니다. GKE는 서비스를 지원하는 포드에 DNS 레코드가 추가되도록 헤드리스 서비스를 추가합니다.

워크로드 준비가 완료되면 TPU를 사용하는 작업을 실행할 수 있습니다.

워크로드에서 TPU 요청

이 섹션에서는 Autopilot에서 TPU를 요청하는 작업을 만드는 방법을 보여줍니다. TPU가 필요한 모든 워크로드에서 다음을 지정해야 합니다.

  • TPU 버전과 토폴로지의 노드 선택기
  • 워크로드의 컨테이너에 대한 TPU 칩 수

지원되는 TPU 버전, 토폴로지, 슬라이스의 해당 TPU 칩 및 노드 수 목록은 Autopilot TPU 구성 선택을 참조하세요.

워크로드에서 TPU 요청 고려사항

포드에 있는 컨테이너 하나만 TPU를 사용할 수 있습니다. 컨테이너에서 요청하는 TPU 칩 수는 슬라이스의 노드에 연결된 TPU 칩 수와 같아야 합니다. 예를 들어 2x4 토폴로지로 TPU v5e(tpu-v5-lite-podslice)를 요청하는 경우 다음 중 하나를 요청할 수 있습니다.

  • 각각 TPU 칩 4개가 있는 멀티 호스트 노드 2개를 만드는 4 칩
  • TPU 칩 8개가 있는 단일 호스트 노드 1개를 만드는 8 칩

경제성을 극대화할 수 있는 권장사항은 항상 요청하는 슬라이스의 모든 TPU를 사용하는 것입니다. 각각 TPU 칩 4개가 있는 노드 2개의 멀티 호스트 슬라이스를 요청하는 경우 두 노드 모두에서 실행되고 슬라이스의 TPU 칩 8개를 모두 사용하는 워크로드를 배포해야 합니다.

TPU를 요청하는 워크로드 만들기

다음 단계에서는 TPU를 요청하는 작업을 만듭니다. 멀티 호스트 TPU 슬라이스에서 실행되는 워크로드가 있으면 이름을 기준으로 워크로드를 선택하는 헤드리스 서비스도 만들어야 합니다. 이 헤드리스 서비스를 사용하면 워크로드의 포드를 가리키도록 Kubernetes DNS 구성을 업데이트하여 멀티 호스트 슬라이스의 서로 다른 노드에 있는 포드가 서로 통신할 수 있습니다.

  1. 다음 매니페스트를 tpu-autopilot.yaml로 저장합니다.

    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_TYPE
        cloud.google.com/gke-tpu-topology: TOPOLOGY
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS

    다음을 바꿉니다.

    • TPU_TYPE: 사용할 TPU 유형입니다(예: tpu-v4-podslice). GKE에서 지원하는 값이어야 합니다.
    • TOPOLOGY: 슬라이스의 TPU 칩 배열입니다(예: 2x2x4). 선택한 TPU 유형에 지원되는 토폴로지여야 합니다.
    • NUMBER_OF_CHIPS: 사용할 컨테이너의 TPU 칩 수입니다. limits 및 requests 값이 같아야 합니다.

  2. 작업을 배포합니다.

    kubectl create -f tpu-autopilot.yaml

이 작업을 만들면 GKE에서 자동으로 다음을 수행합니다.

  1. 포드를 실행할 노드를 프로비저닝합니다. 지정한 TPU 유형, 토폴로지, 리소스 요청에 따라 이러한 노드는 단일 호스트 슬라이스이거나 멀티 호스트 슬라이스입니다.
  2. 다른 워크로드가 TPU 워크로드와 동일한 노드에서 실행되지 않도록 taint를 포드에, 톨러레이션(toleration)을 노드에 추가합니다.

예시: 멀티 호스트 슬라이스의 총 TPU 칩 표시

다음 워크로드는 멀티 호스트 TPU 슬라이스의 모든 노드에서 TPU 칩 수를 반환합니다. 멀티 호스트 슬라이스를 만들기 위한 워크로드 매개변수는 다음과 같습니다.

  • TPU 버전: TPU v4
  • 토폴로지: 2x2x4

이 버전과 토폴로지를 선택하면 멀티 호스트 슬라이스가 생성됩니다.

  1. 다음 매니페스트를 available-chips-multihost.yaml로 저장합니다.
    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
  2. 매니페스트를 배포합니다.
    kubectl create -f available-chips-multihost.yaml

    GKE는 4개의 VM(멀티 호스트 TPU 슬라이스)으로 TPU v4 슬라이스를 실행합니다. 슬라이스에는 상호 연결된 TPU 칩 16개가 있습니다.

  3. 작업으로 4개의 포드가 생성되었는지 확인합니다. 
    kubectl get pods

    출력은 다음과 비슷합니다.

    NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
  4. 포드 중 하나의 로그를 가져옵니다.
    kubectl logs POD_NAME

    POD_NAME을 생성된 포드 중 하나의 이름으로 바꿉니다. 예를 들면 tpu-job-podslice-0-5cd8r입니다.

    출력은 다음과 비슷합니다.

    TPU cores: 16

예시: 단일 노드의 TPU 칩 표시

다음 워크로드는 특정 노드에 연결된 TPU 칩 수를 표시하는 정적 포드입니다. 단일 호스트 노드를 만들기 위한 워크로드 매개변수는 다음과 같습니다.

  • TPU 버전: TPU v5e
  • 토폴로지: 2x4

이 버전과 토폴로지를 선택하면 단일 호스트 슬라이스가 생성됩니다.

  1. 다음 매니페스트를 available-chips-singlehost.yaml로 저장합니다.
    apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8
  2. 매니페스트를 배포합니다.
    kubectl create -f available-chips-singlehost.yaml

    GKE는 TPU v5e를 사용하는 단일 호스트 TPU 슬라이스 8개로 노드를 프로비저닝합니다. 각 TPU 노드에는 8개의 TPU 칩이 있습니다(단일 호스트 TPU 슬라이스).

  3. 포드 로그를 가져옵니다.
    kubectl logs tpu-job-jax-v5

    출력은 다음과 비슷합니다.

    Total TPU chips: 8

관측 가능성 및 측정항목

대시보드

Google Cloud 콘솔의 Kubernetes 클러스터 페이지에서 관측 가능성 탭에는 TPU 관측 가능성 측정항목이 표시됩니다. 자세한 내용은 GKE 관측 가능성 측정항목을 참조하세요.

TPU 대시보드는 GKE 클러스터에서 시스템 측정항목을 사용 설정한 경우에만 채워집니다.

런타임 측정항목

GKE 버전 1.27.4-gke.900 이상에서 JAX 버전 0.4.14 이상을 사용하고 containerPort: 8431을 지정하는 TPU 워크로드는 TPU 사용률 측정항목을 GKE 시스템 측정항목으로 내보냅니다. TPU 워크로드의 런타임 성능을 모니터링하기 위해 Cloud Monitoring에서 다음 측정항목을 사용할 수 있습니다.

  • 가동 주기: 이전 샘플링 기간(60초) 동안 TPU 칩에서 TensorCore가 활발하게 처리된 시간의 백분율입니다. 백분율이 높을수록 TPU 사용률이 높은 것입니다.
  • 메모리 사용량: 할당된 가속기 메모리 양(바이트)입니다. 60초마다 샘플링됩니다.
  • 메모리 용량: 총 가속기 메모리 용량(바이트)입니다. 60초마다 샘플링됩니다.

이러한 측정항목은 Kubernetes 노드(k8s_node) 및 Kubernetes 컨테이너(k8s_container) 스키마에 있습니다.

Kubernetes 컨테이너:

  • kubernetes.io/container/accelerator/duty_cycle
  • kubernetes.io/container/accelerator/memory_used
  • kubernetes.io/container/accelerator/memory_total

Kubernetes 노드:

  • kubernetes.io/node/accelerator/duty_cycle
  • kubernetes.io/node/accelerator/memory_used
  • kubernetes.io/node/accelerator/memory_total

호스트 측정항목

GKE 버전 1.28.1-gke.1066000 이상에서 TPU 슬라이스의 VM은 TPU 사용률 측정항목을 GKE 시스템 측정항목으로 내보냅니다. TPU 호스트의 성능을 모니터링하기 위해 Cloud Monitoring에서 다음 측정항목을 사용할 수 있습니다.

  • TensorCore 사용률: 현재 사용되는 TensorCore의 백분율입니다. TensorCore 값은 행렬 곱셈 단위(MXU)와 벡터 단위의 합과 동일합니다. TensorCore 사용률은 지난 샘플 기간(60초) 동안 수행된 TensorCore 작업을 동일 기간 동안 지원된 TensorCore 작업 수로 나눈 값입니다. 값이 클수록 사용률이 높습니다.
  • 메모리 대역폭 사용률: 사용 중인 가속기 메모리 대역폭의 현재 백분율입니다. 샘플 기간(60초) 동안 사용된 메모리 대역폭을 동일한 샘플 기간 동안 지원된 최대 대역폭으로 나누어 계산합니다.

이러한 측정항목은 Kubernetes 노드(k8s_node) 및 Kubernetes 컨테이너(k8s_container) 스키마에 있습니다.

Kubernetes 컨테이너:

  • kubernetes.io/container/accelerator/tensorcore_utilization
  • kubernetes.io/container/accelerator/memory_bandwidth_utilization

Kubernetes 노드:

  • kubernetes.io/container/node/tensorcore_utilization
  • kubernetes.io/container/node/memory_bandwidth_utilization

자세한 내용은 Kubernetes 측정항목 및 GKE 시스템 측정항목을 참조하세요.

로깅

TPU VM을 포함하여 GKE 노드에서 실행되는 컨테이너에서 내보낸 로그는 GKE Logging 에이전트에 의해 수집되어 Logging으로 전송되고 Logging에 표시됩니다.

Autopilot의 TPU 워크로드 권장사항

다음 권장사항으로 TPU 워크로드 효율성이 향상될 수 있습니다.

  • GKE에서 축소 또는 노드 업그레이드를 위해 포드를 종료하기 전의 유예 기간(최대 7일) 동안 연장된 실행 시간 포드를 사용합니다. 연장된 실행 시간 포드와 함께 유지보수 기간과 제외를 사용하여 자동 노드 업그레이드를 더 지연시킬 수 있습니다.
  • 용량 예약을 사용하면 워크로드가 가용성을 위해 큐에 배치되지 않고 요청된 TPU를 수신할 수 있습니다.