GKE에서 TPU 멀티슬라이스 배포

이 페이지에서는 비용 효율적인 대규모 학습을 위해 Cloud TPU 멀티슬라이스 구성을 사용하여 Google Kubernetes Engine(GKE)에서 워크로드를 배포하는 방법을 보여줍니다.

이 튜토리얼은 Kubernetes 컨테이너 조정을 사용하여 TPU를 사용하는 대규모 모델 학습, 튜닝, 추론 워크로드를 관리하려는 머신러닝 (ML) 엔지니어, 플랫폼 관리자, 운영자를 대상으로 합니다. Google Cloud 콘텐츠에서 참조하는 일반적인 역할 및 예시 태스크에 대해 자세히 알아보려면 일반 GKE Enterprise 사용자 역할 및 태스크를 참조하세요.

GKE에서 멀티슬라이스를 구성하려면 다음 개념을 숙지해야 합니다.

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

TPU 멀티슬라이스란?

TPU 멀티슬라이스는 Cloud TPU 슬라이스 2개 이상이 데이터 센터 네트워크(DCN)를 통해 통신하는 TPU 슬라이스의 VM 아키텍처 구성입니다. 멀티슬라이스는 TPU 칩을 최대 10,000개까지 선형에 가깝게 수직 확장하여 경제적인 풀 스택 대규모 학습을 지원합니다. 멀티슬라이스 구성에서 GKE는 여러 TPU 슬라이스에 멀티슬라이스 워크로드를 배포합니다. 슬라이스 내에서 TPU 칩 간 통신은 내부 칩 상호 연결(ICI)을 통해 수행됩니다. 슬라이스 간 통신은 DCN을 통해 수행됩니다.

작업이 너무 커서 단일 TPU 슬라이스에 맞지 않으면 멀티슬라이스를 사용하는 것이 좋습니다.

GKE의 멀티슬라이스 가용성

  • Standard는 1.27.4-gke.900 이상 버전에서 멀티슬라이스를 지원합니다.
  • Autopilot은 1.29.2-gke.1521000 이상 버전에서 멀티슬라이스를 지원합니다.
  • 멀티슬라이스는 JAX 및 PyTorch 프레임워크를 지원합니다. 지원되는 최소 JAX 버전은 2.1입니다.
  • 멀티슬라이스는 멀티 호스트 TPU 슬라이스 노드 풀만 지원합니다. 예를 들어 2x2x1 토폴로지의 ct4p-hightpu-4t 또는 2x2 토폴로지의 ct5lp-hightpu-4t에는 단일 호스트 TPU 슬라이스 노드 풀이기 때문에 멀티슬라이스를 사용할 수 없습니다.
  • 멀티슬라이스는 동기식 멀티 컨트롤러 학습만 지원합니다.
  • 멀티슬라이스 워크로드는 동일한 TPU 유형, 크기, 토폴로지를 공유하는 TPU 슬라이스에서만 실행될 수 있습니다.
  • 멀티슬라이스는 TPU v3를 지원하지 않습니다.

시작하기 전에

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

  • Google Kubernetes Engine API를 사용 설정합니다.
  • Google Kubernetes Engine API 사용 설정
  • 이 태스크에 Google Cloud CLI를 사용하려면 gcloud CLI를 설치한 후 초기화합니다. 이전에 gcloud CLI를 설치했으면 gcloud components update 명령어를 실행하여 최신 버전을 가져옵니다. 이전 gcloud CLI 버전에서는 이 문서의 명령어를 실행하지 못할 수 있습니다.

GKE 구현 방법 선택

Google Cloud CLI 또는 가속 처리 키트 (XPK)를 사용하여 TPU를 사용하는 클러스터를 만들 수 있습니다.

  • gcloud CLI를 사용하여 GKE 클러스터 인스턴스를 수동으로 만들어 기존 프로덕션 GKE 환경을 정확하게 맞춤설정하거나 확장합니다.
  • XPK를 사용하여 GKE 클러스터를 빠르게 만들고 개념 증명 및 테스트를 위한 워크로드를 실행합니다. 자세한 내용은 XPK README를 참고하세요.

GKE 노드 풀을 수동으로 만들기

이 섹션에서는 멀티슬라이스에서 워크로드를 실행하기 위해 Standard 모드 노드 풀을 준비하는 방법을 보여줍니다. GKE Autopilot 모드를 사용하는 경우 멀티슬라이스 워크로드 실행 섹션으로 건너뜁니다. 버전 1.29.2-gke.1521000 이상을 실행하는 Autopilot 클러스터는 기본적으로 TPU를 사용 설정합니다.

이 섹션에서는 다음 단계에 대해 설명합니다.

  1. 3개의 멀티 호스트 TPU 슬라이스 노드 풀 만들기
  2. 노드 풀 상태 확인
  3. 멀티슬라이스 워크로드 실행

TPU 슬라이스 노드 풀 만들기

멀티 호스트 TPU 노드 풀을 2개 이상 만들 수 있습니다. 이 가이드에서는 멀티 호스트 TPU 슬라이스 노드 풀을 3개 만들어 멀티슬라이스 워크로드를 실행합니다. 멀티 호스트 TPU 슬라이스 노드 풀을 만드는 단계는 Ironwood (TPU7x)를 사용하는지 또는 이전 TPU 버전을 사용하는지에 따라 다릅니다.

Ironwood (TPU7x)

Google Cloud CLI 또는 Terraform을 사용하여 Ironwood (TPU7x) 버전에서 멀티 호스트 TPU 슬라이스 노드 풀을 만들 수 있습니다.

