Ringkasan Cloud TPU Multislice

Penting: Multislice hanya didukung di Cloud TPU v4 dan versi TPU yang lebih baru.

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, hal ini berarti tugas pelatihan dapat menggunakan lebih dari 4.096 chip dalam satu operasi. Untuk tugas pelatihan yang memerlukan kurang dari 4096 chip, satu slice dapat menawarkan performa terbaik. Namun, beberapa slice yang lebih kecil lebih mudah tersedia, sehingga memungkinkan waktu startup yang lebih cepat saat Multislice digunakan dengan slice yang lebih kecil.

Saat di-deploy dalam konfigurasi Multislice, chip TPU di setiap slice berkomunikasi melalui inter-chip-interconnect (ICI). Chip TPU di slice yang berbeda 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 komunikasi dengan komputasi untuk performa maksimum.

Konsep

Jenis akselerator

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

Perbaikan otomatis

Saat slice mengalami peristiwa pemeliharaan, preemptif, atau kegagalan hardware, Cloud TPU akan membuat slice baru. Dalam kasus yang jarang terjadi, jika tidak ada resource yang memadai untuk membuat slice baru, pembuatan tidak akan selesai hingga hardware tersedia. Setelah slice baru dibuat, semua slice lainnya di lingkungan Multislice akan dimulai ulang sehingga pelatihan dapat dilanjutkan.Dengan skrip startup yang dikonfigurasi dengan benar, skrip pelatihan dapat diluncurkan ulang secara otomatis tanpa intervensi pengguna, memuat, dan melanjutkan dari titik pemeriksaan terbaru.

Dataset

Data yang digunakan oleh model untuk pelatihan atau inferensi.

Jaringan Pusat Data (DCN)

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

Penjadwalan gang

Saat semua slice TPU disediakan secara bersamaan, pada saat yang sama, menjamin semua atau tidak ada slice yang 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 latensi rendah yang menghubungkan TPU dalam Pod TPU.

Multislice

Dua atau beberapa slice chip TPU yang dapat berkomunikasi melalui DCN.

Node

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

Pod

Kumpulan chip TPU yang terhubung oleh antarmuka jaringan ICI khusus. Pod memungkinkan Anda mendistribusikan beban pemrosesan ke beberapa TPU.

Resource dalam antrean (QR)

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

Skrip startup

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

Slice TPU

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

TPU VM

Virtual machine yang menjalankan Linux dan 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 oleh Google. Chip ini dirancang untuk menawarkan komputasi yang cepat dan hemat daya untuk tugas machine learning utama seperti perkalian matriks.

Jenis kapasitas Cloud TPU

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

• Reservasi: Untuk menggunakan reservasi, Anda harus memiliki perjanjian reservasi dengan Google. Gunakan flag --reserved saat membuat resource.
• Spot: Menargetkan kuota preemptible menggunakan Spot VM. Resource Anda mungkin didahulukan untuk memberi ruang bagi permintaan tugas dengan prioritas yang lebih tinggi. Gunakan tanda --spot saat membuat resource.
• On-demand: Menargetkan kuota on-demand, yang tidak memerlukan pemesanan dan tidak akan didahului. Permintaan TPU akan dimasukkan ke antrean kuota on-demand yang ditawarkan oleh Cloud TPU, ketersediaan resource tidak dijamin. Dipilih secara default, tidak perlu tanda.

Catatan: Cloud TPU membatasi ukuran slice untuk semua jenis akselerator. Untuk batas ukuran slice untuk 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. Untuk menggunakan Multislice, resource TPU Anda harus dikelola sebagai Resource dalam Antrean.

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

Contoh pengantar

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

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

Paralelisme ICI

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

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

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

Input ini harus dibatasi sehingga ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism sama dengan jumlah chip dalam slice.

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
Paralelisme FSDP 2 arah + Tensor 2 arah	1	2	2

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

Contoh ini mengasumsikan bahwa Anda sudah terbiasa menjalankan kode di satu slice TPU seperti dalam Menjalankan penghitungan di VM Cloud TPU menggunakan JAX. Contoh ini menunjukkan cara menjalankan shardings.py pada satu slice.

