Melatih model machine learning berskala besar di GKE dengan Checkpointing Multi-Tier

Halaman ini menunjukkan cara menggunakan Checkpointing Multi-Tier untuk menyimpan dan mengelola checkpoint secara andal selama pelatihan model machine learning Anda di GKE. Penyimpanan dan pengelolaan titik pemeriksaan sangat penting untuk tugas pelatihan skala besar, yang ditentukan sebagai tugas yang menggunakan lebih dari ribuan node. Gangguan pada tugas berskala besar ini sering terjadi (berpotensi setiap jam) dan pemulihannya bisa lambat.

Manfaat

Menggunakan Pembuatan Checkpoint Multi-Tingkat memberikan manfaat berikut:

  • Pengelolaan data titik pemeriksaan yang diorkestrasi sepenuhnya, termasuk pencadangan, replikasi, dan pemulihan otomatis untuk beban kerja berikut:
  • Pemulihan cepat tugas pelatihan dari checkpoint yang disimpan di node lokal. Anda juga dapat memulihkan menggunakan titik pemeriksaan yang disimpan di node lain dalam cluster pelatihan.
  • Pemulihan cepat tugas pelatihan dari checkpoint yang disimpan di cadangan Cloud Storage dalam skenario terburuk, saat tidak ada checkpoint dalam cluster.

Sebelum memulai

Sebelum memulai, pastikan Anda telah melakukan 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.

Persyaratan

Checkpointing Multi-Tier memerlukan cluster GKE versi 1.32.4-gke.1415000 atau yang lebih baru.

Batasan

  • Cluster Autopilot tidak didukung.

Mengonfigurasi node GKE untuk menggunakan Checkpointing Multi-Tier

Bagian ini membahas cara mengonfigurasi node GKE di cluster baru dan yang sudah ada.

Mengonfigurasi node di cluster baru

  1. Buat cluster dengan Multi-Tier Checkpointing, driver CSI Cloud Storage FUSE, dan Workload Identity Federation untuk GKE diaktifkan. Jika menggunakan slice TPU untuk beban kerja machine learning, Anda harus menyesuaikan perintah pembuatan cluster untuk menyertakan konfigurasi untuk pool node slice TPU.

    gcloud container clusters create CLUSTER_NAME \
    --addons=HighScaleCheckpointing,GcsFuseCsiDriver  \
    --node-locations=NODE_LOCATION \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --cluster-version=CLUSTER_VERSION
    --location=CLUSTER_LOCATION \
    --machine-type=MACHINE_TYPE \
    --num-nodes=NUM_NODES

    Ganti nilai berikut:

    • CLUSTER_NAME: nama cluster.
    • NODE_LOCATION: zona untuk node cluster Anda. Di sinilah kapasitas TPU Anda berada.
    • PROJECT_ID: Google Cloud Project ID Anda.
    • CLUSTER_VERSION: versi cluster Anda. 1.32.4-gke.1415000 adalah versi minimum yang didukung.
    • CLUSTER_LOCATION: region tempat Anda ingin membuat cluster.
    • MACHINE_TYPE: jenis mesin yang digunakan untuk node yang menjalankan komponen seperti pengontrol JobSet dan pengontrol pembuatan checkpoint multi-tingkatan. Untuk pelatihan skala besar, sebaiknya gunakan minimal e2-standard-4 mesin. Anda tidak akan menggunakan jenis mesin ini untuk pelatihan model; sebagai gantinya, Anda akan membuat node pool terpisah untuk tujuan tersebut, yang sering kali memanfaatkan kelompok VM yang dioptimalkan akselerator.
    • NUM_NODES: jumlah node yang akan dibuat di setiap zona cluster.

Mengonfigurasi node di cluster yang ada

Untuk menggunakan Checkpointing Multi-Tier dengan cluster yang ada, aktifkan bersama dengan driver CSI Cloud Storage FUSE, dan Workload Identity Federation untuk GKE dengan perintah berikut. Versi cluster yang ada harus lebih baru dari 1.32.3-gke.1170000.

  1. Aktifkan Workload Identity Federation untuk GKE:

    gcloud container clusters update CLUSTER_NAME \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --location=CLUSTER_LOCATION

    Ganti nilai berikut:

    • CLUSTER_NAME: nama cluster.
    • PROJECT_ID: Google Cloud Project ID Anda.
    • CLUSTER_LOCATION: region cluster.

  2. Aktifkan Multi-Tier Checkpointing dan driver CSI Cloud Storage FUSE:

    gcloud container clusters update CLUSTER_NAME \
    --update-addons=HighScaleCheckpointing=ENABLED,GcsFuseCsiDriver=ENABLED \
    --location=CLUSTER_LOCATION