gcloud

Ironwood (TPU7x)로 멀티 호스트 TPU 슬라이스 노드 풀을 만들려면 먼저 워크로드 정책을 만들어야 합니다.

  1. 워크로드 정책을 만듭니다.

    gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
        --type=HIGH_THROUGHPUT \
        --accelerator-topology=TPU_TOPOLOGY \
        --project=PROJECT_ID \
        --region=REGION
    

    다음을 바꿉니다.

    • WORKLOAD_POLICY_NAME: 워크로드 정책의 이름입니다.
    • TPU_TOPOLOGY: TPU Ironwood (TPU7x) 토폴로지입니다. 예를 들면 2x2x2입니다. 지원되는 모든 Ironwood (TPU7x) 토폴로지를 확인하려면 토폴로지 섹션을 참고하세요.
    • PROJECT_ID: Google Cloud 프로젝트 ID입니다.
    • REGION: 워크로드 정책의 리전입니다. 워크로드 정책은 리전 리소스이며 동일한 토폴로지를 공유하는 노드 풀에서 재사용할 수 있습니다.
  2. 워크로드 정책을 사용하여 노드 풀을 만듭니다.

    gcloud container node-pools create NODE_POOL_NAME \
      --cluster=CLUSTER_NAME \
      --machine-type=tpu7x-standard-4t \
      --placement-policy=WORKLOAD_POLICY_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --node-locations=NODE_ZONE \
      --project=PROJECT_ID \
      --reservation=RESERVATION_NAME \
      --reservation-affinity=specific
    

    다음을 바꿉니다.

    • NODE_POOL_NAME: 새 노드 풀의 이름입니다.
    • CLUSTER_NAME: GKE 클러스터의 이름입니다.
    • WORKLOAD_POLICY_NAME: 생성한 워크로드 정책의 이름입니다.
    • CONTROL_PLANE_LOCATION: 클러스터의 컨트롤 플레인에 대한 Compute Engine 위치입니다. 리전 클러스터의 경우 리전 또는 영역 클러스터의 경우 영역을 제공합니다.
    • NODE_ZONE: 사용할 TPU 버전에 따른 영역의 이름.

      선택사항: us-central1-ai1a와 같은 AI 영역을 사용합니다. AI 영역은 Google Cloud 리전 내에서 AI/ML 워크로드에 최적화된 특수 위치입니다.

      사용 가능한 위치를 확인하려면 GKE의 TPU 가용성을 참조하세요.
    • PROJECT_ID: Google Cloud 프로젝트 ID입니다.
    • RESERVATION_NAME: 사용할 예약의 이름입니다.

    이 명령어에서 --tpu-topology 플래그가 --placement-policy 플래그로 대체되었습니다.

Terraform

  1. google 제공업체 버전 4.84.0 이상을 사용해야 합니다.
  2. 워크로드 정책을 만듭니다.

    resource "google_compute_resource_policy" {
      name   = "WORKLOAD_POLICY_NAME"
      region = CLUSTER_LOCATION
      workload_policy {
        type = "HIGH_THROUGHPUT"
        accelerator_topology = "TPU_TOPOLOGY"
      }
    }
    

    다음을 바꿉니다.

    • WORKLOAD_POLICY_NAME: 워크로드 정책의 이름입니다.
    • CLUSTER_LOCATION: 클러스터의 컴퓨팅 위치. Kubernetes 컨트롤 플레인의 신뢰성 향상을 위해 리전 클러스터를 사용하는 것이 좋습니다. 영역 클러스터를 사용할 수도 있습니다. 자세한 내용은 TPU 버전 및 토폴로지 선택을 참고하세요.
    • TPU_TOPOLOGY: TPU Ironwood (TPU7x) 토폴로지입니다. 예를 들면 2x2x2입니다. 지원되는 모든 Ironwood (TPU7x) 토폴로지를 확인하려면 TPU 계획을 참고하세요.

    google_compute_resource_policy 참조에 대한 자세한 내용은 Terraform 프로바이더를 참고하세요.

  3. Terraform 구성에서 다음 블록을 추가합니다.

    resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
      provider           = google
      project            = PROJECT_ID
      cluster            = CLUSTER_NAME
      name               = POOL_NAME
      location           = CLUSTER_LOCATION
      node_locations     = [NODE_ZONES]
      initial_node_count = NUM_NODES
    
      autoscaling {
        max_node_count = MAX_NODES
        location_policy      = "ANY"
      }
      node_config {
        machine_type = MACHINE_TYPE
        reservation_affinity {
          consume_reservation_type = "SPECIFIC_RESERVATION"
          key = "compute.googleapis.com/reservation-name"
          values = [RESERVATION_LABEL_VALUES]
        }
        flex_start = false
        spot = true
      }
    
      placement_policy {
        policy_name = WORKLOAD_POLICY_NAME
      }
    }
    

    다음을 바꿉니다.

    • NODE_POOL_RESOURCE_NAME: Terraform 템플릿의 노드 풀 리소스 이름입니다.
    • PROJECT_ID: 프로젝트 ID입니다.
    • CLUSTER_NAME: 노드 풀을 추가할 기존 클러스터의 이름입니다.
    • POOL_NAME: 만들 노드 풀의 이름입니다.
    • NODE_ZONES: GKE가 노드 풀을 만드는 하나 이상의 영역을 쉼표로 구분한 목록

      선택사항: us-central1-ai1a와 같은 AI 영역을 사용합니다. AI 영역은 Google Cloud 리전 내에서 AI/ML 워크로드에 최적화된 특수 위치입니다.

    • NUM_NODES: 노드 풀의 노드 수입니다. 0이거나 TPU 칩 수의 곱을 4(멀티 호스트 TPU 슬라이스의 TPU 슬라이스 노드마다 칩이 4개 있기 때문)로 나눈 값이어야 합니다. 예를 들어 TPU_TOPOLOGY4x8이면 칩이 32개 있으므로 NUM_NODES는 8이어야 합니다. TPU 토폴로지에 대한 자세한 내용은 TPU 버전 선택의 표를 참조하세요.
    • TPU_TOPOLOGY: TPU 슬라이스에 대해 선택된 물리적 토폴로지를 나타냅니다. 토폴로지 형식은 사용 중인 TPU 버전에 따라 달라집니다. TPU 토폴로지에 대한 자세한 내용은 토폴로지 선택의 표를 참조하세요.

    선택적으로 다음 변수를 사용할 수도 있습니다.

    • RESERVATION_NAME: TPU 예약을 사용하는 경우 노드 풀을 만들 때 사용할 예약 리소스 라벨 목록을 제공합니다. reservation_affinity 필드에 RESERVATION_LABEL_VALUES를 채우는 방법에 대한 자세한 내용은 Terraform 제공업체를 참조하세요.
    • autoscaling: 자동 확장이 사용 설정된 노드 풀을 만듭니다. GKE는 멀티 호스트 TPU 슬라이스 노드 풀을 확장할 때 노드 풀을 0에서 최대 크기로 원자적으로 수직 확장합니다.
      • MAX_NODES: 노드 풀의 최대 크기입니다. 이 값은 TPU_TOPOLOGY에 정의된 값의 곱 ({A}x{B}x{C})을 각 VM의 칩 수로 나눈 값과 같아야 합니다. 예를 들어 TPU_TOPOLOGY2x2x2이면 곱은 8입니다. tpu7x-standard-4t의 각 VM에는 4개의 칩이 있으므로 노드 수는 2입니다.
    • spot: TPU 슬라이스 노드에 스팟 VM을 사용할 노드 풀입니다. 노드 풀을 만든 후에는 이 설정을 변경할 수 없습니다. 자세한 내용은 스팟 VM을 참조하세요.
    • flex_start: flex-start 소비 옵션을 사용할 노드 풀입니다. spot이 사용 설정된 경우 이 설정을 true로 설정할 수 없습니다.

