Ringkasan Multislice Cloud TPU

Cloud TPU Multislice adalah teknologi penskalaan performa full stack yang memungkinkan tugas pelatihan menggunakan beberapa slice TPU dalam satu Pod atau pada slice di beberapa Pod dengan paralelisme data sederhana. Dengan chip TPU v4, ini berarti tugas pelatihan dapat menggunakan lebih dari 4.096 chip dalam satu proses. Untuk tugas pelatihan yang memerlukan kurang dari 4.096 chip, satu irisan dapat menawarkan performa terbaik. Namun, beberapa irisan yang lebih kecil akan lebih mudah tersedia, sehingga memungkinkan waktu startup yang lebih cepat ketika Multislice digunakan dengan irisan yang lebih kecil.

Saat di-deploy dalam konfigurasi Multislice, chip TPU di setiap slice berkomunikasi melalui inter-chip-interconnect (ICI). Chip TPU di berbagai irisan berkomunikasi dengan mentransfer data ke CPU (host) yang kemudian mengirimkan data melalui jaringan pusat data (DCN).

Developer tidak perlu menulis kode untuk menerapkan komunikasi DCN antar-slice. Compiler XLA membuat kode tersebut untuk Anda dan tumpang tindih dengan komunikasi dengan komputasi untuk performa maksimum.

Konsep

Jenis akselerator

Bentuk setiap slice TPU yang terdiri dari Multislice. Setiap irisan dalam permintaan multislice memiliki jenis akselerator yang sama. Jenis akselerator terdiri dari jenis TPU (v4 atau v5e) yang diikuti dengan jumlah TensorCore. Misalnya, v4-128 menentukan TPU v4 dengan 128 TensorCore.

Perbaikan otomatis

Saat slice mengalami peristiwa pemeliharaan, preemption, atau kegagalan hardware, Cloud TPU akan membuat slice baru. Dalam kasus yang jarang terjadi, jika resource tidak cukup untuk membuat slice baru, pembuatan tidak akan selesai sampai hardware tersedia. Setelah slice baru dibuat, semua slice lain di lingkungan Multislice akan dimulai ulang sehingga pelatihan dapat dilanjutkan.Dengan skrip startup yang dikonfigurasi dengan benar, skrip pelatihan dapat otomatis diluncurkan kembali tanpa intervensi pengguna, memuat, dan melanjutkan dari checkpoint terbaru.

Dataset

Data yang digunakan oleh model untuk pelatihan atau inferensi.

Jaringan Pusat Data (DCN)

Jaringan dengan latensi lebih tinggi dengan throughput lebih rendah (jika dibandingkan dengan ICI) yang menghubungkan slice TPU dalam konfigurasi Multislice.

Penjadwalan geng

Saat semua slice TPU disediakan bersama, pada saat yang sama, menjamin bahwa semua slice atau tidak berhasil disediakan.

Host

Host adalah komputer fisik yang menjalankan VM. Host dapat menjalankan maksimal empat VM sekaligus. Setiap VM memiliki TPU khusus.

Inferensi

Memuat model machine learning terlatih ke host dan membuat prediksi pada data.

Interchip interconnect (ICI)

Link internal berkecepatan tinggi dan berlatensi rendah yang menghubungkan TPU di dalam Pod TPU.

Multiiris

Dua atau lebih potongan chip TPU yang dapat berkomunikasi melalui DCN.

Node

Dalam konteks Multislice, node mengacu pada satu slice TPU. Setiap irisan TPU dalam Multislice diberi ID node.

Pod

Kumpulan chip TPU yang terhubung melalui antarmuka jaringan ICI khusus. Dengan Pod, Anda dapat mendistribusikan beban pemrosesan di beberapa TPU.

Resource dalam antrean (QR)

Representasi resource TPU, yang digunakan untuk mengantrekan dan mengelola permintaan untuk lingkungan TPU slice tunggal atau Multislice.

Skrip startup

Skrip startup Compute Engine standar yang dijalankan setiap kali VM di-booting atau dimulai ulang. Untuk Multislice, parameter ini ditentukan dalam permintaan pembuatan QR. Untuk mengetahui informasi selengkapnya tentang skrip startup Cloud TPU, baca Mengelola resource TPU.

Irisan TPU

Subbagian logis dari Pod TPU yang terdiri dari chip TPU. Semua chip dalam slice berkomunikasi satu sama lain menggunakan jaringan ICI.

VM TPU

Virtual machine yang menjalankan Linux yang memiliki akses ke TPU yang mendasarinya. Untuk TPU v4, setiap VM TPU memiliki akses langsung ke empat chip. Terkadang kita menyebut VM TPU sebagai pekerja.

Tensor

Struktur data yang digunakan untuk merepresentasikan data multidimensi dalam model machine learning.

Tensor processing unit (TPU)

Chip akselerasi ML yang dikembangkan secara internal Google. Fitur ini dirancang untuk menawarkan komputasi yang cepat dan hemat daya untuk tugas-tugas machine learning utama seperti perkalian matriks.

Jenis kapasitas Cloud TPU

TPU dapat dibuat dari tiga jenis kapasitas (lihat Opsi Penggunaan di bagian Cara kerja penetapan harga TPU) :

• Reservasi: Menargetkan kuota yang direservasi. Untuk menggunakan kuota yang dicadangkan, Anda harus memiliki perjanjian reservasi dengan Google. Gunakan flag --reserved saat membuat resource.
• Preemptible: Menargetkan kuota yang dapat dihentikan. Resource Anda mungkin di-preempt untuk memberi ruang bagi permintaan tugas dengan prioritas yang lebih tinggi. Gunakan flag --best-effort saat membuat resource.
• On-demand: Menarget kuota on demand, yang tidak memerlukan reservasi dan tidak akan di-preempt. Permintaan TPU akan diantrekan ke antrean kuota on demand yang ditawarkan oleh Cloud TPU. Ketersediaan resource tidak dijamin. Dipilih secara default, tidak diperlukan tanda.

