TPU v4

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

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, 1.200 GBps
Daya min/rata-rata/maks yang diukur	90/170/192 W
Ukuran Pod TPU	4.096 chip
Topologi interkoneksi	Mesh 3D
Komputasi puncak per Pod	1,1 exaflops (bf16 atau int8)
Bandwidth all-reduce per Pod	1,1 PB/dtk
Bandwidth bisection per Pod	24 TB/dtk

Diagram berikut mengilustrasikan chip TPU v4.

Mesh 3D dan torus 3D

TPU v4 memiliki koneksi langsung ke chip tetangga terdekat dalam 3 dimensi, sehingga menghasilkan mesh 3D 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 informasi selengkapnya, lihat Topologi torus yang dipelintir.

Manfaat performa TPU v4 dibandingkan v3

Bagian ini menjelaskan manfaat performa TPU v4

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 langsung 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 diperlukan 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 binding node NUMA:

1. Instal alat command line numactl.

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

2. Ikat kode skrip Anda ke Node NUMA 0.

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

Aktifkan binding 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, yang memungkinkan koordinasi yang lebih baik antara dua TensorCore di chip.
• Meningkatkan performa HBM menggunakan standar dan kecepatan memori terbaru.
• Profil performa DMA yang ditingkatkan dengan dukungan bawaan untuk langkah berperforma tinggi dengan granularitas 512B.

TensorCores

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

Interkoneksi antar-chip

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

Lainnya

• Antarmuka PCIE gen3 x16 ke host (koneksi langsung).
• Model keamanan yang ditingkatkan.
• Efisiensi energi yang lebih baik.

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 Pod berukuran sama dengan beberapa cara. Saat membuat slice Pod TPU, Anda menentukan versi TPU dan jumlah resource TPU yang diperlukan. Saat membuat slice Pod TPU v4, Anda dapat menentukan jenis dan ukurannya dengan salah satu dari dua cara: AcceleratorType dan AccleratorConfig.

Catatan: AcceleratorType dan AccleratorConfig saling eksklusif; Anda hanya dapat menggunakan salah satunya saat membuat TPU.

Menggunakan AcceleratorType

Gunakan AcceleratorType jika Anda tidak menentukan topologi. Untuk mengonfigurasi TPU v4 menggunakan AcceleratorType, gunakan flag --accelerator-type saat membuat slice Pod TPU. Tetapkan --accelerator-type ke string yang berisi versi TPU dan jumlah TensorCore yang ingin Anda gunakan. Misalnya, untuk membuat slice Pod v4 dengan 32 TensorCore, Anda akan menggunakan --accelerator-type=v4-32.

Perintah berikut akan membuat slice Pod TPU v4 dengan 512 TensorCore menggunakan flag --accelerator-type:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --accelerator-type=v4-512 \
    --version=tpu-vm-tf-2.18.0-pod-pjrt

Angka setelah versi TPU (v4) menentukan jumlah TensorCore. Ada dua TensorCore di TPU v4, sehingga jumlah chip TPU adalah 512/2 = 256.

Menggunakan AcceleratorConfig

Gunakan AcceleratorConfig jika Anda ingin menyesuaikan topologi fisik slice TPU. Hal ini umumnya diperlukan untuk penyesuaian performa dengan slice Pod yang lebih besar dari 256 chip.

Untuk mengonfigurasi TPU v4 menggunakan AcceleratorConfig, gunakan flag --version dan --topology. Tetapkan --version ke versi TPU yang ingin Anda gunakan dan --topology ke pengaturan fisik chip TPU di slice Pod.

Anda menentukan topologi TPU menggunakan 3-tuple, AxBxC dengan A<=B<=C dan A, B, C semuanya <= 4 atau semuanya merupakan kelipatan bilangan bulat dari 4. Nilai A, B, dan C adalah jumlah chip dalam setiap tiga dimensi. Misalnya, untuk membuat slice Pod v4 dengan 16 chip, Anda akan menetapkan --version=v4 dan --topology=2x2x4.