기타 TPU 버전

Google Cloud CLI, Terraform 또는 Google Cloud 콘솔을 사용하여 v3, v4, v5p, v5e, Trillium (v6e) 버전에서 멀티 호스트 TPU 슬라이스 노드 풀을 만들 수 있습니다.

gcloud

  gcloud container node-pools create POOL_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --cluster=CLUSTER_NAME \
      --node-locations=NODE_ZONE \
      --machine-type=MACHINE_TYPE \
      --tpu-topology=TPU_TOPOLOGY \
      [--num-nodes=NUM_NODES] \
      [--spot \]
      [--flex-start \]
      [--enable-autoscaling \
        --max-nodes MAX_NODES]
      [--reservation-affinity=specific \
      --reservation=RESERVATION_NAME] \
      [--node-labels cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME,cloud.google.com/gke-workload-type=HIGH_AVAILABILITY]
      [--placement-type=COMPACT]

다음을 바꿉니다.

  • POOL_NAME: 새 노드 풀의 이름
  • CONTROL_PLANE_LOCATION: 사용할 TPU 버전에 따른 영역의 이름. 사용 가능한 위치를 확인하려면 GKE의 TPU 가용성을 참조하세요.
  • CLUSTER_NAME: 클러스터의 이름
  • NODE_ZONES: GKE가 노드 풀을 만드는 하나 이상의 영역을 쉼표로 구분한 목록

    선택사항: us-central1-ai1a와 같은 AI 영역을 사용합니다. AI 영역은 Google Cloud 리전 내에서 AI/ML 워크로드에 최적화된 특수 위치입니다.

  • MACHINE_TYPE: 노드에 사용할 머신 유형. 사용 가능한 머신 유형에 대해 자세히 알아보려면 TPU 버전 선택을 참조하세요.
  • TPU_TOPOLOGY: TPU 슬라이스의 물리적 토폴로지입니다. 토폴로지 형식은 TPU 버전에 따라 달라집니다. TPU 토폴로지에 대한 자세한 내용은 토폴로지 선택의 표를 참조하세요.

    자세한 내용은 토폴로지를 참고하세요.

    선택적으로 다음 플래그를 사용할 수도 있습니다.

  • NUM_NODES: 노드 풀의 노드 수입니다. 0이거나, TPU_TOPOLOGY에 정의된 값의 곱({A}x{B}x{C})을 각 VM의 칩 수로 나눈 값이어야 합니다. 멀티 호스트 TPU v4 및 TPU v5e의 경우 각 VM의 칩 수는 4개입니다. 따라서 TPU_TOPOLOGY2x4x4(각 VM에 4개의 칩이 포함된 TPU v4)인 경우 NUM_NODES는 8과 동일한 32/4입니다. 이 플래그를 생략하면 토폴로지와 머신 유형에 따라 노드 수가 계산되고 기본값이 지정됩니다.

  • RESERVATION_NAME: 노드 풀을 만들 때 GKE가 사용하는 예약의 이름입니다. 이 플래그를 생략하면 GKE에서 사용 가능한 TPU 슬라이스 노드 풀을 사용합니다. TPU 예약에 대한 자세한 내용은 TPU 예약을 참고하세요.

  • --spot: TPU 슬라이스 노드에 스팟 VM을 사용하도록 노드 풀을 설정합니다. 노드 풀을 만든 후에는 변경할 수 없습니다. 자세한 내용은 스팟 VM을 참조하세요.

  • --flex-start: Flex-start VM을 사용하도록 노드 풀을 설정합니다. Flex-start VM은 GKE 버전 1.33.0-gke.1712000 이상에서 지원되는 flex-start 소비 옵션을 사용하여 생성됩니다.

  • --enable-autoscaling: 자동 확장이 사용 설정된 노드 풀을 만듭니다. GKE는 멀티 호스트 TPU 슬라이스 노드 풀을 확장할 때 노드 풀을 0에서 최대 크기로 원자적으로 수직 확장합니다.

    • MAX_NODES: 노드 풀의 최대 크기입니다. --enable-autoscaling이 제공되는 경우 --max-nodes 플래그가 필요하며 TPU_TOPOLOGY({A}x{B}x{C})에 정의된 값의 곱을 각 VM의 칩 수로 나눈 값과 같아야 합니다.
  • --node-label=cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME, cloud.google.com/gke-workload-type=HIGH_AVAILABILITY: 멀티 호스트 TPU 슬라이스 노드 풀이 컬렉션임을 GKE에 알립니다. 다음 조건이 적용되는 경우 이 플래그를 사용하세요.

    • 노드 풀은 추론 워크로드를 실행합니다.
    • 노드 풀에서 TPU Trillium을 사용합니다.
    • 스팟 VM은 컬렉션 예약을 지원하지 않습니다.

    컬렉션 예약 관리에 대해 자세히 알아보려면 멀티 호스트 TPU 슬라이스에서 컬렉션 예약 관리를 참고하세요.

  • --placement-type=COMPACT: 압축 배치가 사용 설정된 노드 풀을 만듭니다. 이 옵션은 --tpu-topology 플래그와 함께 사용해야 합니다. 자세한 내용은 압축 배치 정책 만들기TPU 토폴로지를 참조하세요.