1. Siapkan lingkungan:

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

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

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

3. Sediakan TPU:

   gcloud

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

   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-id

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

   reserved

   Gunakan reservasi saat membuat slice.

   spot

   Gunakan Spot VM saat membuat slice.

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

   Konsol

   1. Di konsol Google Cloud, buka halaman TPUs:

      Buka TPU

   2. Klik Buat TPU.

   3. Di kolom Nama, masukkan nama untuk TPU Anda.

   4. Di kotak Zone, pilih zona tempat Anda ingin membuat TPU.

   5. Di kotak TPU type, pilih jenis akselerator. 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.

   6. Di kotak Versi software TPU, pilih versi software. Saat membuat VM Cloud TPU, versi software TPU menentukan versi runtime TPU yang akan diinstal. Untuk mengetahui informasi selengkapnya, lihat Gambar VM TPU.

   7. Klik tombol Aktifkan antrean.

   8. Di kolom Queued resource name, masukkan nama untuk permintaan resource yang diantrekan.

   9. Klik Buat untuk membuat permintaan resource yang diantrekan.

4. Tunggu hingga resource yang diantrean berada dalam status ACTIVE yang berarti node pekerja berada dalam status READY. Setelah penyediaan resource dalam antrean dimulai, mungkin perlu waktu satu hingga lima menit untuk menyelesaikannya, bergantung pada ukuran resource dalam antrean. Anda dapat memeriksa status permintaan resource yang diantrekan menggunakan gcloud CLI atau Google Cloud Console:

   gcloud

   $ gcloud compute tpus queued-resources \
       list --filter=YOUR_QR_ID

   Konsol

   1. Di konsol Google Cloud, buka halaman TPUs:

      Buka TPU

   2. Klik tab Resource dalam antrean.

   3. Klik nama permintaan resource yang diantrekan.

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 menyertakan shardings.py) ke VM TPU Anda.

7. Dalam direktori repositori MaxText, jalankan skrip penyiapan untuk menginstal JAX dan dependensi lainnya di slice TPU Anda. Skrip penyiapan memerlukan 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 penggunaan FLOP yang mengesankan sebesar 90%+. Dalam hal ini, kita telah memilih kira-kira batch maksimum yang sesuai dengan Memori Bandwidth Tinggi (HBM) TPU.

9. Jangan ragu untuk mempelajari strategi sharding lainnya di ICI, misalnya, Anda dapat mencoba kombinasi berikut:

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

10. Hapus resource yang diantrekan dan slice TPU setelah selesai. Anda harus menjalankan langkah-langkah pembersihan ini dari lingkungan tempat Anda menyiapkan slice (jalankan exit terlebih dahulu untuk keluar dari sesi SSH). Penghapusan akan memerlukan waktu dua hingga lima menit untuk diselesaikan. Jika menggunakan gcloud CLI, Anda dapat menjalankan perintah ini di latar belakang dengan flag --async opsional.

    Penting: Resource dalam antrean harus dihapus setelah digunakan. Jika tidak, resource tersebut akan terus menggunakan kuota resource dalam antrean yang dilacak secara terpisah dari kuota TPU.

    gcloud

    $ gcloud compute tpus queued-resources
        delete YOUR_QR_ID --force (--async)

    Konsol

    1. Di konsol Google Cloud, buka halaman TPUs:

       Buka TPU

    2. Klik tab Resource dalam antrean.

    3. Centang kotak di samping permintaan resource yang ada dalam antrean.

    4. Klik Delete.

Sharding multislice menggunakan paralelisme DCN

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

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

Nilai parameter ini harus dibatasi sehingga dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism sama dengan jumlah slice.

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

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

dcn_tensor_parallelism harus selalu disetel 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 hal ini bergantung pada aplikasi.