Catatan: Cloud TPU membatasi ukuran irisan untuk semua jenis akselerator. Untuk mengetahui batas ukuran slice bagi permintaan on demand, lihat Ukuran slice default untuk QR on demand.

Mulai

Jika Anda belum pernah menggunakan TPU, mulailah dengan menginstal Google Cloud CLI, dan menyiapkan lingkungan Cloud TPU Anda. Untuk menggunakan Multislice, resource TPU Anda harus dikelola sebagai Resource yang Diantrekan.

Jika Anda adalah pengguna TPU v4 lama dan memiliki reservasi, Anda mungkin perlu memigrasikan reservasi ke sistem reservasi baru. Untuk mengetahui informasi selengkapnya, hubungi perwakilan akun Google Cloud Anda.

Contoh pendahuluan

Tutorial ini menggunakan kode dari repo GitHub MaxText. MaxText adalah LLM dasar berperforma tinggi, skalabel secara arbitrer, open source, dan sudah teruji dengan baik yang ditulis di Python dan Jax. MaxText dirancang untuk berlatih secara efisien di Cloud TPU.

Kode di shardings.py dirancang untuk membantu Anda mulai bereksperimen dengan berbagai opsi paralelisme. Misalnya, paralelisme data, paralelisme data dengan sharding (FSDP), dan paralelisme tensor. Kode diskalakan dari lingkungan satu slice ke Multislice.

Paralelisme ICI

ICI mengacu pada interkoneksi berkecepatan tinggi yang menghubungkan TPU dalam satu irisan. Sharding ICI berkaitan dengan sharding dalam sebuah slice. shardings.py menyediakan tiga parameter paralelisme ICI:

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

Nilai yang Anda tetapkan untuk parameter ini akan menentukan jumlah shard untuk setiap metode paralelisme.

Input ini harus dibatasi agar ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism sama dengan jumlah chip dalam irisan.

Tabel berikut menunjukkan contoh input pengguna untuk paralelisme ICI untuk empat chip yang tersedia di v4-8:

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
FSDP 4 Arah	1	4	1
Paralelisme Tensor 4 arah	1	1	4
FSDP 2 arah + paralelisme Tensor 2 arah	1	2	2

Perlu diperhatikan bahwa ici_data_parallelism harus dibiarkan sebagai 1 dalam sebagian besar kasus karena jaringan ICI cukup cepat untuk hampir selalu lebih memilih FSDP daripada paralelisme data.

Contoh ini mengasumsikan bahwa Anda telah memahami cara menjalankan kode pada satu irisan TPU, seperti dalam artikel Menjalankan penghitungan di VM Cloud TPU menggunakan JAX. Contoh ini menunjukkan cara menjalankan shardings.py pada satu irisan.

1. Menyiapkan lingkungan:

   $ gcloud auth login
   $ gcloud config set project your-project-id
   $ gcloud config set compute/zone your-zone

2. Membuat kunci SSH untuk gcloud. Sebaiknya biarkan sandi kosong (tekan enter dua kali setelah menjalankan perintah berikut). Jika Anda diminta bahwa file google_compute_engine sudah ada, ganti versi yang ada.

   $ ssh-keygen -f ~/.ssh/google_compute_engine

