GKE의 멀티슬라이스 Cloud TPU 개요[공개 미리보기]

요약

이 문서에서는 GKE를 사용하여 멀티슬라이스 구성에서 Cloud TPU를 관리하는 방법을 설명합니다. GKE에서 단일 Cloud TPU 슬라이스를 사용하는 데 익숙하고 멀티슬라이스 Cloud TPU 사용에 대한 일반적인 경험이 있다고 가정합니다. 각 Cloud TPU 슬라이스를 노드 풀로 모델링하고 JobSet API를 사용하여 여러 노드 풀이 포함된 작업을 제출하는 것이 좋습니다.

피드백

질문이나 의견이 있으면 cloudtpu-multislice-preview@google.com으로 이메일을 보내주세요.

개념

자동 복구
슬라이스에서 유지보수 이벤트, 선점 또는 하드웨어 고장이 발생하면 Google Kubernetes Engine이 새 슬라이스를 만듭니다. 새 슬라이스를 만들 용량이 부족하면 하드웨어를 사용할 수 있을 때까지 생성이 완료되지 않습니다. 이후 학습을 계속할 수 있도록 멀티슬라이스 환경의 다른 모든 슬라이스가 다시 시작됩니다. 학습 스크립트에서 학습을 다시 시작하기 전에 체크포인트를 확인하여 로드하는 시작 스크립트를 정의해야 합니다.
데이터 센터 네트워킹(DCN)
멀티 슬라이스 구성에서 Cloud TPU 슬라이스를 연결하는 지연 시간이 길고 처리량이 낮은 네트워크입니다.
데이터 세트
모델이 학습 또는 추론을 위해 사용하는 데이터입니다.
호스트
호스트는 VM을 실행하는 물리적 컴퓨터입니다. 호스트는 한 번에 최대 VM 수를 실행할 수 있습니다. 각 VM에는 전용 Cloud TPU가 있습니다.
추론
호스트에 선행 학습된 머신러닝 모델을 로드하고 데이터에 대한 예측을 수행합니다.
칩 간 상호 연결(ICI)
Cloud TPU Pod 내에서 Cloud TPU를 연결하는 속도가 빠르고 지연 시간이 짧은 내부 링크입니다.
Kubernetes 포드
공유 스토리지 및 네트워크 리소스가 포함된 하나 이상의 컨테이너 그룹과 컨테이너 실행 방법에 대한 사양입니다.
멀티슬라이스
데이터 센터 네트워크(DCN)를 통해 통신할 수 있는 2개 이상의 Cloud TPU 슬라이스입니다.
노드 풀
GKE 컨텍스트에서 노드 풀은 Kubernetes 스케줄러가 워크로드(예: 학습 작업)와 함께 예약할 수 있는 노드(VM) 그룹입니다. GKE 클러스터 아키텍처와 개념에 대한 자세한 내용은 클러스터 아키텍처를 참조하세요.
큐에 추가된 리소스
단일 슬라이스 또는 멀티슬라이스 Cloud TPU 환경에 대한 요청을 큐에 추가하고 관리하는 데 사용되는 Cloud TPU 리소스의 표현입니다.
Tensor
머신러닝 모델에서 다차원 데이터를 나타내는 데 사용되는 데이터 구조입니다.
Tensor Processing Unit(TPU)
Google에서 내부적으로 개발한 ML 가속 칩입니다. 행렬 곱셈과 같은 주요 머신러닝 태스크에 가장 빠르고 효율적인 컴퓨팅을 제공하도록 설계되었습니다.
Cloud TPU Pod
전용 ICI 네트워크 인터페이스로 연결되는 Cloud TPU 칩 모음입니다. Cloud TPU Pod를 사용하면 처리 부하를 여러 Cloud TPU에 분산할 수 있습니다.
Cloud TPU 슬라이스
Cloud TPU 칩으로 구성된 Cloud TPU Pod의 논리적 하위 섹션입니다. 슬라이스의 모든 칩은 ICI 네트워크를 사용하여 서로 통신합니다.
Cloud TPU VM
기본 Cloud TPU에 액세스할 수 있는 Linux를 실행하는 가상 머신입니다. v4 Cloud TPU의 경우 각 Cloud TPU VM이 4개의 칩에 직접 액세스할 수 있습니다. Cloud TPU VM을 작업자라고도 합니다.

기본 요건

이 문서에서는 사용자가 Cloud TPU 멀티슬라이스 개요 및 GKE의 Cloud TPU 소개를 숙지했다고 가정합니다.

