Google menggunakan teknologi AI untuk menerjemahkan konten ke dalam bahasa pilihan Anda. Terjemahan AI mungkin mengandung kesalahan.

Pengoptimalan performa TPU7x (Ironwood)

Panduan ini menjelaskan beberapa metode untuk mengoptimalkan performa dengan TPU7x (Ironwood) dengan mengelola pergerakan data secara efisien di antara sistem memori bertingkatnya. Hal ini mencakup teknik seperti pelatihan presisi rendah, sharding, pengoptimalan komunikasi, rematerialisasi aktivasi, penyesuaian memori vektor (VMEM) tercakup, dan kernel akselerator kustom.

Untuk mengoptimalkan performa dengan TPU7x, Anda harus memahami arsitektur Ironwood terlebih dahulu, khususnya hierarki memori dan topologi interkoneksi. Untuk mengetahui informasi selengkapnya, lihat TPU7x (Ironwood).

Pelatihan presisi rendah dengan FP8

FP8 (floating point 8-bit) adalah format data numerik efisien yang digunakan terutama untuk mempercepat pelatihan dan inferensi model. Dengan merepresentasikan angka menggunakan 8 bit – bukan format 16-bit standar (FP16 atau BF16) dan 32-bit (FP32) – TPU dapat memproses data secara signifikan lebih cepat dan menggunakan lebih sedikit memori.

TPU7x mendukung akselerasi hardware bawaan untuk jenis data FP8, yang menawarkan performa teoritis puncak sebesar 4.614 TFLOPS per chip. Kemampuan ini dapat menghasilkan waktu pelatihan end-to-end yang jauh lebih cepat. Untuk operasi yang kompatibel, terutama perkalian matriks padat yang umum untuk workload AI, penggunaan FP8 dapat menghasilkan peningkatan performa 1,3x lipat dibandingkan pelatihan BF16 standar. Dibandingkan dengan BF16, FP8 menggandakan FLOP puncak dan mengurangi jejak memori hingga setengahnya untuk bobot dan aktivasi. FP8 harus menjadi tuas penyesuaian utama untuk workload yang terikat komputasi dan skenario yang dibatasi oleh kapasitas memori atau bandwidth.

Penggunaan FP8 menawarkan manfaat performa berikut:

Mengurangi tekanan memori bandwidth tinggi (HBM): Jejak memori yang lebih kecil memungkinkan model yang lebih besar, atau model dengan cache KV yang lebih besar selama inferensi, untuk muat sepenuhnya dalam HBM 192 GiB. Hal ini menghindari offload yang mahal ke memori host yang lebih lambat.
Peningkatan ukuran batch efektif: Dengan mengurangi memori yang diperlukan untuk aktivasi, FP8 memungkinkan penggunaan ukuran batch yang lebih besar. Hal ini meningkatkan paralelisme data dan dapat menghasilkan throughput yang lebih tinggi serta pemanfaatan unit komputasi yang lebih baik.
Persyaratan bandwidth memori yang lebih rendah: Memindahkan setengah jumlah data untuk setiap operasi akan mengurangi permintaan pada jalur data HBM-ke-MXU. Pada sistem yang sering mengalami bottleneck dalam pergerakan data, hal ini membantu menjaga MXU tetap sibuk dengan pekerjaan.

Penggunaan FP8 dengan penurunan performa nol atau terbatas memerlukan pemilihan teknik kuantisasi yang cermat. Berikut beberapa praktik terbaik yang perlu dipertimbangkan untuk pelatihan FP8:

Granularitas penskalaan: Mulailah dengan penskalaan per tensor sebagai dasar pengukuran. Jika ada masalah kualitas atau performa, beralihlah ke penskalaan per sumbu. Penskalaan subsaluran mungkin tidak diperlukan.
Mode penskalaan: Penskalaan dinamis, yang menghitung faktor penskalaan saat runtime, adalah setelan default yang baik untuk mempertahankan kualitas. Meskipun penskalaan statis dapat menawarkan peningkatan performa yang signifikan dengan menghilangkan komputasi, penskalaan ini memerlukan pembuatan profil yang cermat untuk menentukan faktor penskalaan yang benar dan mungkin tidak cocok untuk semua kasus penggunaan, terutama saat konfigurasi model berubah. Sebaliknya, beberapa model dan konfigurasi yang tangguh dapat memperbaiki skala ke batas FP8 untuk bobot atau aktivasi, sehingga Anda dapat mengurangi overhead kuantisasi sekaligus mempertahankan akurasi dan meningkatkan performa.
Format FP8 (E4M3 dan E5M2): Pendekatan yang umum dan efektif adalah menggunakan campuran format FP8. Misalnya, gunakan E4M3 untuk bobot dan aktivasi dalam forward pass untuk memanfaatkan presisi E4M3 yang lebih tinggi, dan gunakan E5M2 untuk gradien dalam backward pass guna mengakomodasi rentang dinamis gradien yang lebih luas.
Pembulatan: Menggunakan "pembulatan ke bilangan genap terdekat" (RNE), bukan pembulatan stokastik untuk gradien, dapat mempertahankan kualitas sekaligus menawarkan performa dan reproduksibilitas yang lebih baik.
Mengaktifkan FP8 di MaxText: MaxText mendukung pelatihan FP8 melalui library kuantisasi QWIX. Untuk mengaktifkan kuantisasi, tetapkan flag berikut dalam konfigurasi Anda: use_qwix_quantization=true.