3. Sediakan TPU Anda dengan perintah berikut:

   $ gcloud alpha compute tpus queued-resources \
   create your-qr-id \
   --accelerator-type your-accelerator-type \
   --runtime-version tpu-ubuntu2204-base \
   --node-id qr-id \
   [--reserved |--best-effort]

   Deskripsi flag perintah

   your-qr-id

   String yang ditentukan pengguna yang mengidentifikasi permintaan QR.

   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](/tpu/docs/supported-tpu-configurations#tpu_software_version).

   node-id

   ID resource TPU yang akan dibuat sebagai respons terhadap permintaan QR.

   reserved

   Gunakan kuota yang dicadangkan saat membuat slice.

   best-effort

   Gunakan kuota upaya terbaik saat membuat slice [Default].

   Google Cloud CLI tidak mendukung semua opsi pembuatan QR, seperti tag. Untuk mengetahui informasi selengkapnya, lihat Membuat QR.

4. Tunggu hingga QR berada dalam status ACTIVE yang berarti node pekerja berada dalam status READY. Setelah penyediaan QR dimulai, mungkin perlu waktu satu hingga lima menit untuk menyelesaikan, bergantung pada ukuran QR. Anda dapat memeriksa status permintaan QR menggunakan perintah berikut:

   $ gcloud compute tpus queued-resources \
     list --filter=your-qr-id

5. Slice v4-8 memiliki satu VM TPU. Hubungkan ke VM TPU menggunakan SSH:

   $ gcloud compute tpus tpu-vm ssh your-qr-id

6. Clone MaxText (yang mencakup shardings.py) ke VM TPU Anda.

7. Dalam direktori repositori MaxText, jalankan skrip penyiapan untuk menginstal JAX dan dependensi lainnya pada slice TPU Anda. Skrip penyiapan membutuhkan waktu beberapa menit untuk dijalankan.

   $ bash setup.sh

8. Jalankan perintah berikut untuk menjalankan shardings.py di slice TPU Anda.

   $ python3 pedagogical_examples/shardings.py \
     --ici_fsdp_parallelism 4 \
     --batch_size 131072 \
     --embedding_dimension 2048

   Anda dapat melihat hasilnya di log. TPU Anda akan mencapai sekitar 260 TFLOP per detik atau pemanfaatan FLOP lebih dari 90%secara mengesankan. Dalam hal ini, kita telah memilih kira-kira batch maksimum yang sesuai dengan High Bandwidth Memory (HBM) TPU.

9. Jangan ragu untuk mempelajari strategi sharding lain melalui ICI. Misalnya, Anda dapat mencoba kombinasi berikut:

   $ python3 pedagogical_examples/shardings.py \
     --ici_tensor_parallelism 4 \
     --batch_size 131072 \
     --embedding_dimension 2048

10. Hapus slice QR dan TPU jika sudah selesai. Anda harus menjalankan langkah-langkah pembersihan ini dari lingkungan tempat Anda menyiapkan slice (jalankan exit terlebih dahulu untuk keluar dari sesi SSH). Proses penghapusan akan memerlukan waktu dua hingga lima menit, dan dapat dijalankan di latar belakang dengan flag --async opsional.

    Penting: QR harus dihapus setelah digunakan. Jika tidak, mereka akan terus menggunakan kuota QR Anda yang dilacak secara terpisah dari kuota TPU.

    $ gcloud compute tpus queued-resources
      delete your-qr-id --force (--async)

Sharding multislice menggunakan paralelisme DCN

Skrip shardings.py menggunakan tiga parameter yang menentukan paralelisme DCN, sesuai dengan jumlah shard dari setiap jenis paralelisme data:

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

Nilai parameter ini harus dibatasi agar dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism sama dengan jumlah irisan.

Sebagai contoh untuk dua irisan, gunakan --dcn_data_parallelism = 2.

	dcn_data_parallelism	dcn_fsdp_parallelism	dcn_tensor_parallelism	# irisan
Paralelisme data 2 arah	2	1	1	2

dcn_tensor_parallelism harus selalu ditetapkan ke 1 karena DCN tidak cocok untuk sharding tersebut. Untuk beban kerja LLM standar pada chip v4, dcn_fsdp_parallelism juga harus ditetapkan ke 1, sehingga dcn_data_parallelism harus ditetapkan ke jumlah slice, tetapi cara ini bergantung pada aplikasi.

Meningkatkan jumlah irisan (dengan asumsi ukuran irisan dan batch per irisan konstan), Anda akan meningkatkan jumlah paralelisme data.

Menjalankan shardings.py di lingkungan Multislice

Anda dapat menjalankan shardings.py di lingkungan Multislice menggunakan multihost_runner.py atau dengan menjalankan shardings.py di setiap VM TPU. Di sini kita menggunakan multihost_runner.py. Langkah-langkah berikut sangat mirip dengan Memulai: Eksperimen Cepat di Beberapa slices dari repositori MaxText, hanya saja di sini kami menjalankan shardings.py, bukan LLM yang lebih kompleks di train.py.

Alat multihost_runner.py dioptimalkan untuk eksperimen cepat, yang berulang kali menggunakan kembali TPU yang sama. Karena skrip multihost_runner.py bergantung pada koneksi SSH yang tahan lama, kami tidak merekomendasikannya untuk tugas yang berjalan lama. Jika Anda ingin menjalankan tugas yang lebih lama (misalnya, jam atau hari), sebaiknya gunakan multihost_job.py.

Dalam tutorial ini, kami menggunakan istilah runner untuk menunjukkan mesin tempat Anda menjalankan skrip multihost_runner.py. Kami menggunakan istilah pekerja untuk menunjukkan VM TPU yang membentuk bagian Anda. Anda dapat menjalankan multihost_runner.py di mesin lokal atau VM Compute Engine apa pun dalam project yang sama dengan slice Anda. Menjalankan multihost_runner.py pada pekerja tidak didukung.

multihost_runner.py otomatis terhubung ke pekerja TPU menggunakan SSH.

Dalam contoh ini, kita menjalankan shardings.py pada dua slice v4-16, total empat VM dan 16 chip TPU. Anda dapat mengubah contoh untuk dijalankan di lebih banyak TPU.

Menyiapkan lingkungan Anda

1. Clone MaxText pada mesin runner Anda.

2. Buka direktori repositori.

3. Buat kunci SSH untuk gcloud. Sebaiknya biarkan sandi kosong (tekan enter dua kali setelah menjalankan perintah berikut). Jika Anda diminta bahwa file google_compute_engine sudah ada, pilih untuk tidak mempertahankan versi yang sudah ada.

     $ ssh-keygen -f ~/.ssh/google_compute_engine
     

4. Tambahkan variabel lingkungan untuk menetapkan jumlah slice TPU ke 2.

     $ export SLICE_COUNT=2
     

5. Buat lingkungan Multislice menggunakan queued-resources create.

   Catatan: Pastikan Anda memiliki kuota masing-masing sebelum memilih --reserved, atau --best_effort, atau kuota on-demand default. Untuk mengetahui informasi tentang jenis kuota, lihat Kebijakan Kuota.

   Perintah berikut menunjukkan cara membuat TPU Multislice v4. Untuk menggunakan v5e, tentukan accelerator-type v5e (misalnya v5litepod-16) dan runtime-version v5e (v2-alpha-tpuv5-lite).

     $ gcloud alpha compute tpus queued-resources 
   
     create your-qr-id 
   
     --accelerator-type=your-accelerator-type 
   
     --runtime-version=tpu-vm-runtime-version 
   
     --node-count=node-count 
   
     --node-prefix=your-qr-id 
   
     [--reserved|--best-effort]
     

   Deskripsi flag perintah

   your-qr-id

   String yang ditentukan pengguna yang mengidentifikasi permintaan QR.

   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.

   node-count

   Jumlah irisan yang akan dibuat.

   node-prefix

   Awalan yang digunakan untuk membuat nama setiap irisan. Angka ditambahkan ke awalan untuk setiap irisan. Misalnya, jika Anda menetapkan node-prefix ke mySlice, irisan akan diberi nama: mySlice-0, mySlice-1, dan seterusnya.

   reserved

   Gunakan kuota yang dicadangkan saat membuat slice.

   best-effort

   Gunakan kuota upaya terbaik saat membuat slice [Default].

6. Saat penyediaan QR dimulai, mungkin perlu waktu hingga lima menit untuk menyelesaikannya, bergantung pada ukuran QR. Tunggu hingga Queued Resource (QR) berada dalam status ACTIVE. Anda dapat memeriksa status permintaan QR menggunakan perintah berikut:

   $ gcloud compute tpus queued-resources list \
   --filter=your-qr-id

   Tindakan ini akan menghasilkan output yang terlihat seperti ini:

   NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
   ...
   que-res-id  us-central2-b  4           v4-16             ACTIVE
   ...

   Hubungi perwakilan akun Google Cloud Anda jika status QR menampilkan status WAITING_FOR_RESOURCES atau PROVISIONING selama lebih dari 15 menit.

7. Instal dependensi.

   $ python3 multihost_runner.py \
     --TPU_PREFIX=your-qr-id \
     --COMMAND="bash setup.sh"

8. Menjalankan shardings.py pada setiap pekerja menggunakan multihost_runner.py.

   $ python3 multihost_runner.py \
     --TPU_PREFIX=your-qr-id \
     --COMMAND="python3 pedagogical_examples/shardings.py \
     --dcn_data_parallelism $SLICE_COUNT \
     --ici_fsdp_parallelism 8 \
     --batch_size 131072 \
     --embedding_dimension 2048"

   Anda akan melihat performa sekitar 230 TFLOP per detik dalam file log.

9. Bersihkan TPU dan QR saat selesai. Proses penghapusan akan memerlukan waktu dua hingga lima menit, dan dapat dijalankan di latar belakang dengan flag --async opsional.

Menskalakan workload ke Multislice

Sebelum menjalankan model Anda di lingkungan Multislice, buat perubahan kode berikut:

• Gunakan jax.experimental.mesh_utils.create_hybrid_device_mesh, bukan jax.experimental.mesh_utils.create_device_mesh saat membuat mesh Anda.

Ini akan menjadi satu-satunya perubahan kode yang diperlukan saat beralih ke Multislice. Untuk mencapai performa tinggi, DCN harus dipetakan ke data paralel, sumbu paralel data yang di-sharding sepenuhnya, atau sumbu paralel pipeline. Pertimbangan performa dan strategi sharding dibahas secara lebih mendetail di Melakukan Sharding dengan Multislice untuk Performa Maksimal.

Untuk memvalidasi bahwa kode Anda dapat mengakses semua perangkat, Anda dapat menyatakan bahwa len(jax.devices()) sama dengan jumlah chip di lingkungan Multislice. Misalnya, jika menggunakan empat irisan v4-16, Anda memiliki delapan chip per irisan * 4 irisan, sehingga len(jax.devices()) akan menampilkan 32.

Memilih ukuran irisan untuk lingkungan Multislice

Untuk mendapatkan kecepatan linear, tambahkan irisan baru yang berukuran sama dengan potongan yang ada. Misalnya, jika Anda menggunakan irisan v4-512, Multislice akan mencapai performa sekitar dua kali lipat dengan menambahkan slice v4-512 kedua dan menggandakan ukuran batch global Anda. Untuk informasi selengkapnya, lihat Melakukan Sharding dengan Multislice untuk Performa Maksimal.

Menjalankan Tugas Anda pada beberapa irisan

Ada tiga pendekatan berbeda untuk menjalankan beban kerja kustom di lingkungan Multislice:

1. Menggunakan skrip runner eksperimen, multihost_runner.py
2. Menggunakan skrip runner produksi, multihost_job.py
3. Menggunakan pendekatan manual

Skrip runner eksperimen

Skrip multihost_runner.py mendistribusikan kode ke lingkungan Multislice yang ada, dan menjalankan perintah Anda di setiap host, menyalin kembali log, dan melacak status error setiap perintah. Skrip multihost_runner.py didokumentasikan dalam README MaxText.

Karena mempertahankan koneksi SSH persisten, multihost_runner.py hanya cocok untuk eksperimen berukuran sedang yang relatif singkat. Anda dapat menyesuaikan langkah-langkah dalam tutorial multihost_runner.py untuk konfigurasi workload dan hardware.

Skrip runner produksi

Untuk tugas produksi yang memerlukan ketahanan terhadap kegagalan hardware dan preemption lainnya, sebaiknya langsung integrasikan dengan Create Queued Resource API. Sebagai contoh praktis, kami menyediakan multihost_job.py, yang memicu panggilan API Resource Antrean yang Dibuat dengan skrip startup yang sesuai untuk menjalankan pelatihan dan melanjutkan pelatihan saat terjadi preemption. Skrip multihost_job.py didokumentasikan dalam README MaxText.

Karena harus menyediakan resource untuk setiap run, multihost_job.py tidak menyediakan siklus iterasi secepat multihost_runner.py.

Pendekatan manual

Sebaiknya gunakan atau sesuaikan multihost_runner.py atau multihost_job.py untuk menjalankan beban kerja kustom Anda dalam konfigurasi Multislice. Namun, jika Anda lebih memilih untuk menyediakan dan mengelola lingkungan menggunakan perintah QR secara langsung, lihat Mengelola Lingkungan Multislice.

Mengelola lingkungan Multislice

Untuk menyediakan dan mengelola QR secara manual tanpa menggunakan alat yang disediakan dalam repo MaxText, baca bagian berikut.

Buat QR

Tetapkan variabel lingkungan berikut sebelum menyediakan kapasitas:

  $ export your-qr-id=your-queued-resource-id
  $ export PROJECT=your-project-name
  $ export ZONE=us-central2-b
  $ export NETWORK_NAME=your-network-name
  $ export SUBNETWORK_NAME=your-subnetwork-name
  $ export RUNTIME_VERSION=tpu-ubuntu2204-base
  $ export ACCELERATOR_TYPE=v4-16
  $ export SLICE_COUNT=4
  $ export STARTUP_SCRIPT="#!/bin/bash\n ..."
  $ gcloud config set project project-name
  $ gcloud config set compute/zone zone

Input	Deskripsi
your-qr-id	ID QR yang ditetapkan pengguna.
PROYEK	Nama Project Google Cloud
ZONA	us-central2-b
NETWORK_NAME	Nama jaringan VPC.
SUBNETWORK_NAME	Nama subnet di jaringan VPC
RUNTIME_VERSION	{i>tpu-ubuntu2204-base<i}
ACCELERATOR_TYPE	v4-16
EXAMPLE_TAG_1, EXAMPLE_TAG_2...	Tag yang digunakan untuk mengidentifikasi sumber atau target yang valid untuk firewall jaringan
SLICE_COUNT	Jumlah irisan. Dibatasi maksimum 256 irisan.
STARTUP_SCRIPT	Jika ditambahkan ke permintaan pembuatan, skrip startup dapat dijalankan setiap kali slice TPU disediakan atau dimulai ulang dan jika slice TPU diperbaiki atau direset.

Buat permintaan QR menggunakan gcloud

$ gcloud alpha compute tpus queued-resources \
  create ${your-qr-id} \
  --project your-project-id \
  --zone your-zone \
  --node-count ${SLICE_COUNT} \
  --accelerator-type ${ACCELERATOR_TYPE} \
  --runtime-version ${RUNTIME_VERSION} \
  --network ${NETWORK_NAME} \
  --subnetwork ${SUBNETWORK_NAME} \
  --tags ${EXAMPLE_TAG_1},${EXAMPLE_TAG_2} \ --metadata=startup-script='${STARTUP_SCRIPT}'
  [--reserved|--best-effort]
  

Deskripsi flag perintah

your-qr-id

String yang ditentukan pengguna yang mengidentifikasi permintaan QR.

project

String yang ditentukan pengguna yang mengidentifikasi permintaan QR.

zone

Zona Google Cloud tempat pembuatan QR.

node-count

Jumlah irisan yang akan dibuat.

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.

network

Nama jaringan VPC tempat untuk melampirkan resource TPU.

subnetwork

Nama subnetwork VPC tempat untuk memasang resource TPU.

reserved

Gunakan kuota yang dicadangkan saat membuat slice.

best-effort

Gunakan kuota upaya terbaik saat membuat slice [Default].

Pastikan Anda memiliki kuota masing-masing sebelum memilih --reserved, atau --best_effort, atau kuota on-demand default. Untuk mengetahui informasi mengenai jenis kuota, lihat Kebijakan Kuota.

Buat permintaan QR menggunakan curl

Buat file bernama queued-resource-req.json dan salin JSON berikut ke dalamnya.

{
  "guaranteed": { "reserved": true },
  "tpu": {
    "node_spec": [
    {
      "parent": "projects/your-project-number/locations/your-zone",
        "node": {
          "accelerator_type": "accelerator-type",
          "runtime_version": "tpu-vm-runtime-version",
          "network_config": {
            "network": "your-network-name",
            "subnetwork": "your-subnetwork-name",
            "enable_external_ips": true
          },
          "tags" : ["example-tag-1"]
          "metadata": {
            "startup-script": "your-startup-script"
          }
      },
      "multi_node_params": {
        "node_count": slice-count,
        "node_id_prefix": "your-queued-resource-id"
      }
    }
    ]
  }
}

Catatan: Ganti nilai berikut:

• your-project-number - Nomor project Google Cloud Anda
• your-zone - Zona tempat Anda ingin membuat QR
• accelerator-type - Versi dan ukuran irisan tunggal
• tpu-vm-runtime-version - Versi runtime VM TPU
• your-network-name - Opsional, jaringan tempat QR akan dilampirkan
• your-subnetwork-name - Opsional, subnetwork tempat QR akan dilampirkan
• example-tag-1 - Opsional, string tag arbitrer
• your-startup-script - Skrip startup yang akan dijalankan saat QR dialokasikan
• slice-count - Jumlah slice TPU di lingkungan Multislice Anda
• your-qr-id - ID yang disediakan pengguna untuk QR

Untuk informasi selengkapnya, lihat dokumentasi REST Queued Resource API untuk semua opsi yang tersedia.

Catatan: Meskipun node_spec adalah daftar, hanya satu entri node_spec yang didukung.

Untuk menggunakan kapasitas yang dapat dihentikan, ganti:

"guaranteed": { "reserved": true } dengan "best_effort": {}

Atau hapus garis untuk menggunakan kapasitas on demand default.

Kirim permintaan pembuatan QR dengan payload JSON:

  $ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d @queuedresourcereq.json https://tpu.googleapis.com/v2alpha1/projects/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=your-qr-id

Catatan: Ganti nilai berikut:

• your-project-id - ID Project Google Cloud Anda
• your-zone - Zona tempat Anda ingin membuat QR
• your-qr-id - ID yang disediakan pengguna untuk QR

Responsnya akan terlihat seperti berikut:

{
  "name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qr-guid>",
  "metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-11-01T00:17:05.742546311Z",
    "target": "projects/<your-project-id>/locations/<your-zone>/queuedResources/<your-qa-id>",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
  },
  "done": false
}