Terraform

  1. google 제공업체 버전 4.84.0 이상을 사용해야 합니다.
  2. Terraform 구성에 다음 블록을 추가합니다.

    resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
      provider           = google
      project            = PROJECT_ID
      cluster            = CLUSTER_NAME
      name               = POOL_NAME
      location           = CLUSTER_LOCATION
      node_locations     = [NODE_ZONES]
      initial_node_count = NUM_NODES
    
      autoscaling {
        max_node_count = MAX_NODES
        location_policy      = "ANY"
      }
      node_config {
        machine_type = MACHINE_TYPE
        reservation_affinity {
          consume_reservation_type = "SPECIFIC_RESERVATION"
          key = "compute.googleapis.com/reservation-name"
          values = [RESERVATION_LABEL_VALUES]
        }
        flex_start = false
        spot = true
      }
    
      placement_policy {
        type = "COMPACT"
        tpu_topology = TPU_TOPOLOGY
      }
    }
    

    다음을 바꿉니다.

    • NODE_POOL_RESOURCE_NAME: Terraform 템플릿의 노드 풀 리소스 이름입니다.
    • PROJECT_ID: 프로젝트 ID입니다.
    • CLUSTER_NAME: 노드 풀을 추가할 기존 클러스터의 이름입니다.
    • POOL_NAME: 만들 노드 풀의 이름입니다.
    • CLUSTER_LOCATION: 클러스터의 컴퓨팅 위치입니다. Kubernetes 컨트롤 플레인의 신뢰성 향상을 위해 리전 클러스터를 사용하는 것이 좋습니다. 영역 클러스터를 사용할 수도 있습니다. 자세한 내용은 TPU 버전 및 토폴로지 선택을 참조하세요.
    • NODE_ZONES: GKE가 노드 풀을 만드는 하나 이상의 영역을 쉼표로 구분한 목록

      선택사항: us-central1-ai1a와 같은 AI 영역을 사용합니다. AI 영역은 Google Cloud 리전 내에서 AI/ML 워크로드에 최적화된 특수 위치입니다.

    • NUM_NODES: 노드 풀의 노드 수입니다. 0이거나 TPU 칩 수의 곱을 4(멀티 호스트 TPU 슬라이스의 TPU 슬라이스 노드마다 칩이 4개 있기 때문)로 나눈 값이어야 합니다. 예를 들어 TPU_TOPOLOGY4x8이면 칩이 32개 있으므로 NUM_NODES는 8이어야 합니다. TPU 토폴로지에 대한 자세한 내용은 TPU 버전 선택의 표를 참조하세요.
    • TPU_TOPOLOGY: TPU 슬라이스의 물리적 토폴로지를 나타냅니다. 토폴로지 형식은 사용 중인 TPU 버전에 따라 달라집니다. TPU 토폴로지에 대한 자세한 내용은 토폴로지 선택의 표를 참조하세요.

    선택적으로 다음 변수를 사용할 수도 있습니다.

    • RESERVATION_NAME: TPU 예약을 사용하는 경우 노드 풀을 만들 때 사용할 예약 리소스의 라벨 목록입니다. reservation_affinity 필드에 RESERVATION_LABEL_VALUES를 채우는 방법에 대한 자세한 내용은 Terraform 제공업체를 참조하세요.
    • autoscaling: 자동 확장이 사용 설정된 노드 풀을 만듭니다. GKE는 멀티 호스트 TPU 슬라이스 노드 풀을 확장할 때 노드 풀을 0에서 최대 크기로 원자적으로 수직 확장합니다.
      • MAX_NODES: 노드 풀의 최대 크기입니다. TPU_TOPOLOGY에 정의된 값의 곱 ({A}x{B}x{C})을 각 VM의 칩 수로 나눈 값이어야 합니다.
    • spot: 노드 풀에서 TPU 슬라이스 노드에 스팟 VM을 사용할 수 있게 합니다. 노드 풀을 만든 후에는 변경할 수 없습니다. 자세한 내용은 스팟 VM을 참조하세요.
    • flex_start: flex-start 소비 옵션을 사용하도록 노드 풀을 설정합니다. spot이 사용 설정된 경우 true로 설정할 수 없습니다.

콘솔

TPU가 포함된 노드 풀을 만들려면 다음 안내를 따르세요.

  1. Google Cloud 콘솔에서 Google Kubernetes Engine 페이지로 이동합니다.

    Google Kubernetes Engine으로 이동

  2. 클러스터 목록에서 수정하려는 클러스터 이름을 클릭합니다.

  3. 노드 탭을 클릭합니다.

  4. 사용자 관리형 노드 풀 만들기를 클릭합니다.

  5. 노드 풀 세부정보 섹션에서 노드 위치 지정 체크박스를 선택합니다.

  6. 사용할 TPU 버전에 따른 영역의 이름. 사용 가능한 위치를 확인하려면 GKE의 TPU 가용성을 참조하세요.

  7. 탐색창에서 노드를 클릭합니다.

  8. 머신 구성 섹션에서 TPU를 선택합니다.

  9. 시리즈 드롭다운 메뉴에서 다음 중 하나를 선택합니다.

    • CT3: TPU v3, 단일 호스트 기기
    • CT3P: TPU v3, 멀티 호스트 포드 슬라이스
    • CT4P: TPU v4
    • CT5LP: TPU v5e
    • CT5P: TPU v5p
    • CT6E: TPU Trillium(v6e)
  10. 머신 유형 드롭다운 메뉴에서 노드에 사용할 머신 이름을 선택합니다. TPU 버전 선택 표를 사용하여 멀티 호스트 TPU 슬라이스 노드 풀을 만드는 머신 유형 및 TPU 토폴로지를 정의하는 방법을 알아봅니다.

  11. TPU 토폴로지 드롭다운 메뉴에서 TPU 슬라이스의 물리적 토폴로지를 선택합니다.

  12. 필요한 변경사항 대화상자에서 변경을 클릭합니다.

  13. 부팅 디스크 유형표준 영구 디스크 또는 SSD 영구 디스크인지 확인합니다.

  14. 선택적으로 노드 풀의 노드에 스팟 VM을 사용하려면 스팟 VM에서 노드 사용 설정 체크박스를 선택합니다.

  15. 만들기를 클릭합니다.

노드 풀 상태 확인

  1. kubectl을 사용해 클러스터에 액세스할 수 있도록 사용자 인증 정보를 받습니다.

    gcloud container clusters get-credentials CLUSTER_NAME \
        --project=PROJECT_ID \
        --location=CONTROL_PLANE_LOCATION
    

    다음을 바꿉니다.

    • CLUSTER_NAME: 클러스터의 이름입니다.
    • PROJECT_ID: 프로젝트 ID입니다.
    • CONTROL_PLANE_LOCATION: 클러스터의 컨트롤 플레인에 대한 Compute Engine 위치입니다. 리전 클러스터의 경우 리전 또는 영역 클러스터의 경우 영역을 제공합니다.
  2. Cloud Shell에서 kubectl을 사용하여 TPU 슬라이스 노드를 확인합니다.

    kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=ACCELERATOR_TYPE \
       -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGY
    

    다음을 바꿉니다.

    • ACCELERATOR_TYPE: TPU 슬라이스 노드 풀의 가속기 유형입니다. 이 값은 노드 풀 생성 중에 사용한 머신 유형에 매핑되어야 합니다. 예를 들어 ct5lp-hightpu-4t 머신 유형으로 노드 풀을 만든 경우 tpu-v5-lite-podslice를 사용합니다. 사용할 가속기 유형을 알아보려면 GKE에서 TPU 계획의 표를 참고하세요.
    • TPU_TOPOLOGY: TPU 슬라이스의 물리적 토폴로지입니다.

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

     NAME                                    STATUS   ROLES    AGE    VERSION
     gke-tpu-20ee2cce-5tv6                   Ready    <none>   34h     v1.28.1-gke.1066000
    

멀티슬라이스 워크로드 실행

이 섹션에서는 TPU 슬라이스의 전역 TPU 칩 수를 표시하고 종료되는 JAX 워크로드를 실행합니다.

