Men-deploy workload TPU di Autopilot GKE


Halaman ini menjelaskan cara mempercepat beban kerja machine learning (ML) dengan men-deploy menggunakan akselerator Cloud TPU (TPU) di cluster Autopilot Google Kubernetes Engine (GKE). Sebelum membaca halaman ini, pastikan Anda memahami konsep berikut:

  1. Pengantar Cloud TPU
  2. Arsitektur sistem Cloud TPU
  3. Tentang TPU di GKE

Cara kerja TPU di Autopilot

Untuk menggunakan TPU dalam workload Autopilot, Anda meminta versi TPU dan topologi yang didukung untuk versi TPU tersebut dalam manifes workload. Kemudian, Anda menggunakan kolom resources.requests dan resources.limits Kubernetes untuk menentukan jumlah chip TPU yang akan digunakan oleh workload. Saat Anda men-deploy workload, GKE menyediakan node yang memiliki konfigurasi TPU yang diminta dan menjadwalkan Pod Anda di node tersebut. GKE menempatkan setiap workload di node-nya sendiri sehingga setiap Pod dapat mengakses resource lengkap node dengan risiko gangguan yang diminimalkan.

TPU di Autopilot kompatibel dengan kemampuan berikut:

  1. Spot Pod
  2. Reservasi kapasitas tertentu
  3. Pod runtime yang diperpanjang

Merencanakan konfigurasi TPU

Rencanakan konfigurasi TPU berdasarkan model dan jumlah memori yang diperlukan. Sebelum menggunakan panduan ini untuk men-deploy workload di TPU, selesaikan langkah-langkah perencanaan di Merencanakan konfigurasi TPU.

Harga

Untuk mengetahui informasi harga, lihat Harga Autopilot.

Sebelum memulai

Sebelum memulai, pastikan Anda telah menjalankan tugas berikut:

  • Aktifkan Google Kubernetes Engine API.
  • Aktifkan Google Kubernetes Engine API
  • Jika ingin menggunakan Google Cloud CLI untuk tugas ini, instal lalu lakukan inisialisasi gcloud CLI. Jika sebelumnya Anda telah menginstal gcloud CLI, dapatkan versi terbaru dengan menjalankan gcloud components update.
  • Pastikan Anda memiliki cluster Autopilot yang menjalankan GKE versi 1.29.2-gke.1521000 atau yang lebih baru.
  • Untuk menggunakan TPU yang direservasi, pastikan Anda memiliki reservasi kapasitas tertentu yang sudah ada. Untuk mendapatkan petunjuk, lihat Menggunakan resource zona yang dicadangkan.

Pastikan Anda memiliki kuota TPU

Bagian berikut membantu Anda memastikan bahwa Anda memiliki kuota yang cukup saat menggunakan TPU di GKE.

Untuk membuat node slice TPU, Anda harus memiliki kuota TPU yang tersedia, kecuali jika Anda menggunakan pemesanan kapasitas yang ada. Jika Anda menggunakan TPU yang dicadangkan, lewati bagian ini.

Membuat node slice TPU di GKE memerlukan kuota Compute Engine API (compute.googleapis.com), bukan kuota Cloud TPU API (tpu.googleapis.com). Nama kuota berbeda di Pod Autopilot reguler dan di Pod Spot.

Untuk memeriksa batas dan penggunaan saat ini dari kuota Compute Engine API untuk TPU, ikuti langkah-langkah berikut:

  1. Buka halaman Kuota di Konsol Google Cloud:

    Buka Kuota

  2. Di kotak Filter , lakukan tindakan berikut:

    1. Pilih properti Service, masukkan Compute Engine API, lalu tekan Enter.

    2. Pilih properti Type, lalu pilih Quota.

    3. Pilih properti Name dan masukkan nama kuota berdasarkan versi dan nilai TPU di pemilih node cloud.google.com/gke-tpu-accelerator. Misalnya, jika Anda berencana membuat node TPU v5e on-demand yang nilainya di pemilih node cloud.google.com/gke-tpu-accelerator adalah tpu-v5-lite-podslice, masukkan TPU v5 Lite PodSlice chips.

      Versi TPU cloud.google.com/gke-tpu-accelerator Nama kuota untuk instance on demand Nama kuota untuk instance Spot2
      TPU v3 tpu-v3-device TPU v3 Device chips Preemptible TPU v3 Device chips
      TPU v3 tpu-v3-slice TPU v3 PodSlice chips Preemptible TPU v3 PodSlice chips
      TPU v4 tpu-v4-podslice TPU v4 PodSlice chips Preemptible TPU v4 PodSlice chips
      TPU v5e tpu-v5-lite-device TPU v5 Lite Device chips Preemptible TPU v5 Lite Device chips
      TPU v5e tpu-v5-lite-podslice TPU v5 Lite PodSlice chips Preemptible TPU v5 Lite PodSlice chips
      TPU v5p tpu-v5p-slice TPU v5p chips Preemptible TPU v5p chips
      TPU v6e (Pratinjau) tpu-v6e-slice TPU v6e Slice chips Preemptible TPU v6e Lite PodSlice chips
    4. Pilih properti Dimensions (e.g. locations) dan masukkan region: diikuti dengan nama region tempat Anda berencana membuat TPU di GKE. Misalnya, masukkan region:us-west4 jika Anda berencana membuat node slice TPU di zona us-west4-a. Kuota TPU bersifat regional, sehingga semua zona dalam region yang sama menggunakan kuota TPU yang sama.