Gunakan nilai GUID di akhir nilai string untuk atribut name guna mendapatkan informasi tentang permintaan QR.

Mengambil status QR

Untuk mendapatkan status permintaan QR, gunakan perintah berikut:

  $ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/operations/operation-your-qr-guid

Catatan: Pada perintah sebelumnya, ganti perintah berikut:

• your-project-id - ID Project Google Cloud Anda
• your-zone - Zona tempat membuat QR
• your-qr-guid - GUID yang mengikuti name dalam output dari permintaan pembuatan QR.

Respons perintah ini berisi status operasi:

{
  "name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qa-guid>,
  "metadata": {...},
  "done": true,
  "response": {
    "@type": "type.googleapis.com/google.cloud.tpu.v2.QueuedResource",
    ...
    "state": {
      "state": "WAITING_FOR_RESOURCES"
    }
  }
}

Jika QR berhasil dibuat ("done = true"), status dalam kolom response akan menjadi WAITING_FOR_RESOURCES atau FAILED. Jika QR berada dalam status WAITING_FOR_RESOURCES, artinya QR telah diantrekan dan akan memulai penyediaan jika ada resource yang cukup. Jika QR berada dalam status FAILED, alasan kegagalan akan ada di output. Untuk informasi selengkapnya tentang kemungkinan status lainnya, lihat Panduan pengguna resource yang diantrekan.