JAX 워크로드를 실행하려면 다음을 수행합니다.

  1. 다음 tpu-multislice.yaml 매니페스트를 만듭니다.

    Autopilot

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: multislice-job
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: 2 # This value defines a "Multislice" setup, by creating 2 independent TPU slices.
          template:
            spec:
              parallelism: 4
              completions: 4
              backoffLimit: 0
              template:
                spec:
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                    cloud.google.com/gke-tpu-topology: 4x4
                  containers:
                  - name: jax-tpu
                    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    - containerPort: 8431
                    command:
                    - bash
                    - -c
                    - |
                      python -c 'import jax; print("Global device count:", jax.device_count())'
                      sleep 60
                    resources:
                     limits:
                        google.com/tpu: 4 # This container requires 4 TPU chips
    

    이 매니페스트에서는 4x4 토폴로지를 사용하여 tpu-v5-lite-podslice를 선택합니다. 이 구성은 다음을 의미합니다.

    • 토폴로지에는 총 16개의 칩이 있습니다 (4*4=16).
    • tpu-v5-lite-podslice는 VM당 칩이 4개인 ct5lp-hightpu-4t 머신에 배포됩니다. 즉, 4x4 토폴로지를 사용하려면 VM이 4개 (16/4=4) 필요합니다. 따라서 parallelismcompletions 필드를 4로 설정해야 합니다.
    • 포드의 각 컨테이너는 4개의 TPU 칩을 요청하며 이는 VM당 칩 수에 해당합니다. 이는 google.com/tpu: 4 줄을 사용하여 매니페스트에 지정됩니다.
    • replicas 필드는 2으로 설정되어 슬라이스 두 개를 만듭니다.

    이러한 변수의 가능한 값에 대한 자세한 내용은 GKE에서 TPU 계획의 표를 참고하세요.

    표준

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: multislice-job
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: NUM_SLICES
          template:
            spec:
              parallelism: NUM_NODES
              completions: NUM_NODES
              backoffLimit: 0
              template:
                spec:
                  hostNetwork: true
                  dnsPolicy: ClusterFirstWithHostNet
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                    cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
                  containers:
                  - name: jax-tpu
                    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    - containerPort: 8431
                    securityContext:
                      privileged: true
                    command:
                    - bash
                    - -c
                    - |
                      python -c 'import jax; print("Global device count:", jax.device_count())'
                      sleep 60
                    resources:
                      limits:
                       google.com/tpu: NUM_CHIPS
    

    다음을 바꿉니다.

    • NUM_SLICES: TPU 슬라이스 노드 풀 수입니다. 여기에서 NUM_SLICES3입니다.
    • ACCELERATOR_TYPE: TPU 슬라이스 노드 풀의 가속기 유형입니다. 이 값은 노드 풀 생성 중에 사용한 머신 유형에 매핑되어야 합니다. 예를 들어 ct5lp-hightpu-4t 머신 유형으로 노드 풀을 만든 경우 tpu-v5-lite-podslice를 사용합니다. 사용할 가속기 유형을 알아보려면 GKE에서 TPU 계획의 표를 참고하세요.
    • TPU_TOPOLOGY: TPU 슬라이스의 물리적 토폴로지입니다. 예를 들어 TPU 버전에 따라 4x4x4 또는 2x2입니다.
    • NUM_NODES: 노드 풀의 노드 수입니다. 0이거나 TPU_TOPOLOGY에 정의된 값의 곱({A}x{B}x{C})을 각 VM의 TPU 칩 수로 나눈 값이어야 합니다. 멀티호스트 TPU v4의 경우 각 VM의 TPU 칩 수는 4개입니다. 멀티호스트 TPU v5e의 경우 각 VM의 TPU 칩 수는 1개, 4개 또는 8개입니다. 따라서 TPU_TOPOLOGY2x4x4(VM마다 TPU 칩 4개가 있는 TPU v4)인 경우 NUM_NODES는 8과 같은 32/4입니다.
    • NUM_CHIPS: VM당 칩 수입니다.
  2. 매니페스트를 적용합니다.

    kubectl apply -f tpu-multislice.yaml
    
  3. 워크로드가 수락되었는지 확인합니다.

    kubectl get jobsets
    

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

    NAME            RESTARTS   COMPLETED   AGE
    multislice-job                         3s
    
  4. 프로비저닝된 포드의 상태를 모니터링합니다.

    kubectl get pods
    

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

     NAME                                READY   STATUS      RESTARTS   AGE
     multislice-job-slice-0-0-wzq9t      0/1     Completed   0          2m31s
     multislice-job-slice-0-1-zf4dp      0/1     Completed   0          2m30s
     multislice-job-slice-1-0-hbfn5      0/1     Completed   0          2m31s
     multislice-job-slice-1-1-45fgl      0/1     Completed   0          2m30s
     multislice-job-slice-2-0-wjbp4      0/1     Completed   0          2m30s
     multislice-job-slice-2-1-lwnvs      0/1     Completed   0          2m30s
    

    multislice-job JobSet는 포드를 예약하고, 만들고, 끝까지 실행합니다. 포드 이름은 <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix> 형식입니다. jobsetName 프리픽스에 따라 포드가 속한 JobSet가 결정됩니다.

  5. 선택사항: JAX 워크로드 삭제:

    kubectl delete -f tpu-multislice.yaml
    