Mengonfigurasi izin untuk menggunakan Pembuatan Checkpoint Bertingkat

Bagian ini membahas cara mengonfigurasi izin untuk menggunakan Pembuatan Checkpoint Multi-Tingkat.

Memberikan akses ke bucket Cloud Storage

Volume sementara yang digunakan oleh driver CSI Multi-Tier Checkpointing harus menggunakan bucket Cloud Storage yang ada.

Untuk menyimpan checkpoint di bucket Cloud Storage, Multi-Tier Checkpointing memerlukan akses ke bucket. Berikan peran IAM Storage Object User (roles/storage.objectUser) di bucket ke akun layanan Kubernetes untuk Checkpointing Multi-Tier.

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member "principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/gke-managed-checkpointing/sa/gke-checkpointing-multitier-node" \
    --role "roles/storage.objectUser"

Ganti nilai berikut:

(Opsional) Berikan akses akun layanan default Compute Engine

Jika instance Compute Engine Anda memerlukan akses baca ke bucket Cloud Storage, berikan peran IAM Storage Object Viewer (roles/storage.objectViewer) ke akun layanan default Compute Engine.

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --role roles/storage.objectViewer

Deploy pengontrol JobSet

Pengontrol JobSet bertanggung jawab untuk mengelola tugas batch yang menjalankan pelatihan model Anda di GKE, dan alokasi sumber dayanya disesuaikan untuk menangani beban kerja secara efisien. Pastikan peluncur tugas pelatihan Anda men-deploy dan menggunakan JobSet.

Untuk meningkatkan permintaan memori menjadi 1 Gi, batas memori menjadi 2 Gi, dan permintaan CPU menjadi 1 untuk penampung pengelola dalam deployment JobSet Anda, jalankan perintah patch berikut:

kubectl patch -n jobset-system deploy jobset-controller-manager --type json \
    --patch '[{"op": "add", "path": "/spec/template/spec/containers/0/resources", "value": {"limits": {"memory": "2Gi"}, "requests": {"cpu": "1", "memory": "1Gi"}}}]'

Menginisialisasi driver CSI Multi-Tier Checkpointing

Bagian ini menjelaskan cara menginisialisasi driver CSI Multi-Tier Checkpointing pada node tempat workload Anda akan berjalan. Driver CSI bertanggung jawab untuk menangani penyimpanan dan pengelolaan checkpoint selama proses pelatihan model Anda.

Membuat CheckpointConfiguration

CheckpointConfiguration adalah resource kustom Kubernetes yang menentukan properti untuk men-deploy driver CSI Multi-Tier Checkpointing. Resource ini memiliki cakupan cluster.

  1. Buat manifes checkpoint.yaml berikut.

    apiVersion: checkpointing.gke.io/v1
kind: CheckpointConfiguration
metadata:
  name: MTC_CONFIG_NAME-configuration
spec:
    cloudStorageBucketName: GCS_BUCKET
    nodeSelector:
        node.kubernetes.io/instance-type: MACHINE_TYPE
    tolerations:
    - key: TOLERATION_KEY
        operator: Exists
        effect: NoSchedule
    inMemoryVolumeSize: IN_MEMORY_VOLUME_SIZE
    gcsFuseMountOptions:
    - implicit-dirs
    - metadata-cache:negative-ttl-secs:0
    - metadata-cache:ttl-secs:-1
    - metadata-cache:stat-cache-max-size-mb:-1
    - metadata-cache:type-cache-max-size-mb:-1
    - file-cache:max-size-mb:-1
    - file-cache:cache-file-for-range-read:true
    - file-system:kernel-list-cache-ttl-secs:0
    - file-cache:enable-parallel-downloads:true
    - read_ahead_kb=1024
    - write:enable-streaming-writes:true

    Ganti kode berikut:

    • MTC_CONFIG_NAME: nama CheckpointConfiguration Anda. Nama ini bersifat global untuk cluster dan tidak spesifik per tugas.
    • GCS_BUCKET: nama bucket Cloud Storage tempat Anda akan menyimpan data checkpoint. Gunakan bucket yang Anda siapkan di langkah Siapkan bucket Cloud Storage dengan izin.

    • MACHINE_TYPE: jenis mesin untuk akselerator yang sesuai. Nilainya bisa berupa salah satu dari berikut:

      Untuk mengetahui informasi selengkapnya tentang menjalankan workload terdistribusi di GPU dengan GKE, lihat Menjalankan GPU multi-instance. Untuk TPU, lihat Membuat node pool slice TPU.

    • TOLERATION_KEY: kolom ini memungkinkan driver CSI dijadwalkan di node dengan taint yang cocok. Untuk mengetahui informasi selengkapnya tentang cara kerja taint pada berbagai jenis akselerator, lihat halaman berikut:

    • IN_MEMORY_VOLUME_SIZE: ukuran untuk cache pembuatan titik pemeriksaan dalam memori. Tentukan kuantitas dan unit (misalnya, 200 Gi).Nilai ini harus:

      • Ukuran checkpoint lokal untuk TPU dikalikan dengan 2,2
      • Ukuran checkpoint lokal untuk GPU dengan satu peer dikalikan dengan 6,6.

  2. Terapkan manifes:

    kubectl apply -f checkpoint.yaml

  3. Periksa apakah driver CSI sedang berjalan:

    kubectl get pod -n gke-managed-checkpointing

    Output-nya akan mirip dengan yang berikut: Akan ada beberapa entri, satu per node yang dipercepat.

    NAME                                                          READY   STATUS    RESTARTS   AGE