Setelah operasi selesai, gunakan describe QR untuk memantau tahap QR.

Pada skenario yang jarang terjadi, Anda mungkin menemukan QR dalam status FAILED, sementara beberapa irisan berupa ACTIVE. Jika hal ini terjadi, hapus resource yang telah dibuat, dan coba lagi dalam beberapa menit, atau hubungi tim Cloud TPU untuk menyelesaikan masalah tersebut.

SSH dan menginstal dependensi

Menjalankan kode JAX pada slice Pod TPU menjelaskan cara terhubung ke VM TPU menggunakan SSH dalam satu irisan. Untuk terhubung ke semua VM TPU di lingkungan Multislice melalui SSH dan menginstal dependensi, gunakan perintah gcloud berikut:

  $ gcloud compute tpus queued-resources ssh ${your-qr-id} \
    --zone your-zone \
    --node=all \
    --worker=all \
    --command="command-to-run"
    --batch-size=4

Perintah gcloud ini mengirimkan perintah yang ditentukan ke semua pekerja dan node dalam QR menggunakan SSH. Perintah ini dikelompokkan ke dalam kelompok yang terdiri dari empat orang dan dikirim secara bersamaan. Batch perintah berikutnya dikirim saat batch saat ini menyelesaikan eksekusi. Jika terjadi kegagalan dengan salah satu perintah, pemrosesan akan dihentikan, dan tidak ada batch lebih lanjut yang dikirim. Untuk informasi selengkapnya, lihat Referensi API resource yang diantrekan. Jika jumlah slice yang Anda gunakan melebihi batas threading komputer lokal (juga disebut batas pengelompokan), Anda akan mengalami deadlock. Sebagai contoh, anggaplah batas pengelompokan pada mesin lokal Anda adalah 64. Jika Anda mencoba menjalankan skrip pelatihan di lebih dari 64 slice, misalnya 100, perintah SSH akan membagi slice menjadi beberapa bagian. Skrip ini akan menjalankan skrip pelatihan pada batch pertama dari 64 slice dan menunggu hingga skrip selesai sebelum menjalankan skrip di batch lainnya yang terdiri dari 36 slice. Namun, batch pertama dari 64 irisan tidak dapat selesai hingga 36 irisan yang tersisa mulai menjalankan skrip, sehingga menyebabkan deadlock.