Sharding dan paralelisme

Sharding adalah proses membagi model besar atau data pelatihannya menjadi bagian-bagian yang lebih kecil dan mendistribusikannya ke beberapa chip atau core TPU. Memilih strategi sharding yang tepat penting untuk mencapai performa tinggi di TPU7x.

Pendekatan naif yang murni memaksimalkan tingkat paralelisme sering kali akan menghasilkan performa yang buruk karena terikat komunikasi. Pendekatan terbaik sering kali adalah memilih strategi sharding paling sederhana yang memenuhi batasan memori, karena hal ini meminimalkan overhead komunikasi dan memungkinkan unit komputasi digunakan secara efisien.

Sebelum memilih strategi sharding, langkah pertama dalam upaya penyesuaian performa harus berupa analisis intensitas aritmetika. Analisis ini menentukan apakah komputasi tertentu dibatasi oleh komputasi, bandwidth memori, atau bandwidth interkoneksi. Intensitas aritmetika dihitung sebagai rasio operasi floating point terhadap byte data yang harus dipindahkan.

Intensitas aritmetika yang tinggi menunjukkan workload terikat komputasi. Intensitas aritmetika yang rendah menunjukkan workload yang terikat pada memori atau komunikasi, dengan performa yang dibatasi oleh kecepatan pemindahan data dari HBM atau di seluruh jaringan ICI. Analisis ini memberikan informasi tentang ukuran batch dan strategi sharding yang ideal. Misalnya, workload yang terikat komunikasi tidak akan mendapatkan manfaat dari strategi sharding yang menimbulkan lebih banyak komunikasi, seperti paralelisme tensor tingkat tinggi.

Framework keputusan strategi sharding

MaxText menawarkan berbagai strategi sharding. Pilihan yang optimal bergantung pada arsitektur model, panjang urutan, dan kebutuhan untuk menyeimbangkan beban komputasi dengan overhead komunikasi.

Fully Sharded Data Parallelism (FSDP): Ini adalah strategi default pilihan untuk paralelisme data. FSDP melakukan sharding pada bobot model, gradien, dan status pengoptimal di seluruh perangkat paralel data. Selama komputasi, setiap perangkat melakukan operasi All-Gather untuk mengambil bobot penuh yang diperlukan untuk microbatch lokalnya. FSDP sangat efektif selama ukuran batch per perangkat cukup besar untuk menyembunyikan latensi komunikasi All-Gather ini. Untuk model Mixture-of-Experts (MoE), penghitungan intensitas aritmetika harus memperhitungkan sparsity.
Tensor Parallelism (TP): TP melakukan sharding tensor individual di seluruh perangkat. Biasanya, tensor adalah matriks bobot dalam perceptron multilayer (MLP) dan blok atensi. Intensitas aritmetika hardware yang tinggi (11,5 ribu) membebankan persyaratan yang sangat tinggi pada dimensi model untuk membuat TP layak melalui ICI, dan mencoba menggunakan TP dapat mengakibatkan sistem menjadi terikat komunikasi.
Expert Parallelism (EP): Ini adalah strategi standar dan diperlukan untuk melatih model MoE. EP melakukan sharding pada lapisan "pakar" di seluruh kumpulan perangkat, dan kolektif komunikasi All-to-All digunakan untuk merutekan token ke perangkat pakar yang ditentukan. EP dapat efisien jika dimensi MLP model cukup besar untuk mendekati roofline.
Paralelisme Konteks (CP): CP adalah strategi khusus yang penting untuk melatih model dengan panjang urutan yang sangat panjang. Fungsi utamanya adalah mengelola konsumsi memori aktivasi, yang bertambah secara kuadratik dengan panjang urutan dan dapat melebihi kapasitas HBM. CP melakukan sharding pada dimensi urutan tensor aktivasi, sehingga memungkinkan penggunaan ukuran batch per perangkat pecahan. Karena CP memerlukan lebih banyak komunikasi daripada FSDP, aturan umumnya adalah menggunakan tingkat CP minimum yang diperlukan untuk memenuhi batasan memori dan memastikan shard sumbu batch tetap berupa bilangan bulat.

