Halaman ini diterjemahkan oleh Cloud Translation API.

Mempertahankan progres pelatihan menggunakan Autocheckpoint

Secara historis, saat VM TPU memerlukan pemeliharaan, prosedur akan segera dimulai, tanpa memberi waktu bagi pengguna untuk melakukan tindakan yang mempertahankan progres seperti menyimpan titik pemeriksaan. Hal ini ditunjukkan dalam Gambar 1(a).

Diagram yang menunjukkan dampak pemeliharaan host dengan dan tanpa pemeriksaan otomatis

Gambar 1. Ilustrasi fitur Autocheckpoint: (a) Tanpa Autocheckpoint, progres pelatihan dari checkpoint terakhir akan hilang saat ada peristiwa pemeliharaan mendatang. (b) Dengan Autocheckpoint, progres pelatihan sejak checkpoint terakhir dapat dipertahankan saat ada peristiwa pemeliharaan mendatang.

Anda dapat menggunakan Autocheckpoint (Gambar 1(b)) untuk mempertahankan progres pelatihan dengan mengonfigurasi kode untuk menyimpan titik pemeriksaan yang tidak terjadwal saat peristiwa pemeliharaan terjadi. Saat peristiwa pemeliharaan terjadi, progres sejak titik kontrol terakhir akan otomatis disimpan. Fitur ini berfungsi pada satu irisan dan Multislice.

Fitur Autocheckpoint berfungsi dengan framework yang dapat menangkap sinyal SIGTERM, lalu menyimpan checkpoint. Framework yang didukung meliputi:

MaxText,
Pax,
JAX dengan Orbax.

Menggunakan Titik pemeriksaan otomatis

Fitur Autocheckpoint dinonaktifkan secara default. Saat membuat TPU atau meminta resource dalam antrean, Anda dapat mengaktifkan Autocheckpoint dengan menambahkan flag --autocheckpoint-enabled saat menyediakan TPU. Dengan mengaktifkan fitur ini, Cloud TPU akan melakukan langkah-langkah berikut setelah menerima notifikasi tentang peristiwa pemeliharaan:

Menangkap sinyal SIGTERM yang dikirim ke proses menggunakan perangkat TPU
Tunggu hingga proses keluar, atau 5 menit telah berlalu, mana saja yang lebih dulu
Melakukan pemeliharaan pada slice yang terpengaruh

Infrastruktur yang digunakan oleh Autocheckpoint tidak bergantung pada framework ML. Framework ML apa pun dapat mendukung Autocheckpoint jika dapat menangkap sinyal SIGTERM dan memulai proses pembuatan checkpoint.

Dalam kode aplikasi, Anda perlu mengaktifkan kemampuan Autocheckpoint yang disediakan oleh framework ML. Misalnya, di Pax, hal ini berarti mengaktifkan flag command line saat meluncurkan pelatihan. Untuk informasi selengkapnya, lihat panduan memulai Autocheckpoint dengan Pax. Di balik layar, framework menyimpan titik pemeriksaan yang tidak terjadwal saat sinyal SIGTERM diterima, dan VM TPU yang terpengaruh akan menjalani pemeliharaan saat TPU tidak lagi digunakan.

Panduan memulai: Titik henti sementara otomatis dengan MaxText

MaxText adalah LLM berperforma tinggi, skalabilitas arbitrer, open source, dan telah diuji dengan baik yang ditulis dalam Python/JAX murni yang menargetkan Cloud TPU. MaxText berisi semua penyiapan yang diperlukan untuk menggunakan fitur Autocheckpoint.

File README MaxText menjelaskan dua cara untuk menjalankan MaxText dalam skala besar:

Menggunakan multihost_runner.py, direkomendasikan untuk eksperimen
Menggunakan multihost_job.py, direkomendasikan untuk produksi

Saat menggunakan multihost_runner.py, aktifkan Autocheckpoint dengan menetapkan flag autocheckpoint-enabled saat menyediakan resource yang diantrekan.

Saat menggunakan multihost_job.py, aktifkan Autocheckpoint dengan menentukan flag command line ENABLE_AUTOCHECKPOINT=true saat meluncurkan tugas.

Panduan memulai: Titik periksa otomatis dengan Pax di satu slice

Bagian ini memberikan contoh cara menyiapkan dan menggunakan Autocheckpoint dengan Pax pada satu slice. Dengan penyiapan yang sesuai:

Titik pemeriksaan akan disimpan saat peristiwa pemeliharaan terjadi.
Cloud TPU akan melakukan pemeliharaan pada VM TPU yang terpengaruh setelah checkpoint disimpan.
Setelah Cloud TPU menyelesaikan pemeliharaan, Anda dapat menggunakan VM TPU seperti biasa.