Untuk mencegah skenario ini, Anda dapat menjalankan skrip pelatihan di latar belakang pada setiap VM dengan menambahkan ampersand (&) ke perintah skrip yang Anda tentukan dengan tanda --command. Saat Anda melakukannya, setelah memulai skrip pelatihan pada batch slice pertama, kontrol akan segera kembali ke perintah SSH. Kemudian, perintah SSH dapat mulai menjalankan skrip pelatihan pada batch 36 slice yang tersisa. Anda harus menyalurkan aliran data stdout dan stderr dengan tepat saat menjalankan perintah di latar belakang. Untuk meningkatkan paralelisme dalam QR yang sama, Anda dapat memilih bagian tertentu menggunakan parameter --node.

Penyiapan jaringan

Pastikan slice TPU dapat saling berkomunikasi dengan menjalankan langkah-langkah berikut. Instal JAX pada setiap irisan. Untuk mengetahui informasi selengkapnya, baca Menjalankan kode JAX pada slice Pod TPU. Nyatakan bahwa len(jax.devices()) sama dengan jumlah chip di lingkungan Multislice Anda. Untuk melakukannya, pada setiap irisan, jalankan:

  $ python3 -c 'import jax; print(jax.devices())'

Jika Anda menjalankan kode ini di empat irisan v4-16, ada delapan chip per irisan dan empat irisan, total 32 chip (perangkat) harus ditampilkan oleh jax.devices().

Mencantumkan QR

Anda dapat melihat status QR menggunakan perintah queued-resources list:

$ gcloud compute tpus queued-resources list

NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central2-b  4           v4-16             ACTIVE
...

Menjelaskan QR

Untuk melihat konfigurasi dan status QR yang mendetail, gunakan QR API yang dideskripsikan. Anda dapat memanggil API ini menggunakan gcloud atau curl.

Menggunakan gcloud:

$ gcloud compute tpus queued-resources describe ${your-qr-id}
...state:
 state: ACTIVE
...

Menggunakan curl:

$ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/queuedResources/${your-qr-id}
{
  "name": your-queued-res,
  "tpu": {
    "nodeSpec": [
      {
        ... // node 1
      },
      {
        ... // node 2
      },
      ...
    ]
  },
  ...
  "state": "ACTIVE"
}

state merepresentasikan status QR. Untuk informasi selengkapnya tentang kemungkinan status QR, lihat Resource yang diantrekan.

Memulai pekerjaan Anda di lingkungan yang disediakan

Anda dapat menjalankan beban kerja secara manual dengan menghubungkan ke semua host di setiap bagian melalui SSH dan menjalankan perintah berikut pada semua host.

$ gcloud compute tpus tpu-vm ssh your-qr-id \
  --zone=your-zone \
  --worker=all \
  --node=all \
  --command="command-to-run"

Mereset QR

ResetQueuedResource API dapat digunakan untuk mereset semua VM dalam QR ACTIVE. Mereset VM akan menghapus memori mesin secara paksa dan mereset VM ke status awal. Setiap data yang disimpan secara lokal akan tetap utuh dan skrip startup akan dipanggil setelah reset. ResetQueuedResource API dapat berguna saat Anda ingin memulai ulang semua TPU. Misalnya, saat pelatihan terhenti dan mereset semua VM lebih mudah daripada proses debug.

Reset semua VM dilakukan secara paralel dan operasi ResetQueuedResource memerlukan waktu satu hingga dua menit untuk diselesaikan. Untuk memanggil API, gunakan perintah berikut:

$ gcloud compute tpus queued-resources reset your-qr-id

Menghapus QR

Untuk melepaskan resource di akhir sesi pelatihan, hapus resource yang diantrekan dengan tanda --force. Proses penghapusan akan memerlukan waktu dua hingga lima menit, dan dapat dijalankan di latar belakang dengan flag --async opsional.