Saat Anda meningkatkan jumlah slice (dengan asumsi Anda mempertahankan ukuran slice dan batch per slice tetap 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 langkah-langkah Memulai: Eksperimen Cepat di Beberapa slice dari repositori MaxText, kecuali di sini kita 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 berumur panjang, sebaiknya jangan gunakan skrip ini untuk tugas yang berjalan lama. Jika Anda ingin menjalankan tugas yang lebih lama (misalnya, berjam-jam atau berhari-hari), sebaiknya gunakan multihost_job.py.

Dalam tutorial ini, kita menggunakan istilah runner untuk menunjukkan mesin tempat Anda menjalankan skrip multihost_runner.py. Kami menggunakan istilah pekerja untuk menunjukkan VM TPU yang membentuk slice 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 di atas 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 di 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 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 perintah queued-resources create atau konsol Google Cloud.

   gcloud

   Catatan: Pastikan Anda memiliki kuota masing-masing sebelum memilih --reserved, atau --spot, 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 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|--spot]

   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. Multislice hanya didukung di Cloud TPU v4 dan versi TPU yang lebih baru. 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 slice yang akan dibuat.

   node-prefix

   Awalannya digunakan untuk membuat nama untuk setiap slice. Angka ditambahkan ke awalan untuk setiap slice. Misalnya, jika Anda menetapkan node-prefix ke mySlice, slice akan diberi nama: mySlice-0, mySlice-1, yang berlanjut secara numerik untuk setiap slice.

   reserved

   Gunakan reservasi saat membuat slice.

   spot

   Gunakan Spot VM saat membuat slice.

   Konsol

   1. Di konsol Google Cloud, buka halaman TPUs:

      Buka TPU

   2. Klik Buat TPU.

   3. Di kolom Nama, masukkan nama untuk TPU Anda.

   4. Di kotak Zone, pilih zona tempat Anda ingin membuat TPU.

   5. Di kotak TPU type, pilih jenis akselerator. Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Multislice hanya didukung di Cloud TPU v4 dan versi TPU yang lebih baru. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.

   6. Di kotak Versi software TPU, pilih versi software. Saat membuat VM Cloud TPU, versi software TPU menentukan versi runtime TPU yang akan diinstal. Untuk mengetahui informasi selengkapnya, lihat Gambar VM TPU.

   7. Klik tombol Aktifkan antrean.

   8. Di kolom Queued resource name, masukkan nama untuk permintaan resource yang diantrekan.

   9. Klik kotak centang Make this a Multislice TPU.

   10. Di kolom Slice count, masukkan jumlah slice yang ingin Anda buat.

   11. Klik Buat untuk membuat permintaan resource yang diantrekan.

6. Saat penyediaan resource yang diantrean dimulai, mungkin perlu waktu hingga lima menit untuk selesai, bergantung pada ukuran resource yang diantrean. Tunggu hingga resource yang diantrekan berada dalam status ACTIVE. Anda dapat memeriksa status permintaan resource yang diantrekan menggunakan gcloud CLI atau Google Cloud Console:

   gcloud

   $ gcloud compute tpus queued-resources list \
       --filter=YOUR_QR_ID

   Tindakan ini akan menghasilkan output yang terlihat seperti:

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

   Konsol

   1. Di konsol Google Cloud, buka halaman TPUs:

      Buka TPU

   2. Klik tab Resource dalam antrean.

   3. Klik nama permintaan resource yang diantrekan.

   Hubungi perwakilan akun Google Cloud Anda jika status QR berada dalam 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. Jalankan shardings.py di 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 resource yang diantrekan setelah selesai. Penghapusan akan memerlukan waktu dua hingga lima menit untuk diselesaikan. Jika menggunakan gcloud CLI, Anda dapat menjalankan perintah ini di latar belakang dengan flag --async opsional.

Menskalakan workload ke Multislice

Sebelum menjalankan model 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.

Ini seharusnya adalah satu-satunya perubahan kode yang diperlukan saat beralih ke Multislice. Untuk mencapai performa tinggi, DCN perlu dipetakan ke sumbu paralel data, paralel data yang sepenuhnya di-sharding, atau paralel pipeline. Pertimbangan performa dan strategi sharding dibahas secara lebih mendetail dalam 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 slice v4-16, Anda memiliki delapan chip per slice * 4 slice, sehingga len(jax.devices()) akan menampilkan 32.

Memilih ukuran slice untuk lingkungan Multislice

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