multitier-driver-e2b033a7-a4e7-496a-87a3-ffd7fcc2e57b-2d4fz   5/5     Running   0          114s

Meng-uninstal driver CSI Multi-Tier Checkpointing

Jika Anda ingin membatalkan deployment driver CSI Multi-Tier Checkpointing, hapus resource CheckpointConfiguration. Pengontrol Multi-Tier Checkpointing menghapus driver CSI dari node. Tindakan ini akan menghapus disk RAM dan mengosongkan memori untuk workload lainnya. Contoh:

kubectl delete -f checkpoint.yaml

Mengelola retensi data dan pengumpulan sampah untuk cadangan Cloud Storage

Anda bertanggung jawab untuk menerapkan kebijakan retensi untuk cadangan Cloud Storage dari titik pemeriksaan. Pencadangan Multi-Tier hanya menulis cadangan titik pemeriksaan ke Cloud Storage dan tidak pernah mengubah atau menghapusnya.

Banyak alat open source yang dapat menangani retensi dan pengumpulan sampah, termasuk yang berikut:

Contoh berikut menggunakan backup-warden dengan direktori backup dipasang ke lokasi cadangan yang menggunakan Cloud Storage FUSE:

# Add --delete option to actually delete the backups, as is it only shows what would be deleted (dry-run)
backup-warden -p backup \
    --hourly 24 \
    --daily 7 \
    --weekly 5 \
    --monthly always \
    --yearly always \
    --prefer-recent

Perbarui manifes JobSet workload

Perbarui manifes JobSet untuk tugas Anda agar menyertakan volume checkpoint berskala besar. Detailnya bergantung pada workload Anda.

Misalnya, untuk memperluas JobSet contoh dari Men-deploy TPU Multislices di GKE, lakukan langkah-langkah berikut:

  1. Tambahkan baris berikut ke penampung jax-tpu.

    volumeMounts:
- name: checkpoint
  mountPath: CHECKPOINT_DIR

    Ganti CHECKPOINT_DIR dengan jalur ke direktori titik pemeriksaan Anda. Ini adalah lokasi tempat replicator.yaml dibuat dan Checkpointing Multi-Tingkat melakukan operasi penyimpanan checkpoint. Untuk mengetahui informasi selengkapnya, lihat Mengintegrasikan Pembuatan Checkpoint Multi-Tier di aplikasi Anda.

  2. Tambahkan baris berikut ke kolom spec.template.spec spesifikasi Job.

    volumes:
- name: checkpoint
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

Mengintegrasikan Pembuatan Checkpoint Multi-Tingkat di aplikasi Anda

Untuk membagikan informasi tentang lokasi titik pemeriksaan dan kesiapan replikasi, ubah aplikasi Anda agar menggunakan protokol berikut untuk berkomunikasi dengan Multi-Tier Checkpointing.

Memulai

Bagian ini menjelaskan langkah-langkah awal yang diperlukan aplikasi untuk berinteraksi dengan Checkpointing Multi-Tier.

Replicator adalah komponen inti Multi-Tier Checkpointing, yang berjalan di setiap node sebagai bagian dari driver CSI. Replicator mengelola replikasi titik pemeriksaan di seluruh tingkat penyimpanan, dari disk RAM lokal ke node peer dan ke penyimpanan eksternal seperti Cloud Storage.

File replicator.yaml berfungsi sebagai bidang kontrol dinamis antara tugas pelatihan ML (kode framework) dan komponen Replicator. Aplikasi ML Anda membuat file ini secara terprogram di volume lokal (RAMDisk), yang dapat diakses oleh tugas pelatihan dan layanan Replicator. Manifes ini memungkinkan framework ML memberikan konfigurasi runtime dan petunjuk pengelolaan siklus proses ke Replicator, yang berbeda dari parameter infrastruktur statis ( misalnya, frekuensi upload Cloud Storage) yang ditentukan selama penyiapan backend.