Jika tidak ada kuota yang cocok dengan filter yang Anda masukkan, berarti project belum diberi kuota yang ditentukan untuk region yang Anda butuhkan, dan Anda harus meminta peningkatan kuota TPU.

Saat reservasi TPU dibuat, nilai batas dan penggunaan saat ini untuk kuota yang sesuai akan meningkat sesuai dengan jumlah chip dalam reservasi TPU. Misalnya, saat reservasi dibuat untuk 16 chip TPU v5e yang nilainya di pemilih node cloud.google.com/gke-tpu-accelerator adalah tpu-v5-lite-podslice, Batas dan Penggunaan saat ini untuk kuota TPU v5 Lite PodSlice chips di region yang relevan akan bertambah 16.

Kuota untuk resource GKE tambahan

Anda mungkin perlu meningkatkan kuota terkait GKE berikut di region tempat GKE membuat resource Anda.

  • Kuota SSD Persistent Disk (GB): Disk booting setiap node Kubernetes memerlukan 100 GB secara default. Oleh karena itu, kuota ini harus ditetapkan setidaknya sebesar jumlah maksimum node GKE yang Anda perkirakan akan dibuat dan 100 GB (node * 100 GB).
  • Kuota alamat IP yang digunakan: Setiap node Kubernetes menggunakan satu alamat IP. Oleh karena itu, kuota ini harus ditetapkan setidaknya setinggi jumlah maksimum node GKE yang Anda perkirakan akan dibuat.
  • Pastikan max-pods-per-node selaras dengan rentang subnet: Setiap node Kubernetes menggunakan rentang IP sekunder untuk Pod. Misalnya, max-pods-per-node dari 32 memerlukan 64 alamat IP yang diterjemahkan menjadi subnet /26 per node. Perhatikan bahwa rentang ini tidak boleh dibagikan dengan cluster lain. Untuk menghindari penggunaan rentang alamat IP secara maksimal, gunakan flag --max-pods-per-node untuk membatasi jumlah pod yang diizinkan untuk dijadwalkan di node. Kuota untuk max-pods-per-node harus ditetapkan setidaknya sama dengan jumlah maksimum node GKE yang Anda perkirakan akan dibuat.

Untuk meminta penambahan kuota, lihat Meminta kuota yang lebih tinggi.

Menyiapkan aplikasi TPU

Workload TPU memiliki persyaratan persiapan berikut.

  1. Framework seperti JAX, PyTorch, dan TensorFlow mengakses VM TPU menggunakan library bersama libtpu. libtpu mencakup compiler XLA, software runtime TPU, dan driver TPU. Setiap rilis PyTorch dan JAX memerlukan versi libtpu.so tertentu. Untuk menggunakan TPU di GKE, pastikan Anda menggunakan versi berikut:
    Jenis TPU Versi libtpu.so
    TPU v6e
    tpu-v6e-slice
    TPU v5e
    tpu-v5-lite-podslice
    tpu-v5-lite-device
    TPU v5p
    tpu-v5p-slice
    • Versi jax[tpu] yang direkomendasikan: 0.4.19 atau yang lebih baru.
    • Versi torchxla[tpuvm] yang direkomendasikan: disarankan untuk menggunakan build versi setiap malam pada 23 Oktober 2023.
    TPU v4
    tpu-v4-podslice
    TPU v3
    tpu-v3-slice
    tpu-v3-device
  2. Tetapkan variabel lingkungan berikut untuk penampung yang meminta resource TPU:
    • TPU_WORKER_ID: Bilangan bulat unik untuk setiap Pod. ID ini menunjukkan worker-id unik di slice TPU. Nilai yang didukung untuk kolom ini berkisar dari nol hingga jumlah Pod dikurangi satu.
    • TPU_WORKER_HOSTNAMES: Daftar nama host VM TPU atau alamat IP yang dipisahkan koma yang perlu berkomunikasi satu sama lain dalam slice. Harus ada nama host atau alamat IP untuk setiap VM TPU dalam slice. Daftar alamat IP atau nama host diurutkan dan dimulai dengan indeks nol oleh TPU_WORKER_ID.
    • GKE secara otomatis memasukkan variabel lingkungan ini menggunakan webhook yang berubah saat Tugas dibuat dengan properti completionMode: Indexed, subdomain, parallelism > 1, dan meminta google.com/tpu. GKE menambahkan Service headless sehingga data DNS ditambahkan untuk Pod yang mendukung Service.