Tabel berikut memetakan jenis workload umum ke strategi sharding yang optimal:

Jenis workload	Sharding utama yang direkomendasikan	Sharding sekunder	Hambatan utama	Alasan
Model padat - urutan pendek	FSDP	T/A	Rematerialisasi, Matmul FF	FSDP memberikan keseimbangan terbaik. Dengan urutan pendek, memori aktivasi mungkin bukan masalah utama. Kuncinya adalah batch global yang cukup besar untuk menyembunyikan All-Gather bobot FSDP. Seiring bertambahnya ukuran batch, ukuran aktivasi juga meningkat, dan kebijakan rematerialisasi yang sesuai diperlukan untuk memastikan konfigurasi ini tidak kehabisan memori.
Model padat - urutan panjang	FSDP	CP	Flash attention, memori aktivasi	Memori aktivasi menjadi batasan utama. CP diperlukan untuk mengaktifkan ukuran batch pecahan per perangkat dan menghindari masalah kehabisan memori (OOM). Flash attention adalah sumber utama komputasi dan waktu yang terbuang.
Model MoE - urutan pendek	FSDP + EP	T/A	All-to-All (Pemilihan rute pakar), rematerialisasi	Model MoE memerlukan EP untuk melakukan sharding pada pakar. Komunikasi All-to-All untuk perutean token adalah bottleneck utama yang harus ditumpuk. Rematerialisasi juga merupakan sumber pemborosan yang signifikan.
Model MoE - skala sangat besar	FSDP + EP + PP	Paralelisme model (MP)	Semua bottleneck yang disebutkan sebelumnya, ditambah bubble pipeline	Untuk model yang melebihi memori satu pod, PP diperlukan untuk melakukan sharding lapisan di seluruh pod. Hal ini menimbulkan overhead komunikasi DCN dan bubble pipeline. Ini adalah konfigurasi yang sangat kompleks yang memerlukan penyesuaian yang cermat.

Pengoptimalan komunikasi

Mekanisme utama untuk komunikasi dan komputasi yang tumpang-tindih di TPU7x disebut SparseCore Collective Offloading. Arsitektur Ironwood mencakup unit SparseCore khusus, yang bertindak sebagai thread kontrol independen yang mampu mengelola pergerakan data melalui fabric ICI. Hal ini memungkinkan operasi komunikasi kolektif (seperti All-Gather atau Reduce-Scatter) dijalankan secara paralel dengan komputasi utama yang terjadi di TensorCore. Ini adalah metode yang direkomendasikan untuk kolektif asinkron di TPU7x. Gunakan flag yang direkomendasikan untuk mengaktifkan offload bagi kolektif yang paling umum.

Rematerialisasi aktivasi

Rematerialisasi aktivasi, yang juga dikenal sebagai checkpointing gradien, adalah teknik dasar untuk mengurangi jejak HBM model. Alih-alih menyimpan semua aktivasi perantara dari forward pass di HBM untuk digunakan selama backward pass, teknik ini hanya menyimpan beberapa aktivasi utama (checkpoint) dan menghitung ulang aktivasi lainnya sesuai permintaan selama backward pass. Hal ini menghemat banyak memori dengan mengorbankan peningkatan komputasi (sekitar 25-30% FLOP tambahan untuk blok transformer standar).

Keputusan seberapa agresif menerapkan rematerialisasi adalah parameter penyesuaian penting yang sepenuhnya bergantung pada bottleneck utama, yang sering kali bervariasi dengan panjang urutan.

Untuk workload urutan panjang (seperti 128k): Dalam kasus ini, ukuran tensor aktivasi adalah konsumen dominan HBM. Workload biasanya terikat memori. Oleh karena itu, menerapkan kebijakan rematerialisasi yang agresif sangatlah bermanfaat. Penghematan memori memungkinkan pelatihan dilanjutkan tanpa error kehabisan memori dan juga memungkinkan ukuran batch yang lebih besar, dan overhead komputasi dari penghitungan ulang adalah kompromi yang sepadan.

Untuk workload urutan pendek (seperti 8k): Dalam kasus ini, memori aktivasi jauh lebih tidak menjadi masalah, dan workload lebih cenderung terikat komputasi. Overhead komputasi rematerialisasi dapat menjadi satu-satunya sumber ketidakefisienan terbesar.

Menyesuaikan kebijakan rematerialisasi di MaxText

MaxText memberikan kontrol terperinci atas rematerialisasi melalui serangkaian kebijakan preset dan kustom, yang dikonfigurasi menggunakan flag remat_policy.

Kebijakan preset