Untuk contoh konkret interaksi ini, lihat:

Aplikasi Anda harus melakukan langkah-langkah berikut selama startup:

  1. Tunggu hingga file replicator.yaml tidak ada, yang menunjukkan bahwa Replicator siap dikonfigurasi oleh aplikasi Anda. File replicator.yaml dibuat di lokasi CHECKPOINT_DIR yang Anda konfigurasi di bagian Perbarui manifes JobSet workload.

    Saat tugas pelatihan model pertama kali dibuat, file replicator.yaml tidak akan ada dan aplikasi Anda dapat langsung dilanjutkan. Namun, jika tugas dimulai ulang (misalnya, karena kegagalan atau intervensi manual), sistem mungkin masih memproses instance tugas sebelumnya, dan replicator.yaml dari instance tersebut mungkin masih ada di volume lokal.

  2. Aplikasi atau tugas ML Anda membuat file replicator.yaml dengan konfigurasi yang mirip dengan berikut.

    Orbax

    job-name: orbax
framework: orbax
assume-data-parallelism: 3
node-rank: 0
nodes: 32
peer-ranks: [1, 16] or peers-per-node: 2
backup-interval-minutes: 30

    PyTorch

    job-name: nemo
framework: pytorch.distributed
node-rank: 0
nodes: 32
peer-ranks: [1, 16] or peers-per-node: 2
backup-interval-minutes: 30

    Contoh konfigurasi ini memiliki kolom berikut:

    • name: nama tugas pelatihan.
    • framework: framework ML yang digunakan oleh tugas pelatihan.
    • node-rank: ID unik node saat ini dalam tugas pelatihan terdistribusi. Ini mewakili peringkat node dari node yang membuat file ini. Setiap node yang berpartisipasi dalam proses akan memiliki peringkatnya sendiri.
    • nodes: jumlah total node yang berpartisipasi dalam tugas pelatihan terdistribusi. Nilai ini berasal dari metadata Pod. Tugas pelatihan ML juga dapat melihat nilai ini.
    • peer-ranks atau peers-per-node: dua cara alternatif untuk menentukan topologi replikasi. Hanya salah satu dari dua parameter ini yang boleh ada.
      • peer-ranks: peringkat eksplisit node peer yang data titik pemeriksaan node saat ini harus direplikasi. Hal ini memberikan kontrol mendetail atas node tertentu yang berfungsi sebagai partner replikasi.
      • peers-per-node: jumlah node peer per node yang harus dipilih secara otomatis oleh Replicator untuk replikasi.
    • backup-interval-minutes: frekuensi, dalam menit, saat titik pemeriksaan dicadangkan ke Cloud Storage. Sebaiknya tetapkan nilai ini ke 30 menit atau lebih.

  3. Tunggu hingga file replicator.yaml baru dihapus oleh sistem. Hal ini menandakan bahwa replikator telah dimulai ulang dan melakukan pembersihan. Langkah ini memungkinkan Anda menghindari file sementara atau usang di volume lokal saat aplikasi Anda melakukan langkah-langkah di bagian berikutnya.

Memulihkan dari checkpoint terakhir yang diketahui baik (LKG)

  1. Setelah replikator diinisialisasi, Multi-Tier Checkpointing akan membuat satu link simbolik per pekerja TPU atau GPU. Link simbolis ini dibuat di volume lokal yang sama dengan file replicator.yaml yang ter-mount, tempat tugas menyimpan titik pemeriksaan.

    Link simbolis memiliki bentuk <job-name>-s{step}-n<node-rank>-w<worker-index>.restore.

  2. Pulihkan setiap pekerja dari file .restore yang sesuai. Untuk contohnya, lihat contoh pengelola titik pemeriksaan yang direplikasi Orbax di bagian berikutnya.

Simpan checkpoint

Aplikasi Anda melakukan langkah-langkah ini beberapa kali saat tugas pelatihan berlangsung. Operasi penyimpanan terjadi di lokasi CHECKPOINT_DIR yang Anda konfigurasi di Perbarui manifes JobSet workload.

Orbax

Buat checkpoint Orbax. Direktori diberi nama dengan nomor langkah. Replicator mendeteksi direktori checkpoint yang baru dibuat, melakukan replikasi atau pencadangan sesuai kebutuhan, dan membersihkan secara otomatis.