Setelah menyelesaikan persiapan workload, Anda dapat menjalankan Tugas yang menggunakan TPU.

Meminta TPU dalam workload

Bagian ini menunjukkan cara membuat Tugas yang meminta TPU di Autopilot. Dalam workload apa pun yang memerlukan TPU, Anda harus menentukan hal berikut:

  • Pemilih node untuk versi dan topologi TPU
  • Jumlah TPU chip untuk penampung dalam workload Anda

Untuk mengetahui daftar versi TPU, topologi, dan jumlah chip serta node TPU yang didukung dalam slice, lihat Memilih konfigurasi TPU Autopilot.

Pertimbangan untuk permintaan TPU dalam workload

Hanya satu penampung di Pod yang dapat menggunakan TPU. Jumlah TPU chip yang diminta penampung harus sama dengan jumlah TPU chip yang terpasang ke node dalam slice. Misalnya, jika Anda meminta TPU v5e (tpu-v5-lite-podslice) dengan topologi 2x4, Anda dapat meminta salah satu dari hal berikut:

  • Chip 4, yang membuat dua node multi-host dengan masing-masing 4 chip TPU
  • Chip 8, yang membuat satu node host tunggal dengan 8 chip TPU

Sebagai praktik terbaik untuk memaksimalkan efisiensi biaya, selalu gunakan semua TPU dalam slice yang Anda minta. Jika Anda meminta slice multi-host yang terdiri dari dua node dengan masing-masing 4 TPU chip, Anda harus men-deploy workload yang berjalan di kedua node dan menggunakan semua 8 TPU chip dalam slice.

Membuat workload yang meminta TPU

Langkah-langkah berikut membuat Tugas yang meminta TPU. Jika memiliki workload yang berjalan di slice TPU multi-host, Anda juga harus membuat Layanan headless yang memilih workload menurut namanya. Layanan headless ini memungkinkan Pod di node yang berbeda dalam slice multi-host untuk berkomunikasi satu sama lain dengan memperbarui konfigurasi DNS Kubernetes agar mengarah ke Pod dalam beban kerja.

  1. Simpan manifes berikut sebagai 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
    

    Ganti kode berikut:

    • TPU_TYPE: jenis TPU yang akan digunakan, seperti tpu-v4-podslice. Harus berupa nilai yang didukung oleh GKE.
    • TOPOLOGY: pengaturan chip TPU dalam slice, seperti 2x2x4. Harus berupa topologi yang didukung untuk jenis TPU yang dipilih.
    • NUMBER_OF_CHIPS: jumlah chip TPU yang akan digunakan oleh penampung. Harus sama dengan nilai untuk limits dan requests.
  2. Deploy Job:

    kubectl create -f tpu-autopilot.yaml
    

Saat Anda membuat Tugas ini, GKE akan otomatis melakukan hal berikut:

  1. Menyediakan node untuk menjalankan Pod. Bergantung pada jenis TPU, topologi, dan permintaan resource yang Anda tentukan, node ini adalah slice satu host atau slice multi-host.
  2. Menambahkan taint ke Pod dan toleransi ke node untuk mencegah beban kerja lainnya berjalan di node yang sama dengan beban kerja TPU.

Contoh: Menampilkan total TPU chip dalam slice multi-host

Workload berikut menampilkan jumlah TPU chip di semua node dalam slice TPU multi-host. Untuk membuat slice multi-host, workload memiliki parameter berikut:

  • Versi TPU: TPU v4
  • Topologi: 2x2x4

Pilihan versi dan topologi ini menghasilkan slice multi-host.

  1. Simpan manifes berikut sebagai 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. Deploy manifes:
    kubectl create -f available-chips-multihost.yaml
    

    GKE menjalankan slice TPU v4 dengan empat VM (slice TPU multi-host). Slice memiliki 16 chip TPU yang saling terhubung.

  3. Pastikan bahwa Tugas membuat empat Pod:
    kubectl get pods
    

    Outputnya mirip dengan yang berikut ini:

    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. Dapatkan log salah satu Pod:
    kubectl logs POD_NAME
    

    Ganti POD_NAME dengan nama salah satu Pod yang dibuat. Contoh, tpu-job-podslice-0-5cd8r.

    Outputnya mirip dengan hal berikut ini:

    TPU cores: 16
    