GKE에서 멀티슬라이스 Cloud TPU 관리

GKE 멀티슬라이스에 사용할 Cloud TPU 용량 프로비저닝

GKE에 Cloud TPU 워크로드 배포에서는 Cloud TPU 예약 방법에 대해 자세히 설명합니다. 그러나 Cloud TPU API를 통해 직접 Cloud TPU 용량을 사용하고 있고 GKE 멀티슬라이스에 동일한 용량을 사용하려는 경우에는 gke-tpu-support@google.com에 문의하여 용량의 일부를 GKE로 이전하세요.

멀티슬라이스 환경에서 노드 풀 만들기

멀티슬라이스 환경을 만들려면 여러 GKE 노드 풀을 한 멀티슬라이스 그룹으로 그룹화해야 합니다. 멀티슬라이스 환경에 포함되는 슬라이스마다 개별 노드 풀을 만들어야 합니다. 각 노드 풀을 한 멀티슬라이스 그룹에 추가해야 합니다. 멀티슬라이스 그룹은 여러 노드 풀을 함께 그룹화하여 관리하는 방법입니다. 멀티슬라이스 그룹은 멀티슬라이스 환경의 일부인 각 노드 풀에 라벨(이름/값 쌍)을 적용하여 정의됩니다. 이 튜토리얼에서는 MultisliceGroup 및 MultisliceGroupSize라는 라벨을 사용하지만 원하는 이름을 사용할 수 있습니다. 한 멀티슬라이스 그룹의 모든 노드 풀이 동일한 라벨(이름/값 쌍)과 동일한 Cloud TPU 토폴로지를 가져야 합니다.

생성할 때 멀티슬라이스 라벨을 추가할 수 있습니다. 이 예시에서는 라벨 이름으로 MultisliceGroup 및 MultisliceGroupSize를 사용합니다.

$ gcloud beta container node-pools create pool-name \
  --region=cluster-region \
  --cluster=cluster-name \
  --node-locations=node-zone \
  --machine-type=machine-type \
  --tpu-topology=tpu-topology \
  --node-labels=MultisliceGroup=multislice-group-name,MultisliceGroupSize=num-slices \
  --num-nodes=num-nodes \
  [--reservation-affinity=specific \
     --reservation=reservation-name]

참고:

  • 한 번에 하나의 노드 풀 작업만 수행할 수 있으므로 다음 노드 풀의 생성/삭제를 시작하려면 각 노드 풀이 생성/삭제될 때까지 기다려야 합니다.

  • 단일 호스트 슬라이스로 구성된 멀티슬라이스 환경은 지원되지 않습니다. 예를 들어 여러 v4-16, (1,4) 토폴로지를 사용하는 v5e 또는 v5e-8 슬라이스를 사용해 멀티슬라이스를 만들 수 없습니다.

기존 노드 풀에 MultisliceGroup 및 MultisliceGroupSize 라벨을 추가할 수도 있습니다.

$ gcloud beta container node-pools update pool-name \
  --region=cluster-region \
  --cluster=cluster-name \
  --node-labels=MultisliceGroup=multislice-group-name,MultisliceGroupSize=num-slices

노드 풀 상태 확인

kubectl로 노드 풀의 상태를 확인할 수 있습니다.

$ kubectl get nodes -l MultisliceGroup=multislice-group-1 | grep "Ready" | wc -l

참고:

  1. 명령어를 실행하기 전에 GKE 클러스터에 연결해야 합니다. 자세한 내용은 클러스터 액세스를 참조하세요.

  2. get nodes 명령어를 실행할 때 다음 오류가 표시될 수 있습니다. 

    E0608 13:25:01.586645 21293 memcache.go:255] couldn't get resource list for
http://metrics.k8s.io/v1beta1 : the server is currently unable to handle the
request

    이 오류는 일반적으로 GKE 클러스터에 Cloud TPU가 아닌 노드 풀이 없음을 의미합니다. 클러스터에 Cloud TPU가 아닌 노드 풀을 추가하고 명령어를 다시 시도합니다.

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

JobSet API

JobSet API는 작업 그룹을 함께 관리하는 데 사용되는 오픈소스 API입니다. 멀티 슬라이스 워크로드에 적합합니다. 단일 슬라이스 워크로드에 대한 GKE의 Cloud TPU 소개에 설명된 대로 단일 Kubernetes IndexedJob을 사용하여 단일 Cloud TPU 슬라이스에서 실행합니다. 멀티슬라이스 환경과 마찬가지로, 여러 Kubernetes IndexedJobs가 함께 작동해야 하는데 JobSet가 이를 달성하도록 도와줍니다.

