Melatih model menggunakan TPU v5e

Dengan footprint 256 chip yang lebih kecil per Pod, TPU v5e dioptimalkan untuk menjadi produk bernilai tinggi untuk pelatihan, fine-tuning, dan inferensi Transformer, text-to-image, dan Jaringan Neural Konvolusional (CNN). Untuk mengetahui informasi selengkapnya tentang penggunaan Cloud TPU v5e untuk penayangan, lihat Inferensi menggunakan v5e.

Untuk mengetahui informasi selengkapnya tentang hardware dan konfigurasi TPU v5e Cloud TPU, lihat TPU v5e.

Mulai

Bagian berikut menjelaskan cara mulai menggunakan TPU v5e.

Kuota permintaan

Anda memerlukan kuota untuk menggunakan TPU v5e untuk pelatihan. Ada berbagai jenis kuota untuk TPU sesuai permintaan, TPU yang dicadangkan, dan VM Spot TPU. Ada kuota terpisah yang diperlukan jika Anda menggunakan TPU v5e untuk inferensi. Untuk mengetahui informasi selengkapnya tentang kuota, lihat Kuota. Untuk meminta kuota TPU v5e, hubungi Tim Penjualan Cloud.

Buat akun dan project Google Cloud

Anda memerlukan Google Cloud akun dan project untuk menggunakan Cloud TPU. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan lingkungan Cloud TPU.

Buat Cloud TPU

Praktik terbaiknya adalah menyediakan Cloud TPU v5es sebagai resource dalam antrean menggunakan perintah queued-resource create. Untuk mengetahui informasi selengkapnya, lihat Mengelola resource dalam antrean.

Anda juga dapat menggunakan Create Node API (gcloud compute tpus tpu-vm create) untuk menyediakan Cloud TPU v5e. Untuk mengetahui informasi selengkapnya, lihat Mengelola resource TPU.

Untuk mengetahui informasi selengkapnya tentang konfigurasi v5e yang tersedia untuk pelatihan, lihat Jenis Cloud TPU v5e untuk pelatihan.

Penyiapan framework

Bagian ini menjelaskan proses penyiapan umum untuk pelatihan model kustom menggunakan JAX atau PyTorch dengan TPU v5e.

Untuk mengetahui petunjuk penyiapan inferensi, lihat pengantar inferensi v5e.

Tentukan beberapa variabel lingkungan:

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id

Penyiapan untuk JAX

Jika Anda memiliki bentuk irisan yang lebih besar dari 8 chip, Anda akan memiliki beberapa VM dalam satu irisan. Dalam hal ini, Anda harus menggunakan tanda --worker=all untuk menjalankan penginstalan di semua VM TPU dalam satu langkah tanpa menggunakan SSH untuk login ke setiap VM secara terpisah:

gcloud compute tpus tpu-vm ssh $TPU_NAME  \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Deskripsi tanda perintah

  • TPU_NAME: ID teks yang ditetapkan pengguna dari TPU yang dibuat saat permintaan resource dalam antrean dialokasikan.
  • PROJECT_ID: Google Cloud Nama Project. Gunakan project yang ada atau buat project baru di Siapkan Google Cloud project
  • ZONE: Lihat dokumen Region dan zona TPU untuk mengetahui zona yang didukung.
  • worker: VM TPU yang memiliki akses ke TPU yang mendasarinya.

Anda dapat menjalankan perintah berikut untuk memeriksa jumlah perangkat (output yang ditampilkan di sini dihasilkan dengan slice v5litepod-16). Kode ini menguji bahwa semuanya diinstal dengan benar dengan memeriksa apakah JAX melihat TensorCore Cloud TPU dan dapat menjalankan operasi dasar:

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID