TPU v4

Dokumen ini menjelaskan arsitektur dan konfigurasi yang didukung Cloud TPU v4.

Catatan: Semua permintaan kuota di us-central2-b (saat ini, hanya TPU v4) harus disetujui secara manual oleh Google. Tidak ada kuota default yang diberikan. Jika Anda memerlukan kuota di us-central2-b, lihat Melihat dan meminta kuota tambahan.

Arsitektur sistem

Setiap chip TPU v4 berisi dua TensorCore. Setiap TensorCore memiliki empat unit perkalian matriks (MXU), unit vektor, dan unit skalar. Tabel berikut menunjukkan spesifikasi utama untuk Pod TPU v4.

Spesifikasi utama	Nilai Pod v4
Komputasi puncak per chip	275 teraflop (bf16 atau int8)
Kapasitas dan bandwidth HBM2	32 GiB, 1200 GBps
Daya minimum/rata-rata/maksimum yang diukur	90/170/192 W
Ukuran Pod TPU	4096 chip
Topologi Interconnect	Mesh 3D
Komputasi puncak per Pod	1,1 exaflop (bf16 atau int8)
Bandwidth all-reduce per Pod	1,1 PB/dtk
Bandwidth biseksi per Pod	24 TB/dtk

Diagram berikut menggambarkan chip TPU v4.

Untuk mengetahui informasi selengkapnya tentang detail arsitektur dan karakteristik performa TPU v4, lihat TPU v4: Superkomputer yang Dapat Dikonfigurasi Ulang Secara Optik untuk Machine Learning dengan Dukungan Hardware untuk Embedding.

Mesh 3D dan torus 3D

TPU v4 memiliki koneksi langsung ke chip tetangga terdekat dalam 3 dimensi, sehingga menghasilkan mesh 3D dari koneksi jaringan. Koneksi dapat dikonfigurasi sebagai torus 3D pada slice dengan topologi, AxBxC, adalah 2A=B=C atau 2A=2B=C, dengan setiap dimensi adalah kelipatan 4. Misalnya, 4x4x8, 4x8x8, atau 12x12x24. Secara umum, performa konfigurasi torus 3D akan lebih baik daripada konfigurasi mesh 3D. Untuk mengetahui informasi selengkapnya, lihat Topologi tori terpilin.

Manfaat performa TPU v4 dibandingkan v3

Bagian ini menunjukkan cara yang hemat memori untuk menjalankan skrip pelatihan contoh di TPU v4, serta peningkatan performa untuk TPU v4 dibandingkan dengan TPU v3.

Sistem memori

Non Uniform Memory Access (NUMA) adalah arsitektur memori komputer untuk mesin yang memiliki beberapa CPU. Setiap CPU memiliki akses langsung ke blok memori berkecepatan tinggi. CPU dan memorinya disebut node NUMA. Node NUMA terhubung ke node NUMA lain yang berdekatan satu sama lain. CPU dari satu node NUMA dapat mengakses memori di node NUMA lain, tetapi akses ini lebih lambat daripada mengakses memori dalam node NUMA.

Software yang berjalan di komputer multi-CPU dapat menempatkan data yang dibutuhkan oleh CPU dalam node NUMA-nya, sehingga meningkatkan throughput memori. Untuk mengetahui informasi selengkapnya tentang NUMA, lihat Non Uniform Memory Access di Wikipedia.

Anda dapat memanfaatkan manfaat lokalitas NUMA dengan mengikat skrip pelatihan ke Node NUMA 0.

Untuk mengaktifkan pengikatan node NUMA:

1. Buat TPU v4.

2. Hubungkan ke TPU menggunakan SSH.

3. Instal alat command line numactl. numactl memungkinkan Anda menjalankan proses dengan kebijakan penempatan memori atau penjadwalan NUMA tertentu.

    $ sudo apt-get update
    $ sudo apt-get install numactl

4. Ikat kode skrip Anda ke NUMA Node 0. Ganti your-training-script dengan jalur ke skrip pelatihan Anda.

    $ numactl --cpunodebind=0 python3 your-training-script

Aktifkan pengikatan node NUMA jika:

• Jika beban kerja Anda sangat bergantung pada beban kerja CPU (misalnya, klasifikasi gambar, beban kerja rekomendasi) terlepas dari framework.
• Jika Anda menggunakan versi runtime TPU tanpa akhiran -pod (misalnya, tpu-vm-tf-2.10.0-v4).

Perbedaan sistem memori lainnya:

• Chip TPU v4 memiliki ruang memori HBM 32-GiB terpadu di seluruh chip, sehingga memungkinkan koordinasi yang lebih baik antara dua TensorCore di chip.
• Peningkatan performa HBM menggunakan standar dan kecepatan memori terbaru.
• Profil performa DMA yang ditingkatkan dengan dukungan bawaan untuk melangkah berperforma tinggi dengan perincian 512 B.

TensorCores