Gunakan flag autocheckpoint-enabled saat membuat VM TPU atau meminta resource dalam antrean.

Contoh:

Menetapkan variabel lingkungan:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=zone-you-want-to-use
export ACCELERATOR_TYPE=your-accelerator-type
export RUNTIME_VERSION=tpu-ubuntu2204-base

Deskripsi variabel lingkungan

Variabel	Deskripsi
`PROJECT_ID`	Project ID Google Cloud Anda. Gunakan project yang ada atau buat project baru.
`TPU_NAME`	Nama TPU.
`ZONE`	Zona tempat VM TPU akan dibuat. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona TPU.
`ACCELERATOR_TYPE`	Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.
`RUNTIME_VERSION`	Versi software Cloud TPU.

Tetapkan project ID dan zona Anda dalam konfigurasi aktif:

gcloud config set project $PROJECT_ID
gcloud config set compute/zone $ZONE

Buat TPU:

gcloud alpha compute tpus tpu-vm create $TPU_NAME \
    --accelerator-type $ACCELERATOR_TYPE \
    --version $RUNTIME_VERSION \
    --autocheckpoint-enabled

Hubungkan ke TPU menggunakan SSH:

gcloud compute tpus tpu-vm ssh $TPU_NAME

Menginstal Pax di satu slice

Fitur Autocheckpoint berfungsi di Pax versi 1.1.0 dan yang lebih baru. Di VM TPU, instal jax[tpu] dan paxml terbaru:
```
pip install paxml && pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```

Konfigurasikan model LmCloudSpmd2B. Sebelum menjalankan skrip pelatihan, ubah ICI_MESH_SHAPE menjadi [1, 8, 1]:

@experiment_registry.register
class LmCloudSpmd2B(LmCloudSpmd):

    """SPMD model with 2B params.

    Global batch size = 2 * 2 * 1 * 32 = 128
    """
    PERCORE_BATCH_SIZE = 8

    NUM_LAYERS = 18
    MODEL_DIMS = 3072
    HIDDEN_DIMS = MODEL_DIMS * 4

    CHECKPOINT_POLICY = layers.AutodiffCheckpointType.SAVE_NOTHING
    ICI_MESH_SHAPE = [1, 8, 1]

Luncurkan pelatihan dengan konfigurasi yang sesuai.

Contoh berikut menunjukkan cara mengonfigurasi model LmCloudSpmd2B untuk menyimpan titik pemeriksaan yang dipicu oleh Autocheckpoint ke bucket Cloud Storage. Ganti your-storage-bucket dengan nama bucket yang ada, atau buat bucket baru.
```
export JOB_LOG_DIR=gs://your-storage-bucket

{ python3 .local/lib/python3.10/site-packages/paxml/main.py \
    --jax_fully_async_checkpoint=1 \
    --exit_after_ondemand_checkpoint=1 \
    --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2B \
    --job_log_dir=$JOB_LOG_DIR; } 2>&1 | tee pax_logs.txt
```
Perhatikan dua flag yang diteruskan ke perintah:
- jax_fully_async_checkpoint: Jika flag ini diaktifkan, orbax.checkpoint.AsyncCheckpointer akan digunakan. Class AsyncCheckpointer otomatis menyimpan titik pemeriksaan saat skrip pelatihan menerima sinyal SIGTERM.
- exit_after_ondemand_checkpoint: Dengan flag ini diaktifkan, proses TPU akan keluar setelah Autocheckpoint berhasil disimpan, yang memicu pemeliharaan untuk segera dilakukan. Jika Anda tidak menggunakan flag ini, pelatihan akan dilanjutkan setelah titik pemeriksaan disimpan dan Cloud TPU akan menunggu waktu tunggu habis (5 menit) sebelum melakukan pemeliharaan yang diperlukan.

Titik pemeriksaan otomatis dengan Orbax

Fitur Autocheckpoint tidak terbatas pada MaxText atau Pax. Setiap framework yang dapat menangkap sinyal SIGTERM dan memulai proses checkpointing berfungsi dengan infrastruktur yang disediakan oleh Autocheckpoint. Orbax, namespace yang menyediakan library utilitas umum untuk pengguna JAX, menyediakan kemampuan ini.

Seperti yang dijelaskan dalam dokumentasi Orbax, kemampuan ini diaktifkan secara default untuk pengguna orbax.checkpoint.CheckpointManager. Metode save yang dipanggil setelah setiap langkah akan otomatis memeriksa apakah peristiwa pemeliharaan akan segera terjadi, dan jika ya, akan menyimpan titik pemeriksaan meskipun nomor langkah bukan kelipatan save_interval_steps. Dokumentasi GitHub juga mengilustrasikan cara membuat pelatihan keluar setelah menyimpan Autocheckpoint, dengan modifikasi pada kode pengguna.