XPK를 사용하여 GKE 리소스 만들기

  1. XPK 리드미 파일의 단계를 완료하여 XPK를 설치합니다.
  2. 사용 가능한 TPU가 있는 리전에서 표준 모드로 GKE 클러스터를 만듭니다. 다음 명령어는 지정된 TPU 리소스가 포함된 노드 풀이 있는 GKE 클러스터를 만듭니다.

    xpk cluster create \
        --cluster=$CLUSTER_NAME \
        --project=$PROJECT_ID \
        --zone=$ZONE \
        --tpu-type=ACCELERATOR_TYPE \
        --num-slices=NUM_SLICES \
        --reservation=$RESERVATION_NAME
    

    다음을 바꿉니다.

    • CLUSTER_NAME: 클러스터 이름입니다.
    • PROJECT_ID: Google Cloud 프로젝트 ID입니다.
    • ZONE: 사용할 TPU 버전에 따른 영역의 이름. 사용 가능한 위치를 확인하려면 GKE의 TPU 가용성을 참조하세요. 예약된 용량을 사용하려면 용량을 예약한 영역을 사용해야 합니다.
    • ACCELERATOR_TYPE: TPU 버전 및 유형입니다. 예를 들어 Ironwood (TPU7x)의 경우 tpu7x-standard-4t을 사용합니다.
    • NUM_SLICES: 클러스터에 사용할 TPU 슬라이스 수입니다. Multislice의 경우 1보다 커야 합니다.
    • RESERVATION_NAME: 예약의 이름

    출력의 마지막 두 줄은 다음과 비슷합니다.

    See your GKE Cluster here:  https://console.cloud.google.com/kubernetes/clusters/details/<location>/CLUSTER_NAME/details?project=PROJECT_ID
    
  3. 연결된 TPU 기기 수를 표시하는 워크로드를 실행합니다.

    xpk workload create \
        --cluster=$CLUSTER_NAME \
        --project=$PROJECT_ID \
        --zone=$ZONE \
        --tpu-type=ACCELERATOR_TYPE \
        --num-slices=NUM_SLICES \
        --docker-image=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest \
        --workload=WORKLOAD_NAME \
        --command="python -c 'import jax; print(\"TPU cores:\", jax.device_count())'"
    

    다음을 바꿉니다.

    • CLUSTER_NAME: 클러스터 이름입니다.
    • PROJECT_ID: Google Cloud 프로젝트 ID입니다.
    • ZONE: 선택한 머신의 노드 풀 영역입니다.
    • ACCELERATOR_TYPE: 선택한 TPU 유형입니다. bash 자동 완성을 사용하여 사용 가능한 옵션을 확인합니다.
    • NUM_SLICES: 클러스터에 구성된 TPU 슬라이스 수입니다. 이 값은 클러스터 생성 중에 사용된 값과 일치해야 합니다.
    • WORKLOAD_NAME: 워크로드의 이름입니다.

추가 설정 구성

다음 섹션에서는 멀티슬라이스에 적용할 수 있는 추가 구성에 대해 설명합니다.

hostNetwork로 네트워크 성능 개선

TPU 슬라이스 간의 네트워크 성능을 향상시키려면 hostNetworking을 사용 설정하는 것이 좋습니다. 포드 사양에서 hostNetwork: true를 사용해서 Kubernetes 네트워킹 스택을 모두 건너뛰고 Kubernetes 포드가 VM간 통신에 호스트 네트워크를 직접 사용할 수 있게 합니다.

hostNetworking을 사용 설정하려면 포드 사양에 다음 두 줄을 추가합니다.

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

hostNetwork를 사용하는 워커 노드 검색에 podHostnames를 계속 사용하려면 dnsPolicy: ClusterFirstWithHostNet를 설정합니다. 이는 자동 재개 학습 작업을 실행 중이고 동일한 체크포인트를 새로고침하기 위해 동일한 이름을 사용해야 하는 경우에 중요합니다.

TPU Trillium (v6e) 또는 Ironwood (TPU7x) (프리뷰)을 사용하고 포드에서 hostNetworking을 사용하는 경우 노드에서 /proc/sys/net/ipv4/tcp_rmem을 조정하기 위해 다음 DaemonSet을 설치합니다.

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/ai-on-gke/51bf3dcab6ff658cf62cc32867f96860bf58dfdc/scripts/network-setup/v6e-increase-rmem.yaml

TPU Trillium 또는 Ironwood (TPU7x)에서 hostNetwork 없이 네트워크 성능 개선

TPU Trillium 또는 Ironwood (TPU7x) (프리뷰)을 사용하고 포드에서 hostNetworking을 사용할 수 없는 경우 최적의 네트워크 성능을 위해 GKE DRANET을 구현하여 네트워크 기기를 할당하세요. GKE DRANET 기능을 사용하면 Kubernetes 및 GKE Dataplane V2를 우회하여 포드에 고성능 네트워킹 리소스를 요청할 수 있습니다.

ct6e-standard-4t 머신 유형은 두 개의 물리적 NIC로 지원됩니다. Kubernetes에는 포드에 전달할 수 없는 vNIC가 한 개 필요합니다. 따라서 포드가 두 vNIC에 직접 액세스하여 두 물리적 NIC의 최상의 성능을 달성할 수 있도록 각 노드에는 vNIC가 3개 있어야 합니다.

DRANET 기능을 사용 설정하고 ct6e-standard-4t에 네트워크 기기를 할당하려면 다음 단계를 완료하세요.

  1. 다음 명령어를 사용하여 VPC를 두 개 더 만듭니다.

    gcloud compute networks subnets create ADDITIONAL_SUBNET_1 \
        --project=PROJECT_ID \
        --network=ADDITIONAL_VPC_1 \
        --location=CONTROL_PLANE_LOCATION
    
    gcloud compute networks subnets create ADDITIONAL_SUBNET_2 \
        --project=PROJECT_ID \
        --network=ADDITIONAL_VPC_2 \
        --location=CONTROL_PLANE_LOCATION
    
  2. GKE Dataplane V2를 사용하는 GKE 클러스터를 만듭니다.

  3. cloud.google.com/gke-networking-dra-driver=true 라벨을 추가하여 생성된 네트워크와 GKE DRANET가 사용 설정된 노드 풀을 만듭니다.

    gcloud container node-pools create NODE_POOL_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --cluster=CLUSTER_NAME \
      --machine-type=MACHINE_TYPE \
      --num-nodes=NUM_NODES \
      --additional-node-network network=ADDITIONAL_VPC_1,subnetwork=ADDITIONAL_SUBNET_1 \
      --additional-node-network network=ADDITIONAL_VPC_2,subnetwork=ADDITIONAL_SUBNET_2 \
      --node-labels=cloud.google.com/gke-networking-dra-driver=true
    
  4. 노드 풀 상태를 확인한 후 ResourceClaimTemplate을 적용하여 추가 네트워크 기기로 워크로드를 만듭니다.

    apiVersion: resource.k8s.io/v1
    kind: ResourceClaimTemplate
    metadata:
      name: two-netdev
    spec:
      spec:
        devices:
          requests:
          - name: req-netdev
            exactly:
              deviceClassName: netdev.google.com
              allocationMode: ExactCount
              count: 2
    ---
    apiVersion: v1
    kind: Pod
    metadata:
      name: sample-netdevice-pod-1
    spec:
      containers:
      - name: sample-netdevice-pod
        image: busybox
        command: ["sleep", "infinity"]
        resources:
          claims:
          - name: netdev
          limits:
            google.com/tpu: 4
      restartPolicy: Always
      resourceClaims:
      - name: netdev
        resourceClaimTemplateName: two-netdev
    
  5. 인터페이스가 포드에서 사용 가능한지 확인합니다.

    kubectl exec -it pods/sample-netdevice-pod-1 -- ip a
    