Menjalankan Tugas di beberapa slice

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 log kembali, dan melacak status error setiap perintah. Skrip multihost_runner.py didokumentasikan dalam README MaxText.

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

Skrip runner produksi

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

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

Pendekatan manual

Sebaiknya gunakan atau sesuaikan multihost_runner.py atau multihost_job.py untuk menjalankan beban kerja kustom di 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 kode QR secara manual tanpa menggunakan alat yang disediakan di repo MaxText, baca bagian berikut.

Membuat resource dalam antrean

gcloud

1. 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

   Deskripsi variabel

   Input	Deskripsi
   YOUR_QR_ID	ID yang ditetapkan pengguna untuk resource yang diantrekan.
   PROJECT	Nama project Google Cloud
   ZONA	Menentukan zona tempat membuat resource.
   NETWORK_NAME	Nama jaringan VPC.
   SUBNETWORK_NAME	Nama subnet di jaringan VPC
   RUNTIME_VERSION	Versi software Cloud TPU.
   ACCELERATOR_TYPE	v4-16
   EXAMPLE_TAG_1, EXAMPLE_TAG_2 …	Tag yang digunakan untuk mengidentifikasi sumber atau target valid untuk firewall jaringan
   SLICE_COUNT	Jumlah slice. Dibatasi hingga maksimum 256 slice.
   STARTUP_SCRIPT	Jika Anda menentukan skrip startup, skrip akan berjalan saat slice TPU disediakan atau dimulai ulang.

2. Buat permintaan resource yang diantrekan menggunakan perintah berikut:

   $ gcloud 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|--spot]

   Deskripsi flag perintah

   YOUR_QR_ID

   String yang ditentukan pengguna yang mengidentifikasi permintaan resource yang diantrekan.

   project

   Project Google Cloud tempat Anda membuat permintaan resource yang diantrekan.

   zone

   Zona Google Cloud tempat membuat resource yang diantrekan.

   node-count

   Jumlah slice yang akan dibuat.

   accelerator-type

   Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Multislice hanya didukung di Cloud TPU v4 dan versi TPU yang lebih baru. 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 resource TPU akan dilampirkan.

   subnetwork

   Nama subnetwork VPC tempat resource TPU akan dilampirkan.

   reserved

   Gunakan reservasi saat membuat slice.

   spot

   Gunakan Spot VM saat membuat slice.

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

curl

1. 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"
       }
       }
       ]
   }
   }

   Ganti nilai berikut:

   • your-project-number - Nomor project Google Cloud Anda
   • your-zone - Zona tempat Anda ingin membuat resource yang diantrekan
   • accelerator-type - Versi dan ukuran satu slice. Multislice hanya didukung di Cloud TPU v4 dan versi TPU yang lebih baru.
   • tpu-vm-runtime-version - Versi runtime VM TPU yang ingin Anda gunakan.
   • your-network-name - Opsional, jaringan tempat resource yang diantrekan akan dilampirkan
   • your-subnetwork-name - Opsional, subnetwork tempat resource yang diantrekan akan dilampirkan
   • example-tag-1 - Opsional, string tag arbitrer
   • your-startup-script - Skrip startup yang akan dijalankan saat resource yang diantrekan dialokasikan
   • slice-count - Jumlah slice TPU di lingkungan Multislice Anda
   • YOUR_QR_ID - ID yang disediakan pengguna untuk resource yang diantrekan

   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 Spot, ganti:

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

   Hapus baris untuk menggunakan kapasitas on demand default.

2. Kirim permintaan pembuatan resource yang diantrekan 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

   Ganti nilai berikut:

   • your-project-id - Project ID Google Cloud Anda
   • your-zone - Zona tempat Anda ingin membuat resource yang diantrekan
   • YOUR_QR_ID - ID yang disediakan pengguna untuk resource yang diantrekan

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 resource yang diantrekan.

Konsol

1. Di konsol Google Cloud, buka halaman TPUs:

   Buka TPU

2. Klik Buat TPU.

3. Di kolom Nama, masukkan nama untuk TPU Anda.

4. Di kotak Zone, pilih zona tempat Anda ingin membuat TPU.