Contoh: Menampilkan chip TPU dalam satu node

Workload berikut adalah Pod statis yang menampilkan jumlah chip TPU yang terpasang ke node tertentu. Untuk membuat node satu host, workload memiliki parameter berikut:

  • Versi TPU: TPU v5e
  • Topologi: 2x4

Pilihan versi dan topologi ini menghasilkan slice satu host.

  1. Simpan manifes berikut sebagai 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. Deploy manifes:
    kubectl create -f available-chips-singlehost.yaml
    

    GKE menyediakan node dengan delapan slice TPU host tunggal yang menggunakan TPU v5e. Setiap node TPU memiliki delapan chip TPU (slice TPU host tunggal).

  3. Dapatkan log Pod:
    kubectl logs tpu-job-jax-v5
    

    Outputnya mirip dengan hal berikut ini:

    Total TPU chips: 8
    

Kemampuan observasi dan metrik

Dasbor

Di halaman Kubernetes Clusters di konsol Google Cloud, tab Observability menampilkan metrik kemampuan observasi TPU. Untuk informasi selengkapnya, lihat Metrik kemampuan observasi GKE.

Dasbor TPU hanya diisi jika Anda telah mengaktifkan metrik sistem di cluster GKE.

Metrik runtime

Pada GKE versi 1.27.4-gke.900 atau yang lebih baru, workload TPU yang menggunakan JAX versi 0.4.14 atau yang lebih baru dan menentukan containerPort: 8431 akan mengekspor metrik pemakaian TPU sebagai metrik sistem GKE. Metrik berikut tersedia di Cloud Monitoring untuk memantau performa runtime workload TPU Anda:

  • Siklus tugas: Persentase waktu dalam periode pengambilan sampel terakhir (60 detik) saat TensorCore secara aktif memproses TPU chip. Persentase yang lebih besar berarti pemakaian TPU yang lebih baik.
  • Memori yang digunakan: Jumlah memori akselerator yang dialokasikan dalam byte. Sampelnya dibuat setiap 60 detik.
  • Total memori: Total memori akselerator dalam byte. Dibuat sampelnya setiap 60 detik.

Metrik ini terletak di skema node Kubernetes (k8s_node) dan penampung Kubernetes (k8s_container).

Container Kubernetes:

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

Node Kubernetes:

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

Metrik host

Pada GKE versi 1.28.1-gke.1066000 atau yang lebih baru, VM dalam slice TPU mengekspor metrik penggunaan TPU sebagai metrik sistem GKE. Metrik berikut tersedia di Cloud Monitoring untuk memantau performa host TPU Anda:

  • Pemakaian TensorCore: Persentase TensorCore yang saat ini dipakai. Nilai TensorCore sama dengan jumlah unit perkalian matriks (MXU) ditambah unit vektor. Nilai pemakaian TensorCore adalah pembagian dari operasi TensorCore yang dilakukan selama periode sampel terakhir (60 detik) dengan jumlah operasi TensorCore yang didukung selama periode yang sama. Nilai yang lebih besar berarti pemakaian yang lebih baik.
  • Pemakaian Bandwidth Memori: Persentase dari bandwidth memori akselerator yang saat ini sedang digunakan. Dihitung dengan membagi bandwidth memori yang digunakan selama periode sampel (60 detik) dengan bandwidth maksimum yang didukung selama periode sampel yang sama.

Metrik ini terletak di skema node Kubernetes (k8s_node) dan penampung Kubernetes (k8s_container).

Container Kubernetes:

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

Node Kubernetes:

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

Untuk mengetahui informasi selengkapnya, lihat Metrik Kubernetes dan Metrik sistem GKE.

Logging

Log yang dikeluarkan oleh container yang berjalan di node GKE, termasuk VM TPU, dikumpulkan oleh agen logging GKE, dikirim ke Logging, dan terlihat di Logging.

Rekomendasi untuk workload TPU di Autopilot

Rekomendasi berikut dapat meningkatkan efisiensi workload TPU Anda:

  • Gunakan Pod waktu proses yang diperpanjang untuk masa tenggang hingga tujuh hari sebelum GKE menghentikan Pod Anda untuk penurunan skala atau upgrade node. Anda dapat menggunakan periode pengecualian dan pemeliharaan dengan Pod waktu proses yang diperpanjang untuk menunda upgrade node otomatis lebih lanjut.
  • Gunakan reservasi kapasitas untuk memastikan bahwa workload Anda menerima TPU yang diminta tanpa ditempatkan dalam antrean untuk ketersediaan.