Perintah berikut membuat slice Pod TPU v4 dengan 128 chip TPU yang disusun dalam array 4x4x8:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-vm-tf-2.18.0-pod-pjrt

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. Ini dikenal sebagai topologi torus yang dipelintir.

Ilustrasi berikut menunjukkan beberapa topologi TPU v4 umum.

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

Topologi torus yang terpuntir

Beberapa bentuk slice torus 3D v4 memiliki opsi untuk menggunakan topologi yang dikenal sebagai torus bengkok. Misalnya, dua kubus v4 dapat diatur sebagai slice 4x4x8 atau 4x4x8_twisted. Topologi twisted menawarkan bandwidth bisekan yang jauh lebih tinggi. Peningkatan bandwidth bisection berguna untuk beban kerja 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 twisted mungkin berperforma sedikit lebih baik. Untuk LLM, performa yang menggunakan topologi twisted dapat bervariasi bergantung pada jenis paralelisme (DP, MP, dll.). Praktik terbaik adalah melatih LLM dengan dan tanpa topologi twisted untuk menentukan mana yang memberikan performa terbaik untuk model Anda. Beberapa eksperimen pada model FSDP MaxText telah mengalami peningkatan 1-2 MFU menggunakan topologi twisted.

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

• Load balancing yang ditingkatkan
• Bandwidth bisection yang lebih tinggi
• Rute paket yang lebih singkat

Manfaat ini pada akhirnya menghasilkan peningkatan performa untuk banyak pola komunikasi global.

Software TPU mendukung torus bengkok 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) di slice:

Tepi dalam grafik topologi ini ditampilkan sebagai tepi tidak terarah untuk kejelasan. Dalam prakteknya, setiap tepi adalah koneksi dua arah antara TPU. Kita menyebut tepi antara satu sisi petak ini dan sisi yang berlawanan sebagai tepi yang memutar, seperti yang ditunjukkan dalam diagram.

Dengan memutar topologi ini, kita akan mendapatkan topologi torus twisted 4x2 yang sepenuhnya simetris:

Semua yang berubah antara diagram ini dan diagram sebelumnya adalah tepi yang memutar Y. Alih-alih terhubung ke TPU lain dengan koordinat X yang sama, TPU tersebut telah digeser untuk terhubung ke TPU dengan koordinat X+2 mod 4.

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

Lihat menggunakan AcceleratorConfig untuk mengetahui detail tentang cara menentukan konfigurasi torus bengkok saat membuat Cloud TPU.

Tabel berikut menunjukkan topologi twisted yang didukung dan peningkatan teoretis bandwidth bisection dengan topologi tersebut dibandingkan dengan topologi untwisted.

Topologi Twisted	Peningkatan teoretis dalam bandwidth bisection versus torus non-twisted
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 Pod dengan 512 chip (1024 TensorCore) dapat dikonfigurasi menggunakan topologi berikut: 4x4x32, 4x8x16, atau 8x8x8. Slice Pod 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 mungkin merupakan pilihan terbaik untuk pelatihan ML paralel data. Topologi lainnya dapat berguna untuk beban kerja dengan beberapa jenis paralelisme (misalnya, paralelisme model dan data, atau partisi spasial simulasi). Performa workload ini akan optimal jika topologi dicocokkan 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 paralelisme yang dipetakan ke dimensi topologi TPU. Model ini biasanya data dan model paralel model bahasa besar (LLM). Misalnya, untuk slice Pod 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 menghasilkan performa yang lebih baik. Paralelisme model 4 arah akan kurang optimal dengan topologi ini, karena tidak selaras dengan dimensi topologi TPU apa pun, tetapi akan optimal dengan topologi 4x16x32 pada jumlah chip yang sama.

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

Topologi v4 kecil

Cloud TPU mendukung slice Pod 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 Pod TPU v4 tersedia dalam penambahan 64 chip, dengan bentuk yang merupakan kelipatan 4 pada ketiga dimensi. Dimensi harus dalam urutan yang meningkat. 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
…	…	…