5. Di kotak TPU type, pilih jenis akselerator. Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Multislice hanya didukung di Cloud TPU v4 dan versi TPU yang lebih baru. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.

6. Di kotak Versi software TPU, pilih versi software. Saat membuat VM Cloud TPU, versi software TPU menentukan versi runtime TPU yang akan diinstal. Untuk mengetahui informasi selengkapnya, lihat Gambar VM TPU.

7. Klik tombol Aktifkan antrean.

8. Di kolom Queued resource name, masukkan nama untuk permintaan resource yang diantrekan.

9. Klik kotak centang Make this a Multislice TPU.

10. Di kolom Slice count, masukkan jumlah slice yang ingin Anda buat.

11. Klik Buat untuk membuat permintaan resource yang diantrekan.

Mengambil status resource dalam antrean

gcloud

$ gcloud compute tpus queued-resources describe ${YOUR_QR_ID}

Untuk resource yang diantrekan dalam status ACTIVE, output-nya akan terlihat seperti berikut:

...
state:
    state: ACTIVE
...

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}

Untuk resource yang diantrekan dalam status ACTIVE, output-nya akan terlihat seperti berikut:

{
"name": your-queued-res,
"tpu": {
    "nodeSpec": [
    {
        ... // node 1
    },
    {
        ... // node 2
    },
    ...
    ]
},
...
"state": "ACTIVE"
}

Konsol

1. Di konsol Google Cloud, buka halaman TPUs:

   Buka TPU

2. Klik tab Resource dalam antrean.

3. Klik nama permintaan resource yang diantrekan.

Setelah TPU disediakan, Anda juga dapat melihat detail tentang permintaan resource yang diantrekan dengan membuka halaman TPU, menemukan TPU, dan mengklik nama permintaan resource yang diantrekan yang sesuai.

Dalam skenario yang jarang terjadi, Anda mungkin menemukan resource yang diantrekan dalam status FAILED, sementara beberapa slice adalah ACTIVE. Jika hal ini terjadi, hapus resource yang dibuat, dan coba lagi dalam beberapa menit atau hubungi Dukungan Google Cloud.

SSH dan menginstal dependensi

Menjalankan kode JAX pada slice Pod TPU menjelaskan cara terhubung ke VM TPU menggunakan SSH dalam satu slice. 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 mengirim perintah yang ditentukan ke semua pekerja dan node di QR menggunakan SSH. Perintah dikelompokkan menjadi grup yang terdiri dari empat perintah dan dikirim secara bersamaan. Batch perintah berikutnya akan dikirim saat batch saat ini menyelesaikan eksekusi. Jika terjadi kegagalan pada salah satu perintah, pemrosesan akan berhenti, 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. Misalnya, anggap batas pengelompokan di komputer lokal Anda adalah 64. Jika Anda mencoba menjalankan skrip pelatihan pada lebih dari 64 slice, misalnya 100, perintah SSH akan membagi slice menjadi beberapa batch. Tindakan ini akan menjalankan skrip pelatihan pada batch pertama yang terdiri dari 64 slice dan menunggu skrip selesai sebelum menjalankan skrip pada batch yang tersisa yang terdiri dari 36 slice. Namun, batch pertama yang terdiri dari 64 slice tidak dapat selesai hingga 36 slice yang tersisa mulai menjalankan skrip, sehingga menyebabkan deadlock.

Untuk mencegah skenario ini, Anda dapat menjalankan skrip pelatihan di latar belakang di 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. Perintah SSH kemudian dapat mulai menjalankan skrip pelatihan pada batch yang tersisa dari 36 slice. Anda harus menyalurkan streaming stdout dan stderr dengan benar saat menjalankan perintah di latar belakang. Untuk meningkatkan paralelisme dalam QR yang sama, Anda dapat memilih slice tertentu menggunakan parameter --node.

Penyiapan jaringan

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

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

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

Mencantumkan resource dalam antrean

gcloud

Anda dapat melihat status resource yang diantrekan menggunakan perintah queued-resources list:

$ gcloud compute tpus queued-resources list

Outputnya terlihat mirip dengan yang berikut ini:

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

Konsol

1. Di konsol Google Cloud, buka halaman TPUs:

   Buka TPU