$ gcloud compute tpus queued-resources \
delete your-qr-id --force (--async)

Pemulihan kegagalan otomatis

Jika terjadi gangguan, Multislice menawarkan perbaikan bebas intervensi untuk slice yang terpengaruh dan mereset semua slice setelahnya. Slice yang terpengaruh akan diganti dengan slice baru dan slice lain yang seharusnya responsif akan direset. Jika tidak ada kapasitas yang tersedia untuk mengalokasikan slice pengganti, pelatihan akan berhenti.

Untuk melanjutkan pelatihan secara otomatis setelah terjadi gangguan, Anda harus menentukan skrip startup yang memeriksa dan memuat checkpoint tersimpan terakhir. Skrip startup Anda otomatis berjalan setiap kali slice dialokasikan ulang atau VM direset. Anda menentukan skrip startup dalam payload JSON yang Anda kirim ke API permintaan QR pembuatan.

Skrip startup berikut (digunakan dalam Membuat QR) memungkinkan Anda pulih secara otomatis dari kegagalan dan melanjutkan pelatihan dari checkpoint yang disimpan di bucket Cloud Storage selama pelatihan MaxText:

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 MaxText/train.py MaxText/configs/base.yml run_name=run_test_failure_recovery dcn_data_parallelism=4 ici_fsdp_parallelism=8 steps=10000 save_period=10 base_output_directory='gs://user1-us-central2'' EOF"
         }
     ...
     }
   ]
 }
}

Clone repo MaxText sebelum mencobanya.

Pembuatan profil dan proses debug

Pembuatan profil sama di lingkungan irisan tunggal dan Multislice. Untuk mengetahui informasi selengkapnya, lihat Membuat profil program JAX.

Mengoptimalkan pelatihan

Sharding dengan Multislice untuk performa maksimum

Untuk mencapai performa maksimum di lingkungan Multislice, Anda harus mempertimbangkan cara melakukan sharding di beberapa slice. Biasanya ada tiga pilihan (paralelisme data, paralelisme data dengan shard sepenuhnya, dan paralelisme pipeline). Sebaiknya jangan melakukan aktivasi sharding di seluruh dimensi model (terkadang disebut paralelisme tensor) karena memerlukan terlalu banyak bandwidth antar-slice. Untuk semua strategi ini, Anda dapat mempertahankan strategi sharding yang sama dalam slice yang pernah bekerja untuk Anda sebelumnya.

Sebaiknya mulai dengan paralelisme data yang murni. Menggunakan paralelisme data sharding sepenuhnya berguna untuk mengosongkan penggunaan memori. Kekurangannya adalah komunikasi antarslice menggunakan jaringan DCN dan akan memperlambat beban kerja Anda. Gunakan paralelisme pipeline hanya jika diperlukan berdasarkan ukuran batch (seperti yang dianalisis di bawah).

Kapan harus menggunakan paralelisme data

Paralelisme data murni akan berfungsi dengan baik jika Anda memiliki beban kerja yang berjalan dengan baik, tetapi Anda ingin meningkatkan performanya dengan menskalakan di beberapa bagian.

Untuk mencapai penskalaan yang kuat di beberapa bagian, jumlah waktu yang diperlukan untuk melakukan semua pengurangan melalui DCN harus kurang dari jumlah waktu yang diperlukan untuk melakukan penerusan mundur. DCN digunakan untuk komunikasi antara potongan dan merupakan faktor pembatas dalam throughput beban kerja.

Setiap chip TPU v4 berperforma pada puncak 275 * 10¹² FLOPS per detik.

Ada empat chip per host TPU, dan setiap host memiliki bandwidth jaringan maksimum 50 Gbps.

Artinya, intensitas aritmetika adalah 4 * 275 * 10¹² FLOPS / 50 Gbps = 22.000 FLOPS / bit.

Model Anda akan menggunakan bandwidth DCN 32 hingga 64 bit untuk setiap parameter per langkah. Jika Anda menggunakan dua irisan, model akan menggunakan bandwidth DCN 32 bit. Jika Anda menggunakan lebih dari dua irisan, compiler akan melakukan operasi pengurangan penuh pengacakan penuh dan Anda akan menggunakan bandwidth DCN hingga 64 bit untuk setiap parameter per langkah. Jumlah FLOPS yang diperlukan untuk setiap parameter akan bervariasi bergantung pada model Anda. Secara khusus, untuk model bahasa berbasis Transformer, jumlah FLOPS yang diperlukan untuk penerusan maju dan mundur sekitar 6 * B * P, dengan:

• B adalah ukuran tumpukan dalam token
• P adalah jumlah parameter

Jumlah FLOPS per parameter adalah 6 * B dan jumlah FLOPS per parameter selama penerusan mundur adalah 4 * B.

Untuk memastikan penskalaan yang kuat di beberapa bagian, pastikan intensitas operasional melebihi intensitas aritmetika hardware TPU. Untuk menghitung intensitas operasional, bagi jumlah FLOPS per parameter selama proses penerusan mundur dengan bandwidth jaringan (dalam bit) per parameter per langkah: Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

Oleh karena itu, untuk model bahasa berbasis Transformer, jika Anda menggunakan dua irisan: Operational intensity = 4 * B / 32

Jika Anda menggunakan lebih dari dua irisan: Operational intensity = 4 * B/64

Hal ini menyarankan ukuran batch minimum antara 176 ribu dan 352 ribu untuk model bahasa berbasis Transformer. Karena jaringan DCN dapat menghapus paket dalam waktu singkat, sebaiknya pertahankan margin error yang signifikan, yaitu men-deploy paralelisme data hanya jika ukuran batch per Pod minimal 350 ribu (dua Pod) hingga 700 ribu (banyak Pod).

Untuk arsitektur model lainnya, Anda harus memperkirakan runtime penerusan mundur per irisan (dengan mengatur waktu menggunakan profiler atau dengan menghitung FLOPS). Kemudian, Anda dapat membandingkannya dengan waktu proses yang diharapkan untuk mengurangi keseluruhan DCN dan mendapatkan perkiraan yang akurat terkait apakah paralelisme data cocok bagi Anda.