이미지

JobSet는 다음을 제공합니다.

  • 실패 처리: 하위 작업이 실패할 경우 모든 하위 작업을 자동으로 다시 만듭니다.
  • 슬라이스 색인 생성: JobSet는 ID로 각 작업의 색인을 생성하여 작업 및 포드의 주석/라벨에 색인을 추가합니다.
  • 워크로드 배치: 단일 작업을 특정 도메인에 자동으로 할당합니다(이 경우 단일 노드 풀의 모든 작업).
  • 헤드리스 서비스 생성: JobSet의 모든 작업에 대한 headless-svc의 수명 주기를 관리합니다.

설치

  1. 클러스터에 JobSet API를 설치합니다.

    $ kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.2.1/manifests.yaml

  2. JobSet 컨트롤러가 실행 중인지 확인합니다.

    $ kubectl get pods -n jobset-system

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

다음 YAML 매니페스트는 v4-16의 4개 슬라이스에서 멀티슬라이스 워크로드를 실행하는 방법을 보여줍니다. 다음 YAML 파일을 복사하여 파일에 저장한 후 kubectl apply -f <file-name.yaml>을 실행합니다.

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-job  # JobSet name (${JOBSET_NAME})
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool # 1:1 job replica to node pool assignment
spec:
  failurePolicy:
    maxRestarts: 4  # The set will be restarted on failures up to 4 times.
  replicatedJobs:
    - name: slice    # Part of the name of the child Jobs (<replicateJobName>)
      replicas: 4    # Number of slices
      template:
        spec:
          parallelism: 2   # Must be set to number of nodes in each node pool
          completions: 2   # Must be set to number of nodes in each node pool
          backoffLimit: 0   # Must be set to 0. Fail the job when any pod fails.
          template:
            spec:
              affinity: # The affinity section is to make sure there is only one Multislice job running in a single Multislice Group. More on this below.
                podAffinity:
                  requiredDuringSchedulingIgnoredDuringExecution:
                  - labelSelector:
                      matchExpressions:
                      - key: jobset.sigs.k8s.io/jobset-name
                        operator: In
                        values:
                        - multislice-job # JobSet name
                    topologyKey: MultisliceGroup
                podAntiAffinity:
                  requiredDuringSchedulingIgnoredDuringExecution:
                  - labelSelector:
                      matchExpressions:
                      - key: jobset.sigs.k8s.io/jobset-name
                        operator: NotIn
                        values:
                        - multislice-job # JobSet name
                    topologyKey: MultisliceGroup
                    namespaceSelector:
                      matchExpressions:
                      - key: jobset.sigs.k8s.io/jobset-name
                        operator: Exists
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
                cloud.google.com/gke-tpu-topology: 2x2x2
                MultisliceGroupSize: "4"
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080 # Port for MXLA coordinator
                securityContext:
                  privileged: true
                command:
                - bash
                - -c
                - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  python -c 'import jax; print("Global device count:", jax.device_count())'
                resources:
                  limits:
                    google.com/tpu: 4 # Number of Cloud TPU VMs per worker

  • JobSet에서 Kubernetes 색인 생성된 작업 4개(슬라이스 수)를 만듭니다. Kubernetes 색인 생성된 작업은 이름 지정 패턴 <jobSetName>-<replicateJobName>-<job-index>를 따릅니다(이 경우 multislice-job-slice-{0,1,2,3}). 생성된 작업은 Google Cloud CLI에서 Workloads에 표시됩니다.

  • JobSet는 JobSet 이름과 동일한 이름(이 경우 multislice-job)의 헤드리스 서비스입니다.

  • 각 색인 생성된 작업에서 Kubernetes 포드 2개(작업자 수)를 만듭니다. 작업은 이름 지정 패턴 <jobSetName>-<replicateJobName>-<job-index>-<worker-index>-<5 letter suffix>(이 경우 multislice-job-slice-{0,1,2,3}-{1,2}-{5 letter suffix})를 따릅니다.

  • Google Kubernetes Engine의 멀티슬라이스는 현재 동기식 멀티 컨트롤러 학습만 지원합니다. 이를 위해 parallelism 및 completions를 각 노드 풀의 노드 수로 설정하고 backoff는 0으로 설정해야 합니다. 노드 수는 코어 수를 코어/노드 수로 나눈 값입니다. v4 TPU의 경우 8이므로 v4-16에서는 parallelism 및 completions를 2로 설정합니다.

  • 최소 JAX 버전은 v0.4.9입니다.