2. Klik tab Resource dalam antrean.

Memulai tugas di lingkungan yang disediakan

Anda dapat menjalankan beban kerja secara manual dengan menghubungkan ke semua host di setiap slice melalui SSH dan menjalankan perintah berikut di 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 total 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 macet dan mereset semua VM lebih mudah daripada men-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 resource dalam antrean

Untuk merilis resource di akhir sesi pelatihan, hapus resource yang diantrekan. Penghapusan akan memerlukan waktu dua hingga lima menit. Jika menggunakan gcloud CLI, Anda dapat menjalankan perintah ini di latar belakang dengan flag --async opsional.

gcloud

$ gcloud compute tpus queued-resources \
    delete YOUR_QR_ID --force (--async)

Konsol

1. Di konsol Google Cloud, buka halaman TPUs:

   Buka TPU

2. Klik tab Resource dalam antrean.

3. Centang kotak di samping permintaan resource yang ada dalam antrean.

4. Klik Delete.

Pemulihan kegagalan otomatis

Jika terjadi gangguan, Multislice menawarkan perbaikan tanpa intervensi pada slice yang terpengaruh dan mereset semua slice setelahnya. Slice yang terpengaruh akan diganti dengan slice baru dan slice yang tersisa yang sehat akan direset. Jika tidak ada kapasitas yang tersedia untuk mengalokasikan slice pengganti, pelatihan akan berhenti.

Untuk melanjutkan pelatihan secara otomatis setelah gangguan, Anda harus menentukan skrip startup yang memeriksa dan memuat titik pemeriksaan terakhir yang disimpan. Skrip startup Anda akan otomatis dijalankan setiap kali slice dialokasikan ulang atau VM direset. Anda menentukan skrip startup dalam payload JSON yang dikirim ke API permintaan buat QR.

Skrip startup berikut (digunakan di Membuat QR) memungkinkan Anda otomatis pulih dari kegagalan dan melanjutkan pelatihan dari titik pemeriksaan 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 satu slice 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 shard di beberapa slice. Biasanya ada tiga pilihan (paralelisme data, paralelisme data yang sepenuhnya di-shard, dan paralelisme pipeline). Sebaiknya jangan melakukan sharding aktivasi 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 telah berhasil digunakan sebelumnya.

Sebaiknya mulai dengan paralelisme data murni. Menggunakan paralelisme data yang sepenuhnya di-shard berguna untuk mengosongkan penggunaan memori. Kelemahannya adalah komunikasi antar-slice menggunakan jaringan DCN dan akan memperlambat workload 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 slice.

Untuk mencapai penskalaan yang kuat di beberapa slice, jumlah waktu yang diperlukan untuk melakukan all-reduce melalui DCN harus kurang dari jumlah waktu yang diperlukan untuk melakukan pass mundur. DCN digunakan untuk komunikasi antar-slice dan merupakan faktor pembatas dalam throughput beban kerja.

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

Ada empat chip per host TPU dan setiap host memiliki bandwidth jaringan maksimum sebesar 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 slice, model Anda akan menggunakan bandwidth DCN 32 bit. Jika Anda menggunakan lebih dari dua slice, compiler akan melakukan operasi shuffle all-reduce lengkap 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 adalah sekitar 6 * B * P dengan:

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

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

Untuk memastikan penskalaan yang kuat di beberapa slice, pastikan intensitas operasional melebihi intensitas aritmetika hardware TPU. Untuk menghitung intensitas operasional, bagi jumlah FLOPS per parameter selama hitungan 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 slice: Operational intensity = 4 * B / 32

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

Hal ini menunjukkan ukuran batch minimum antara 176 ribu dan 352 ribu untuk model bahasa berbasis Transformer. Karena jaringan DCN dapat menjatuhkan paket untuk sementara, sebaiknya pertahankan margin error yang signifikan, dengan 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 pass mundur per slice (baik dengan mengatur waktu menggunakan profiler atau dengan menghitung FLOPS). Kemudian, Anda dapat membandingkannya dengan waktu proses yang diharapkan untuk mengurangi semua DCN dan mendapatkan estimasi yang baik tentang apakah paralelisme data akan sesuai untuk Anda.

