TPU v4

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

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 kunci untuk Pod TPU v4.

Spesifikasi utama Nilai Pod v4
Komputasi puncak per chip 275 teraflops (bf16 atau int8)
Kapasitas dan bandwidth HBM2 32 GiB, 1200 GBps
Daya min/rata-rata/maks yang diukur 90/170/192 W
Ukuran Pod TPU 4096 chip
Topologi interkoneksi Mesh 3D
Komputasi puncak per Pod 1.1 exaflops (bf16 atau int8)
Semua bandwidth yang dikurangi per Pod 1,1 PB/dtk
Bandwidth Bisection per Pod 24 TB/dtk

Diagram berikut mengilustrasikan chip TPU v4.

Diagram 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 beberapa potongan dengan topologi, AxBxC, adalah 2A=B=C atau 2A=2B=C, dengan setiap dimensi merupakan 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 twisted.

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 {i>node<i} NUMA. NUMA node 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 mesin multi-CPU dapat menempatkan data yang dibutuhkan oleh CPU di dalam node NUMA-nya, sehingga meningkatkan throughput memori. Untuk informasi selengkapnya tentang NUMA, lihat Non Uniform Memory Access di Wikipedia.

Anda dapat memanfaatkan manfaat NUMA-lokalitas dengan mengikat skrip pelatihan Anda ke NUMA Node 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 NUMA Node 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, workload rekomendasi), terlepas dari framework-nya.
  • 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 kedua TensorCore pada chip.
  • Peningkatan performa HBM menggunakan kecepatan dan standar memori terbaru.
  • Peningkatan profil performa DMA dengan dukungan bawaan untuk performa yang tinggi dengan perincian 512 miliar.

TensorCores

  • Dua kali jumlah MXU dan kecepatan clock yang lebih tinggi yang memberikan 275 TFLOPS maks.
  • 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.

Inter-chip interkoneksi

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

Other

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

Konfigurasi

Pod TPU v4 terdiri dari 4096 chip yang saling terhubung dengan link berkecepatan tinggi yang dapat dikonfigurasi ulang. Jaringan fleksibel TPU v4 memungkinkan Anda menghubungkan chip dalam irisan Pod yang berukuran sama melalui beberapa cara. Saat membuat slice Pod TPU, Anda harus 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.

Menggunakan AcceleratorType

Gunakan AcceleratorType saat 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 digunakan. Misalnya, untuk membuat slice Pod v4 dengan 32 TensorCore, Anda harus menggunakan --accelerator-type=v4-32.

Perintah berikut 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.16.1-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 saat 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 --version dan flag --topology. Tetapkan --version ke versi TPU yang ingin Anda gunakan dan --topology ke susunan fisik chip TPU di slice Pod.

Anda menentukan topologi TPU menggunakan AxBxC 3-tuple dengan A<=B<=C dan A, B, C adalah semua <= 4 atau semuanya merupakan kelipatan bilangan bulat 4. Nilai A, B, dan C adalah jumlah chip di ketiga dimensi tersebut. Misalnya, untuk membuat slice Pod v4 dengan 16 chip, Anda harus 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.16.1-pod-pjrt

Topologi dengan 2A=B=C atau 2A=2B=C juga memiliki varian topologi yang dioptimalkan untuk komunikasi menyeluruh, misalnya, 4×4×8, 8×8×16, dan 12×12×24. Ini dikenal sebagai topologi tori tori.

Ilustrasi berikut menunjukkan beberapa topologi TPU v4 yang umum.

Topologi konfigurasi TPU v4 umum

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

Topologi tori bengkok

Beberapa bentuk irisan torus v4 3D memiliki opsi untuk menggunakan topologi torus terpilin. Misalnya, dua kubus v4 dapat disusun sebagai irisan 4x4x8 atau 4x4x8_twisted. Topologi twisted menawarkan bandwidth dua bagian yang jauh lebih tinggi. Peningkatan bandwidth bagian berguna untuk workload yang menggunakan pola komunikasi global. Topologi terpilin dapat meningkatkan performa untuk sebagian besar model, dengan beban kerja penyematan TPU yang paling bermanfaat.

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

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

  • Load balancing yang lebih baik
  • Bandwidth dua bagian yang lebih tinggi
  • Rute paket yang lebih singkat

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

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

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

Tepi dalam grafik topologi ini ditampilkan sebagai tepi yang tidak mengarah agar lebih jelas. Dalam praktiknya, setiap edge merupakan koneksi dua arah antara TPU. Kami menyebut tepi antara satu sisi petak ini dan sisi yang berlawanan sebagai tepi menyeluruh, seperti yang ditunjukkan dalam diagram.

drawing

Dengan memutar topologi ini, kita akan mendapatkan topologi torus 4×2 yang benar-benar simetris:

drawing

Semua yang berubah antara diagram ini dan sebelumnya adalah tepi wrap-around Y. Alih-alih menghubungkan 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 akan menggeneralisasi ke ukuran dimensi yang berbeda dan jumlah dimensi yang berbeda. Jaringan yang dihasilkan bersifat simetris, asalkan setiap dimensi sama dengan atau dua kali ukuran dimensi terkecil.

Lihat menggunakan AcceleratorConfig untuk detail tentang cara menentukan konfigurasi tori terpilin saat membuat Cloud TPU.

Tabel berikut menunjukkan topologi terpilin yang didukung dan peningkatan teoretis bandwidth bagian keduanya dibandingkan dengan topologi tidak terpilin.

Topologi Twisted Peningkatan secara teoritis dalam dua bagian
bandwidth 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

Varian topologi TPU v4

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

Topologi default yang terkait dengan jumlah chip tertentu adalah yang paling mirip dengan kubus. Bentuk ini kemungkinan merupakan pilihan terbaik untuk pelatihan ML paralel data. Topologi lainnya dapat berguna untuk workload dengan berbagai jenis paralelisme (misalnya, paralelisme model dan data, atau partisi spasial simulasi). Workload ini akan berfungsi optimal jika topologinya 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 berbagai dimensi paralelisme memiliki performa terbaik dengan dimensi paralelismenya yang dipetakan ke dimensi topologi TPU. Ini biasanya data dan model bahasa besar (LLM) paralel. Misalnya, untuk slice Pod TPU v4 dengan topologi 8x16x16, dimensi topologi TPU adalah 8, 16, dan 16. Akan lebih efektif jika menggunakan paralelisme model 8 arah atau 16 arah (dipetakan ke salah satu dimensi topologi TPU fisik). Paralelisme model 4 arah akan kurang optimal dengan topologi ini, karena tidak selaras dengan dimensi topologi TPU apa pun, tetapi akan lebih optimal dengan topologi 4x16x32 pada jumlah chip yang sama.

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

Topologi v4 kecil

Cloud TPU mendukung slice Pod TPU v4 berikut yang lebih kecil dari 64 chip, yaitu 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 kelipatan 64 chip, dengan bentuk yang merupakan perkalian 4 pada ketiga dimensi. Dimensi harus dalam urutan naik. 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
T/A - harus menggunakan tanda --type dan --topology 256 4x4x16
v4-1024 512 8x8x8
v4-1536 768 8x8x12
v4-2048 1024 8x8x16
T/A - harus menggunakan tanda --type dan --topology 1024 4x16x16
v4-4096 2048 8x16x16
… … …