• Jumlah MXU dua kali lipat dan kecepatan clock yang lebih tinggi menghasilkan TFLOPS maksimum 275.
• Bandwidth transposisi dan permutasi 2x.
• Model akses memori load-store untuk Memori Umum (Cmem).
• Bandwidth pemuatan bobot MXU yang lebih cepat dan dukungan mode 8-bit untuk memungkinkan ukuran batch yang lebih kecil dan latensi inferensi yang lebih baik.

Interkoneksi antar-chip

Enam link interkoneksi per chip untuk mengaktifkan topologi jaringan yang memiliki diameter jaringan yang lebih kecil.

Lainnya

• Antarmuka PCIE gen3 x16 ke host (koneksi langsung).
• Model keamanan yang ditingkatkan.
• Peningkatan efisiensi energi.

Konfigurasi

Pod TPU v4 terdiri dari 4.096 chip yang saling terhubung dengan link berkecepatan tinggi yang dapat dikonfigurasi ulang. Jaringan fleksibel TPU v4 memungkinkan Anda menghubungkan chip dalam slice berukuran sama dengan berbagai cara. Saat membuat slice TPU, Anda menentukan versi TPU dan jumlah resource TPU yang Anda perlukan.

Anda menentukan topologi TPU menggunakan 3-tuple, AxBxC dengan A<=B<=C dan A, B, C semuanya <= 4 atau semuanya kelipatan bilangan bulat 4. Nilai A, B, dan C adalah jumlah chip di setiap tiga dimensi.

Topologi dengan 2A=B=C atau 2A=2B=C juga memiliki varian topologi yang dioptimalkan untuk komunikasi semua-ke-semua, misalnya, 4×4×8, 8×8×16, dan 12×12×24. Topologi ini dikenal sebagai topologi torus terpelintir.

Ilustrasi berikut menunjukkan beberapa topologi TPU v4 umum.

Slice yang lebih besar dapat dibuat dari satu atau beberapa "kubus" chip 4x4x4.

Untuk mengetahui informasi selengkapnya tentang cara mengelola TPU, lihat Mengelola TPU. Untuk mengetahui informasi selengkapnya tentang arsitektur sistem Cloud TPU, lihat Arsitektur sistem.

Topologi tori yang berbelit

Beberapa bentuk slice torus 3D v4 memiliki opsi untuk menggunakan topologi yang dikenal sebagai torus terpilin. Misalnya, dua kubus v4 dapat disusun sebagai slice 4x4x8 atau 4x4x8_twisted. Topologi twisted menawarkan bandwidth biseksi yang jauh lebih tinggi. Misalnya, slice dengan topologi 4x4x8_twisted memberikan peningkatan bandwidth biseksi teoretis sebesar 70% dibandingkan slice 4x4x8 yang tidak di-twist. Bandwidth biseksi yang ditingkatkan berguna untuk workload yang menggunakan pola komunikasi global. Topologi twisted dapat meningkatkan performa untuk sebagian besar model, dengan beban kerja penyematan TPU yang besar mendapatkan manfaat paling besar.

Untuk beban kerja yang menggunakan paralelisme data sebagai satu-satunya strategi paralelisme, topologi yang tidak teratur mungkin berperforma sedikit lebih baik. Untuk LLM, performa menggunakan topologi twisted dapat bervariasi bergantung pada jenis paralelisme (DP, MP, dll.). Praktik terbaiknya adalah melatih LLM Anda dengan dan tanpa topologi yang diubah untuk menentukan mana yang memberikan performa terbaik untuk model Anda. Beberapa eksperimen pada model FSDP MaxText telah melihat peningkatan MFU sebesar 1-2 menggunakan topologi twisted.

Manfaat utama topologi berpilin adalah mengubah topologi torus asimetris (misalnya, 4×4×8) menjadi topologi simetris yang terkait erat. Topologi simetris memiliki banyak manfaat:

• Load balancing yang ditingkatkan
• Bandwidth biseksi yang lebih tinggi
• Rute paket yang lebih pendek

Manfaat ini pada akhirnya akan meningkatkan performa untuk banyak pola komunikasi global.

Software TPU mendukung tori terpilin pada slice dengan ukuran setiap dimensi sama dengan atau dua kali ukuran dimensi terkecil. Misalnya, 4x4x8, 4×8×8, atau 12x12x24.

Sebagai contoh, pertimbangkan topologi torus 4x2 ini dengan TPU yang diberi label dengan koordinat (X,Y) dalam slice:

Tepi dalam grafik topologi ini ditampilkan sebagai tepi tidak berarah agar lebih jelas. Dalam praktiknya, setiap tepi adalah koneksi dua arah antara TPU. Kita menyebut tepi antara satu sisi petak ini dan sisi yang berlawanan sebagai tepi melingkar, seperti yang ditunjukkan dalam diagram.

Dengan memutar topologi ini, kita akan mendapatkan topologi torus terpilin 4×2 yang sepenuhnya simetris:

Satu-satunya yang berubah antara diagram ini dan diagram sebelumnya adalah tepi wrap-around Y. Daripada terhubung ke TPU lain dengan koordinat X yang sama, TPU tersebut telah dipindahkan untuk terhubung ke TPU dengan koordinat X+2 mod 4.

Ide yang sama berlaku untuk ukuran dimensi yang berbeda dan jumlah dimensi yang berbeda. Jaringan yang dihasilkan bersifat simetris, selama setiap dimensi sama dengan atau dua kali ukuran dimensi terkecil.

Tabel berikut menunjukkan topologi twisted yang didukung dan peningkatan teoretis bandwidth biseksi dengan topologi tersebut dibandingkan dengan topologi yang tidak di-twist.

Topologi yang Tidak Sesuai	Peningkatan teoretis dalam bandwidth biseksi dibandingkan dengan torus yang tidak terpilin
4×4×8_twisted	~70%
8x8x16_twisted
12×12×24_twisted
4×8×8_twisted	~40%
8×16×16_twisted

Catatan: Topologi twisted tidak akan digunakan secara otomatis kecuali jika akhiran _twisted ditentukan dengan flag --topology. Misalnya, untuk menentukan topologi twisted 4x4x8, Anda akan menentukan flag topology sebagai: 4x4x8_twisted.

Varian topologi TPU v4

Beberapa topologi yang berisi jumlah chip yang sama dapat diatur dengan cara yang berbeda. Misalnya, slice TPU dengan 512 chip (1024 TensorCore) dapat dikonfigurasi menggunakan topologi berikut: 4x4x32, 4x8x16, atau 8x8x8. Slice TPU dengan 2.048 chip (4.096 TensorCore) menawarkan lebih banyak opsi topologi: 4x4x128, 4x8x64, 4x16x32, dan 8x16x16.

Topologi default yang terkait dengan jumlah chip tertentu adalah topologi yang paling mirip dengan kubus. Bentuk ini kemungkinan merupakan pilihan terbaik untuk pelatihan ML paralel data. Topologi lain dapat berguna untuk workload dengan beberapa jenis paralelisme (misalnya, paralelisme model dan data, atau partisi spasial simulasi). Workload ini akan berfungsi optimal jika topologinya cocok dengan paralelisme yang digunakan. Misalnya, menempatkan paralelisme model 4 arah pada dimensi X dan paralelisme data 256 arah pada dimensi Y dan Z cocok dengan topologi 4x16x16.

Model dengan beberapa dimensi paralelisme berperforma terbaik dengan dimensi paralelismenya yang dipetakan ke dimensi topologi TPU. Model ini biasanya merupakan model bahasa besar (LLM) paralel data dan model. Misalnya, untuk slice TPU v4 dengan topologi 8x16x16, dimensi topologi TPU adalah 8, 16, dan 16. Penggunaan paralelisme model 8 arah atau 16 arah (dipetakan ke salah satu dimensi topologi TPU fisik) akan lebih efektif. Paralelisme model 4 arah akan menjadi sub-optimal dengan topologi ini, karena tidak selaras dengan dimensi topologi TPU mana pun, tetapi akan optimal dengan topologi 4x16x32 pada jumlah chip yang sama.

Konfigurasi TPU v4 terdiri dari dua grup, yaitu yang memiliki topologi lebih kecil dari 64 chip (topologi kecil) dan yang memiliki topologi lebih besar dari 64 chip (topologi besar).

Topologi Small v4

Cloud TPU mendukung slice TPU v4 berikut yang lebih kecil dari 64 chip, kubus 4x4x4. Anda dapat membuat topologi v4 kecil ini menggunakan nama berbasis TensorCore (misalnya, v4-32), atau topologinya (misalnya, 2x2x4):

Nama (berdasarkan jumlah TensorCore)	Jumlah chip	Topologi
v4-8	4	2x2x1
v4-16	8	2x2x2
v4-32	16	2x2x4
v4-64	32	2x4x4

Topologi v4 besar

Slice TPU v4 tersedia dalam kenaikan 64 chip, dengan bentuk yang merupakan kelipatan 4 di ketiga dimensinya. Dimensi harus dalam urutan menaik. Beberapa contoh ditampilkan dalam tabel berikut. Beberapa topologi ini adalah topologi "kustom" yang hanya dapat dibuat menggunakan flag --type dan --topology karena ada lebih dari satu cara untuk mengatur chip.

Nama (berdasarkan jumlah TensorCore)	Jumlah chip	Topologi
v4-128	64	4x4x4
v4-256	128	4x4x8
v4-512	256	4x8x8
topologi kustom: harus menggunakan flag --type dan --topology	256	4x4x16
v4-1024	512	8x8x8
v4-1536	768	8x8x12
v4-2048	1024	8x8x16
topologi kustom: harus menggunakan flag --type dan --topology	1024	4x16x16
v4-4096	2048	8x16x16
…	…	…