Kapan harus menggunakan paralelisme data yang sepenuhnya di-shard (FSDP)

Paralelisme data yang di-shard sepenuhnya (FSDP) menggabungkan paralelisme data (sharding data di seluruh node) dengan sharding bobot di seluruh node. Untuk setiap operasi dalam langkah maju dan mundur, semua bobot dikumpulkan sehingga setiap slice memiliki bobot yang diperlukan. Alih-alih menyinkronkan gradien menggunakan all-reduce, gradien akan di-reduce-scatter saat dihasilkan. Dengan cara ini, setiap slice hanya mendapatkan gradien untuk bobot yang menjadi tanggung jawabnya.

Serupa dengan paralelisme data, FSDP akan memerlukan penskalaan ukuran batch global secara linear dengan jumlah slice. FSDP akan mengurangi tekanan memori saat Anda meningkatkan jumlah slice. Hal ini karena jumlah bobot dan status pengoptimal per slice berkurang, tetapi hal ini dilakukan dengan mengorbankan traffic jaringan yang meningkat dan kemungkinan pemblokiran yang lebih besar karena kolektif yang tertunda.

Dalam praktiknya, FSDP di seluruh slice paling baik jika Anda meningkatkan batch per slice, menyimpan lebih banyak aktivasi untuk meminimalkan rematerialisasi selama pass mundur atau meningkatkan jumlah parameter dalam jaringan saraf Anda.

Operasi all-gather dan all-reduce di FSDP berfungsi mirip dengan operasi di DP, sehingga Anda dapat menentukan apakah beban kerja FSDP 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 yang lebih besar dari ukuran batch maksimum yang Anda inginkan. Paralelisme pipeline memungkinkan slice yang terdiri dari pipeline untuk "membagikan" batch. Namun, paralelisme pipeline memiliki dua kelemahan signifikan:

1. Hal ini menyebabkan "gelembung pipeline" saat chip tidak ada aktivitas karena menunggu data.
2. Hal ini memerlukan pengelompokan mikro yang mengurangi ukuran batch 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 batch yang diperlukan untuk mencapai FSDP berperforma tinggi. FSDP cenderung mencapai penggunaan FLOP model yang lebih tinggi, tetapi jika konvergensi per sampel melambat seiring peningkatan ukuran batch, paralelisme pipeline mungkin masih menjadi pilihan yang lebih baik. Sebagian besar beban kerja dapat mentoleransi ukuran batch 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. Hal ini akan memungkinkan Anda meminimalkan kedalaman pipeline sekaligus meningkatkan ukuran batch per pipeline hingga latensi DCN menjadi kurang menjadi faktor dalam throughput. Secara konkret, jika Anda memiliki N slice, pertimbangkan pipeline kedalaman 2 dan replika N/2 paralelisme data, lalu pipeline kedalaman 4 dan replika N/4 paralelisme data, dan seterusnya, hingga batch per pipeline menjadi cukup besar sehingga kolektif DCN dapat disembunyikan di balik aritmetika dalam penerusan mundur. Hal ini akan meminimalkan pelambatan yang disebabkan oleh paralelisme pipeline sekaligus memungkinkan Anda menskalakan melebihi 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 mengelompokkan batch di seluruh host sangat penting untuk menghindari kekurangan tugas TPU. Loader data saat ini di MaxText membuat setiap host memuat subkumpulan contoh yang sama. Solusi ini cukup untuk teks, tetapi memerlukan reshard dalam model. Selain itu, MaxText belum menawarkan snapshot deterministik yang akan memungkinkan iterator data memuat data yang sama sebelum dan sesudah preemptif.

Checkpoint

Library checkpointing Orbax menyediakan primitif untuk melakukan checkpointing JAX PyTrees ke penyimpanan lokal atau penyimpanan Google Cloud. Kami menyediakan integrasi referensi dengan pemeriksaan titik henti sementara sinkron ke MaxText di checkpointing.py.

Konfigurasi yang didukung

Bentuk

Semua slice harus memiliki bentuk yang sama (misalnya, AcceleratorType yang sama). Bentuk slice 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.