권장사항:

eth0 인터페이스 대신 eth1eth2 인터페이스를 사용하여 네트워킹 성능을 개선합니다. 워크로드 사양에 export LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS --megascale_grpc_interface_prefixes=eth1,eth2,lo"를 추가하면 됩니다.

로깅 사용 설정

TPU 슬라이스 노드를 포함하여 GKE 노드에서 실행되는 컨테이너에서 내보낸 로그는 클러스터에서 GKE 시스템 로깅을 사용 설정한 경우에 로그 탐색기에 표시됩니다.

다음 필터와 함께 로그 탐색기를 사용해서 GKE에서 로그를 확인하여 워크로드의 컨테이너 로그를 볼 수 있습니다.

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME

TPU 슬라이스 및 작업자에 대해 다음 필터를 사용합니다.

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME
resource.labels.pod_name:<jobSetName>-<replicateJobName>-<job-index>-<worker-index>

자세한 내용은 GKE TPU 로그 보기를 참조하세요.

추가 측정항목 사용 설정

일반적인 TPU 측정항목 외에도 4개의 추가 멀티슬라이스 특정 TPU 런타임 측정항목이 있습니다. 이러한 측정항목은 GKE 버전 1.29.1-gke.1016000 이상에서 사용할 수 있습니다. TPU 워크로드는 JAX 버전 0.4.24를 사용해야 합니다.

다음은 사용 가능한 멀티슬라이스 측정항목입니다.

  • DCN(데이터 센터 네트워크) 전송 지연 시간: 멀티슬라이스 트래픽의 네트워크 전송 지연 시간 분포입니다.
  • 집합 지연 시간: 멀티슬라이스 트래픽의 엔드 투 엔드 집합 지연 시간의 분포입니다.
  • 호스트-기기 전송 지연 시간: 멀티슬라이스 트래픽의 각 데이터 청크에 대한 호스트-기기 전송 지연 시간 분포입니다.
  • 기기-호스트 전송 지연 시간: 멀티슬라이스 트래픽의 각 데이터 청크에 대한 기기-호스트 전송 지연 시간 분포입니다.

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

  • kubernetes.io/container/multislice/network/dcn_transfer_latencies
  • kubernetes.io/container/multislice/network/collective_end_to_end_latencies
  • kubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencies
  • kubernetes.io/container/multislice/accelerator/device_to_host_transfer_latencies

TPU 슬라이스와 멀티슬라이스 비교

다음 표에서는 TPU 슬라이스와 멀티슬라이스의 아키텍처 구성을 비교해서 보여줍니다.

TPU 슬라이스 멀티슬라이스
상호 연결 워크로드가 단일 TPU 슬라이스에서 실행됩니다. 슬라이스의 모든 TPU 칩이 ICI로 연결됩니다. 워크로드가 여러 TPU 슬라이스에서 실행됩니다. 슬라이스 내 통신이 ICI를 통해 수행됩니다. 슬라이스 간 통신이 DCN을 통해 수행됩니다.
지원되는 노드 풀 단일 호스트 TPU 슬라이스와 멀티 호스트 TPU 슬라이스 멀티 호스트 TPU 슬라이스의 그룹
권장 워크로드 유형 IndexedJob 또는 JobSet JobSet

리소스 삭제

비용이 청구되지 않도록 하는 가장 쉬운 방법은 튜토리얼에서 만든 Google Cloud 프로젝트를 삭제하는 것입니다. 또는 개별 리소스를 삭제할 수 있습니다.

프로젝트 삭제

    Google Cloud 프로젝트를 삭제합니다.

    gcloud projects delete PROJECT_ID

개별 리소스 삭제

GKE 클러스터를 삭제합니다.

```sh
gcloud container clusters delete  CLUSTER_NAME \
   --project=PROJECT_ID  \
   --location=CONTROL_PLANE_LOCATION
```

XPK 생성 리소스 삭제

XPK를 사용하여 클러스터를 만든 경우 추가 요금이 청구되지 않도록 클러스터를 삭제합니다.

xpk cluster delete \
    --cluster=$CLUSTER_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

다음을 바꿉니다.

  • CLUSTER_NAME: 클러스터의 이름입니다.
  • PROJECT_ID: Google Cloud 프로젝트 ID입니다.
  • ZONE: 클러스터가 생성된 영역입니다.

다음 단계