- TPU adalah chip kustom yang dikembangkan oleh Google untuk pelatihan dan inferensi AI skala besar, dengan filosofi desain yang berbeda dari GPU
- Menekankan skalabilitas dan efisiensi energi, dengan merancang perangkat keras (misalnya konfigurasi system-on-chip, memori on-chip berukuran besar) dan perangkat lunak (XLA compiler) secara bersamaan
- Struktur intinya adalah systolic array, pipelining, dan pendekatan kompilasi terlebih dahulu, yang dioptimalkan untuk sebagian besar operasi deep learning (terutama perkalian matriks)
- Berkat teknologi OCI dan OCS, dimungkinkan konfigurasi node yang fleksibel, pemrosesan paralel berperforma tinggi, dan pemilihan berbagai topologi
- Sistem TPU dapat diskalakan secara hierarkis dari satu chip hingga multipod raksasa, sehingga memungkinkan pelatihan model AI berskala besar dan pemanfaatan sumber daya yang efisien
Gambaran Umum dan Latar Belakang
- TPU adalah chip berbasis ASIC yang dikembangkan Google untuk menghadapi ekspansi besar-besaran layanan AI
- Pada awalnya Google mempertimbangkan GPU, FPGA, dan ASIC, tetapi seiring meluasnya pemanfaatan deep learning seperti pencarian suara pada 2013, kebutuhan akan perangkat keras kustom meningkat
- Saat ini TPU berperan penting dalam layanan AI utama seperti Gemini dan Veo, serta banyak digunakan pada model rekomendasi seperti DLRM
Struktur Satu Chip TPU
Konfigurasi Dasar
- Berdasarkan TPUv4, satu chip memiliki 2 TensorCore (TPU yang dioptimalkan untuk inferensi memiliki 1)
- Masing-masing TensorCore terhubung ke perangkat memori CMEM (128MiB) dan HBM (32GiB)
Struktur Internal TensorCore
- Matrix Multiply Unit (MXU) : berbasis systolic array 128x128, khusus untuk perkalian matriks
- Vector Unit (VPU) : menjalankan operasi elemen demi elemen umum
- Vector Memory (VMEM; 32MiB) : menyalin data dari HBM dan digunakan sebagai memori persiapan kerja
- Scalar Unit + Scalar Memory (SMEM; 10MiB) : menangani alur kontrol, operasi skalar, dan pengelolaan alamat memori
Perbedaan Struktural dengan GPU
- TPU memiliki memori on-chip (CMEM, VMEM, SMEM) yang jauh lebih besar dibanding GPU
- Kapasitas HBM lebih besar pada GPU, dan jumlah core komputasinya juga jauh lebih banyak pada GPU
- Berdasarkan TPUv5p, performa tinggi dicapai pada level 500 TFLOPs/sec per chip, dan 4.45 ExaFLOPs/sec untuk satu pod penuh (8960 chip)
Filosofi Desain TPU
1. Systolic Array dan Pipelining
- Systolic array adalah susunan elemen pemrosesan (PE) yang meneruskan hasil perhitungan ke elemen di sebelahnya
- Setelah data dimasukkan, komputasi berkelanjutan dapat dilakukan tanpa kontrol tambahan, dan operasi baca/tulis memori hanya terjadi pada input dan output
- Terutama dioptimalkan untuk perkalian matriks dan konvolusi
- Melalui pemrosesan pipeline, komputasi dan perpindahan data berlangsung secara bersamaan untuk mengoptimalkan throughput
Kekurangan Systolic Array - Sparsity
- Semua elemen komputasi selalu aktif sehingga kurang cocok untuk matriks sparse
- Jika model DL di masa depan makin mengarah pada sparsity tidak teratur, hal ini bisa menjadi keterbatasan
2. Kompilasi di Muka (AoT) dan Minimasi Cache
- Melalui co-design TPU-XLA, ketergantungan pada cache yang diperlukan untuk akses memori tidak teratur diminimalkan sehingga konsumsi energi berkurang
- Compiler XLA menganalisis graph komputasi untuk menghitung pola akses memori sebelumnya, lalu mengandalkan scratchpad memory alih-alih cache
@jit pada JAX adalah bentuk peralihan antara JIT dan AoT; saat eksekusi pertama, graph statis dibuat lalu dikompilasi AoT oleh XLA- Jika bentuk input berubah, kompilasi ulang diperlukan, sehingga ada ketidakefisienan pada padding atau iterasi dinamis
- Mencapai efisiensi energi yang tinggi, tetapi kekurangan fleksibilitas menjadi kelemahannya
Efisiensi Energi TPUv4
- Chip modern menggunakan memori HBM3 untuk mengurangi konsumsi energi
- Dibanding komputasi, pekerjaan memori mengonsumsi energi puluhan hingga ratusan kali lebih besar, sehingga meminimalkan akses memori dapat sangat meningkatkan efisiensi
Struktur Multi-Chip TPU
Level Tray (Tray/Board; 4 chip)
- 1 tray terdiri dari 4 chip TPU (8 TensorCore) dan CPU host
- Koneksi Host↔Chip menggunakan PCIe, sedangkan koneksi Chip↔Chip menggunakan Inter-Core Interconnect (ICI) yang menyediakan bandwidth lebih tinggi
Level Rack (Rack; 4x4x4=64 chip)
- 1 rack terdiri dari 64 chip TPU, terhubung melalui ICI dan OCS (Optical Circuit Switching) dalam 3D torus (4x4x4)
- Di Google, rack, pod, dan slice dibedakan sebagai berikut
- Rack: unit fisik 64 chip (=kubus)
- Pod: unit maksimum yang dapat dihubungkan dengan ICI dan OCS (misalnya TPUv4=4096 chip=64 rack)
- Slice: unit abstrak dengan konfigurasi sewenang-wenang dari 4 chip hingga Superpod
Keunggulan OCS
- Wraparound : setiap sumbu dibentuk menjadi ring (1D torus) untuk mengurangi jumlah hop terburuk antar-node
- Slice nonkontigu yang fleksibel: dengan struktur switching OCS, node yang terpisah secara fisik tetap dapat dibentuk menjadi satu slice, sehingga memudahkan pemanfaatan sumber daya dan pemeliharaan
- Topologi twisted: dengan jumlah chip yang sama (x,y,z tetap), struktur koneksi dapat diubah untuk mengoptimalkan kecepatan pada pola komputasi tertentu (misalnya twisted torus)
Contoh Pemanfaatan Topologi
- Bentuk kubus: cocok untuk paralelisme data/tensor (bandwidth maksimum)
- Bentuk garis lurus: menguntungkan untuk paralelisme pipeline
- Twisted torus: meningkatkan kecepatan saat komunikasi all-to-all diperlukan (misalnya tensor parallel)
Superpod (Full Pod/Superpod; TPUv4: 4096 chip/64 rack)
- Beberapa rack dihubungkan dengan ICI dan OCS untuk membangun sistem berskala sangat besar
- Bergantung pada jenis topologi slice, bandwidth komunikasi dan performa pemrosesan paralel dapat berbeda
- Berkat OCS, diperoleh fleksibilitas seperti slice nonkontigu dan topologi twisted
Level Multipod (Multi-pod/Multislice; TPUv4: lebih dari 4096)
- Beberapa pod dihubungkan melalui data center network (DCN) untuk menyediakan infrastruktur pelatihan skala besar, meski bandwidth-nya lebih rendah dibanding ICI
- Pelatihan PaLM menggunakan 2 pod (6144 TPUv4) dan dijalankan dengan pengelolaan sumber daya total 6 pod
- Saat mengembangkan model berskala besar, compiler XLA mengoptimalkan pola komunikasi berdasarkan pengaturan peneliti (seperti dimensi paralelisme)
- XLA menyisipkan operasi komunikasi antar-slice dan antar-pod, sehingga pelatihan terdistribusi skala besar dapat diwujudkan dengan perubahan kode minimal
Contoh Perangkat Keras Nyata dan Diagram
- Rack TPU: 3D torus 4x4x4 sebagai satu unit, dengan 2 tray (8 chip) di setiap baris
- Tray TPUv4: secara nyata memiliki 4 port PCIe (masing-masing untuk 1 TPU)
- Chip TPUv4: ASIC di tengah, dikelilingi 4 stack HBM (berdasarkan 2 TensorCore)
- Floorplan chip TPUv4i (tipe inferensi): 1 TensorCore, dengan area CMEM yang luas
Penutup
- Menyampaikan apresiasi kepada Google TPU Research Cloud (TRC) atas dukungan riset
Referensi
Belum ada komentar.