GKE에 환경 변수 삽입

GKE 웹훅은 색인 생성된 작업의 포드 사양에 환경 변수를 자동으로 삽입합니다. 이는 나중에 멀티슬라이스 워크로드를 실행하는 Kubernetes 포드에 상속됩니다. GKE는 다음 환경 변수를 삽입합니다.

환경 변수 값
TPU_WORKER_ID metadata.annotations['batch.kubernetes.io/job-completion-index']
TPU_WORKER_HOSTNAMES 작업의 포드 색인을 다음과 같이 쉼표로 구분된 목록으로 계산합니다. <job-name>-0.<subdomain>,........,<job-name>-<n-1>.<subdomain>
MEGASCALE_NUM_SLICES metadata.annotations['jobset.sigs.k8s.io/replicatedjob-replicas']
MEGASCALE_SLICE_ID metadata.annotations['jobset.sigs.k8s.io/job-index']
MEGASCALE_COORDINATOR_ADDRESS Generate the string using the jobSetName and the replicatedJobName:

<jobSetName>-<replicatedJobName>-0-0.<subdomain>

선택사항: GKE 포드에서 hostNetwork 사용 중지

포드 사양에서 hostNetwork: true를 사용하여 모든 Kubernetes 네트워킹 스택을 건너뛰고 Kubernetes 포드가 VM 간 통신에 호스트 네트워크를 직접 사용하도록 허용합니다. 이렇게 하면 슬라이스 간에 네트워크 성능이 향상됩니다. hostNetwork를 사용하는 작업자 검색에 계속 podHostnames를 사용하려면 dnsPolicy: ClusterFirstWithHostNet를 설정합니다. 이는 학습 작업을 자동 재개할 때 중요하며 동일한 체크포인트를 다시 로드하기 위해서는 이름이 동일해야 합니다.

hostNetworking을 사용 중지하려면 포드 사양에서 다음 두 줄을 삭제합니다.

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

멀티슬라이스 작업 조정

예약 교착 상태 방지

동일한 MultisliceGroupSize의 멀티슬라이스 작업(JobSet A 및 B)을 제출하는 경우 두 JobSet가 일부만 예약되는 상황이 발생할 수 있습니다(예: JobSet A에서 하위 작업 8개 중 5개가 예약되고 JobSet B에서 하위 작업 8개 중 3개가 예약됨). 이로 인해 노드 풀이 유휴 상태가 되고 JobSet A와 B가 모두 정지 상태가 됩니다.

podAffinity 및 podAntiAffinity를 활용하여 이를 방지할 수 있습니다. 멀티슬라이스 워크로드의 경우 topologyKey를 cloud.google.com/gke-nodepool 라벨 대신 멀티슬라이스 그룹을 나타내는 라벨(MultisliceGroup)로 변경합니다. 예를 들면 다음과 같습니다.

affinity:
  podAffinity:
    requiredDuringSchedulingIgnoredDuringExecution:
      - labelSelector:
          matchExpressions:
          - key: jobset.sigs.k8s.io/jobset-name
            operator: In
            values:
            - multislice-job # JobSet name
          topologyKey: MultisliceGroup
  podAntiAffinity:
    requiredDuringSchedulingIgnoredDuringExecution:
    - labelSelector:
        matchExpressions:
        - key: jobset.sigs.k8s.io/jobset-name
          operator: NotIn
          values:
          - multislice-job # JobSet name
      topologyKey: MultisliceGroup
      namespaceSelector:
        matchExpressions:
        - key: jobset.sigs.k8s.io/jobset-name
          operator: Exists

관측 가능성

Cloud Logging 로그

JobSet의 GKE 포드 이름은 다음 패턴을 사용하여 지정됩니다.

<jobSetName>-<replicateJobName>-<job-index>-<worker-index>-<5 letter suffix>

Cloud Logging 로그 탐색기와 다음 필터를 사용하여 로그를 보고 워크로드의 컨테이너 로그를 확인할 수 있습니다.

resource.type="k8s_container"
resource.labels.cluster_name=<cluster-name>
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=<jobSetName>

슬라이스 x 및 작업자 y의 로그를 필터링하려면 다음 필터를 사용합니다.

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