Kapan harus menggunakan paralelisme data dengan sharding penuh (FSDP)

Paralelisme data dengan sharding penuh (FSDP) menggabungkan paralelisme data (mem-sharding data di seluruh node) dengan sharding bobot di seluruh node. Untuk setiap operasi pada penerusan maju dan mundur, bobot dikumpulkan sehingga setiap irisan memiliki bobot yang dibutuhkan. Daripada menyinkronkan gradien menggunakan semua pengurangan, gradien akan mengurangi tersebar saat dihasilkan. Dengan cara ini, setiap irisan hanya mendapatkan gradien untuk bobot yang menjadi tanggung jawabnya.

Serupa dengan paralelisme data, FSDP akan memerlukan penskalaan ukuran batch global secara linear dengan jumlah irisan. FSDP akan menurunkan tekanan memori saat Anda menambah jumlah slice. Hal ini karena jumlah bobot dan status pengoptimal per bagian menurun, tetapi hal ini terjadi karena adanya kenaikan traffic jaringan dan kemungkinan pemblokiran yang lebih besar karena adanya kolektif yang tertunda.

Dalam praktiknya, FSDP lintas irisan adalah pilihan terbaik jika Anda meningkatkan batch per irisan, menyimpan lebih banyak aktivasi untuk meminimalkan re-materialisasi selama proses mundur atau meningkatkan jumlah parameter dalam jaringan neural Anda.

Operasi yang mencakup semua dan semua pengurangan di FSDP berfungsi mirip dengan operasi di DP, sehingga Anda dapat menentukan apakah beban kerja FSDP Anda dibatasi oleh performa DCN dengan cara yang sama seperti yang dijelaskan di bagian sebelumnya.

Kapan harus menggunakan paralelisme pipeline

Paralelisme pipeline menjadi relevan saat mencapai performa tinggi dengan strategi paralelisme lainnya yang memerlukan ukuran batch global lebih besar dari ukuran batch maksimum yang Anda pilih. Paralelisme pipeline memungkinkan irisan yang terdiri dari pipeline untuk "membagikan" batch. Namun, paralelisme pipeline memiliki dua kelemahan signifikan:

1. Hal ini menyebabkan "gelembung pipeline" ketika chip tidak ada aktivitas karena sedang menunggu data.
2. Implementasi ini memerlukan micro-batching yang mengurangi ukuran tumpukan yang efektif, intensitas aritmetika, dan pada akhirnya memodelkan penggunaan FLOP.

Paralelisme pipeline hanya boleh digunakan jika strategi paralelisme lainnya memerlukan ukuran batch global yang terlalu besar. Sebelum mencoba paralelisme pipeline, sebaiknya lakukan eksperimen untuk melihat secara empiris apakah konvergensi per sampel melambat pada ukuran tumpukan yang diperlukan untuk mencapai FSDP berperforma tinggi. FSDP cenderung mencapai pemanfaatan FLOP model yang lebih tinggi, tetapi jika konvergensi per sampel melambat saat ukuran tumpukan bertambah, paralelisme pipeline mungkin masih menjadi pilihan yang lebih baik. Sebagian besar beban kerja dapat menoleransi ukuran tumpukan yang cukup besar agar tidak mendapatkan manfaat dari paralelisme pipeline, tetapi beban kerja Anda mungkin berbeda.

Jika paralelisme pipeline diperlukan, sebaiknya gabungkan dengan paralelisme data atau FSDP. Dengan begitu, Anda dapat meminimalkan kedalaman pipeline sekaligus meningkatkan ukuran batch per pipeline hingga latensi DCN menjadi faktor yang lebih kecil dalam throughput. Konkretnya, jika Anda memiliki irisan N, pertimbangkan pipeline dengan replika kedalaman 2 dan N/2 dari paralelisme data, lalu pipeline kedalaman 4 dan replika N/4 dari paralelisme data, dan seterusnya, hingga batch per pipeline menjadi cukup besar sehingga kolektif DCN dapat disembunyikan di balik aritmetika dalam proses mundur. Hal ini akan meminimalkan perlambatan yang ditimbulkan oleh paralelisme pipeline sekaligus memungkinkan Anda melakukan penskalaan melewati batas ukuran batch global.

Praktik terbaik multislice

Pemuatan data

Selama pelatihan, kita berulang kali memuat batch dari set data untuk dimasukkan ke dalam model. Memiliki loader data asinkron yang efisien yang melakukan sharding batch di seluruh host sangat penting untuk menghindari kelaparan TPU pekerjaan. Loader data saat ini di MaxText membuat setiap host memuat subset contoh yang sama. Solusi ini memadai untuk teks, tetapi memerlukan pengoptimalan dalam model. Selain itu, MaxText belum menawarkan snapshot deterministik yang memungkinkan iterator data memuat data yang sama sebelum dan sesudah preemption.

Checkpoint

Library checkpoint Orbax menyediakan primitif untuk checkpoint JAX PyTrees ke penyimpanan lokal atau Google Cloud Storage. Kami menyediakan integrasi referensi dengan checkpoint sinkron ke MaxText di checkpointing.py.

Konfigurasi yang didukung

Bentuk

Semua irisan harus memiliki bentuk yang sama (misalnya, AcceleratorType yang sama). Bentuk irisan heterogen tidak didukung.

Orkestrasi

Orkestrasi didukung dengan GKE. Untuk mengetahui informasi selengkapnya, lihat TPU di GKE.

Framework

Multislice hanya mendukung workload JAX dan PyTorch.

Keparalelan

Sebaiknya pengguna menguji Multislice dengan paralelisme data. Untuk mempelajari lebih lanjut cara menerapkan paralelisme pipeline dengan Multislice, hubungi perwakilan akun Google Cloud Anda.

Dukungan dan Masukan

Kami menerima semua masukan! Untuk memberikan masukan atau meminta dukungan, hubungi kami menggunakan formulir Dukungan atau Masukan Cloud TPU.