MaxText menawarkan kebijakan bawaan berikut:

full: Kebijakan paling agresif, yang merematerialisasi hampir semua hal. Hal ini meminimalkan penggunaan HBM, tetapi memaksimalkan overhead komputasi ulang. Ideal untuk skenario urutan panjang dengan batasan memori yang sangat ketat.
minimal: Kebijakan yang paling tidak agresif, yang menyimpan sebagian besar aktivasi. Hal ini memaksimalkan penggunaan HBM, tetapi meminimalkan komputasi ulang. Paling cocok untuk workload terikat komputasi berurutan pendek, yang tidak memerlukan banyak memori.
Kebijakan menengah: Opsi seperti save_dot_with_context_except_mlp, save_qkv_proj, dan save_out_proj memberikan berbagai kompromi dengan secara selektif melakukan checkpoint pada output operasi dot-product yang mahal sambil merematerialisasi operasi berbasis elemen yang lebih murah.

Kebijakan kustom

Untuk tingkat kontrol yang lebih tinggi, Anda dapat menetapkan remat_policy ke custom. Dengan demikian, Anda dapat menentukan perilaku untuk setiap lapisan dalam modul dekode model. Setiap lapisan dapat diberi salah satu dari tiga perilaku:

device: Aktivasi disimpan di HBM pada perangkat TPU.
remat: Aktivasi akan dihapus dan akan dimaterialisasi ulang selama backward pass.
offload: Aktivasi dipindahkan dari HBM ke memori host CPU, sehingga mengosongkan HBM dengan mengorbankan latensi transfer PCIe.

Penyesuaian VMEM tercakup

Performa kernel, seperti flash attention, bergantung pada ukuran tile yang dipilih di kernel, yang ukurannya dibatasi oleh memori vektor (VMEM) yang tersedia. Setiap dari dua TensorCore dalam chip TPU7x memiliki memori vektor (VMEM) sebesar 64 MiB. Kapasitas VMEM ini dapat dibagi antara cakupan saat ini (VMEM tercakup) dan pengambilan data bobot mendatang. Meningkatkan VMEM tercakup memungkinkan peningkatan ukuran ubin di kernel, yang berpotensi mengurangi kemacetan memori dan meningkatkan performa kernel. Anda dapat mengubah ukuran VMEM tercakup dengan menetapkan xla_tpu_scoped_vmem_limit_kib (dalam LIBTPU_INIT_ARGS), yang dapat digunakan untuk mempelajari performa kernel serta batas performa end-to-end. Mengoptimalkan ukuran VMEM tercakup dapat secara tidak langsung memengaruhi performa kernel Pallas kustom karena peningkatan VMEM tercakup akan membuka ruang penelusuran hyperparameter yang lebih besar untuk ukuran tile dalam kernel.

Kernel Tokamax

Tokamax, library kernel JAX berperforma tinggi dengan banyak kernel TPU yang sangat dioptimalkan, mengatasi beberapa bottleneck umum yang spesifik untuk hardware:

Splash attention: Splash attention digunakan sebagai implementasi attention utama untuk menghilangkan bottleneck HBM dari attention standar dan menggunakan implementasi attention yang paling efisien di TPU.
Megablox Grouped Matrix Multiplication (GMM): Untuk workload MoE, Megablox menangani perkalian matriks yang dikelompokkan secara efisien dengan menghitung representasi aktivasi yang tidak beraturan. Lapisan ini secara efisien memetakan dimensi ragged, menghitung perkalian matriks antara grup baris ragged di LHS, dan matriks pakar yang sesuai, sehingga tidak perlu melakukan padding batch ke ukuran tetap.
Penyesuaian empiris dengan tune-jax: Library tune-jax memiliki utilitas untuk melakukan penelusuran empiris untuk ukuran blok yang optimal. Ukuran kernel default sering kali kurang optimal; penyesuaian memungkinkan pemilihan ukuran ubin VMEM yang sesuai dengan hardware untuk memaksimalkan pemanfaatan hardware.
Estimasi logit maksimum: Kernel Tokamax Splash attention dapat dioptimalkan lebih lanjut dengan menetapkan nilai untuk max_logit_const. Jika ditetapkan, nilai ini akan menggantikan penghitungan pengurangan logit maksimum selama operasi softmax atensi (softmax(Q * KT)), sehingga mengurangi beberapa overhead komputasi dan sinkronisasi. Di MaxText, hal ini diimplementasikan oleh konfigurasi use_max_logits_estimate, yang dapat ditetapkan ke None (dinonaktifkan) atau nilai floating point. Pastikan rentang logit model spesifik Anda tetap kompatibel dengan estimasi untuk mencegah overflow numerik. Pengujian konvergensi direkomendasikan jika nilai ini ditetapkan.