Untuk mengetahui informasi selengkapnya tentang cara menggunakan pengelola checkpoint replikator Orbax, lihat file MaxtTest checkpointing. Untuk contoh interaksi layanan replikator, lihat file MaxText max_utils.

PyTorch

Gunakan InClusterLocalCheckpointIO sebagai pytorch_lightning.CheckpointIO kustom untuk mengaktifkan pembuatan titik pemeriksaan terdistribusi yang benar dengan penyimpanan lokal. Contoh perintah berikut mengaktifkan pembuatan titik pemeriksaan multi-tingkatan menggunakan implementasi referensi yang dibangun di framework NVIDIA NeMo:

torchrun train.py <other_train_flags> \
    --local-ckpt-dir=CHECKPOINT_DIR \
    --local-ckpt-interval=20 \
    --job-name=JOB_NAME \
    --enable-high-scale-ckpt

Ganti kode berikut:

  • CHECKPOINT_DIR: jalur ke direktori titik pemeriksaan Anda.
  • JOB_NAME: nama workload tugas pelatihan Anda.

Upgrade cluster

Untuk upgrade cluster, Anda dapat menghapus dan membuat ulang objek CheckpointConfiguration sebelum atau setelah upgrade. Tindakan ini diperlukan karena daemonset checkpointing node, yang di-deploy secara dinamis oleh objek ini, tidak akan diupgrade secara otomatis.

Jika Anda secara khusus ingin mempertahankan spesifikasi daemonset yang sama, Anda tidak perlu melakukan apa pun.

Memecahkan masalah

Bagian ini berisi panduan pemecahan masalah untuk menyelesaikan masalah terkait Checkpointing Multi-Tier. Untuk pemecahan masalah penyimpanan umum, lihat Memecahkan masalah Cloud Storage di GKE.

Checkpointing Multi-Tier tidak diaktifkan

Error berikut menunjukkan bahwa Multi-Tier Checkpointing tidak diaktifkan di cluster Anda:

error: unable to recognize "checkpoint.yaml": no matches for kind "CheckpointConfiguration" in version "checkpointing.gke.io/v1"

Anda mungkin mengalami error ini setelah menjalankan kubectl apply -f checkpoint.yaml di langkah Buat CheckpointConfiguration.

Untuk mengatasi masalah ini, periksa apakah Anda telah mengaktifkan Pembuatan Checkpoint Multi-Tingkat di cluster Anda dengan perintah berikut:

gcloud container clusters describe CLUSTER_NAME \
    --project PROJECT_ID
    --location CLUSTER_LOCATION

Jika Multi-Tier Checkpointing diaktifkan, output-nya akan mirip dengan berikut ini:

addonsConfig:
  gcePersistentDiskCsiDriverConfig:
    enabled: true
  gcsFuseCsiDriverConfig:
    enabled: true
  highScaleCheckpointingConfig:
    enabled: true
  kubernetesDashboard:
    disabled: true
  networkPolicyConfig:
    disabled: true

Jika Multi-Tier Checkpointing tidak diaktifkan, perbarui cluster Anda untuk mengaktifkan Multi-Tier Checkpointing.

Driver CSI Multi-Tier Checkpointing tidak dapat memasang volume

Anda mungkin mengalami masalah ini jika driver CSI tidak dapat memasang volume Cloud Storage. Mungkin ada beberapa baris yang serupa dengan ini.

kubectl get pod -n gke-managed-checkpointing
NAME                                                          READY   STATUS     RESTARTS   AGE
multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9   0/5     Init:0/1   0          6m32s

Untuk mengatasi masalah ini, periksa peristiwa Pod driver CSI, seperti yang ditunjukkan dalam contoh berikut:

kubectl describe pod multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 -n gke-managed-checkpointing

Events:
  Type     Reason       Age                 From               Message
  ----     ------       ----                ----               -------
  Normal   Scheduled    17m                 default-scheduler  Successfully assigned gke-managed-checkpointing/multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 to gke-my-cluster-default-pool-353c773f-6d8q
  Warning  FailedMount  82s (x16 over 17m)  kubelet            MountVolume.SetUp failed for volume "gcs" : rpc error: code = PermissionDenied desc = failed to get GCS bucket "checkpointing-test-bucket": googleapi: Error 403: Caller does not have storage.objects.list access to the Google Cloud Storage bucket. Permission 'storage.objects.list' denied on resource (or it may not exist)., forbidden

Jika masalah terjadi karena error PermissionDenied bucket Cloud Storage, seperti yang ditunjukkan dalam contoh, Anda dapat menyelesaikan masalah dengan menyiapkan izin dengan benar.

Langkah berikutnya