Pengantar Cloud TPU

Tensor Processing Unit (TPU) adalah sirkuit terintegrasi khusus aplikasi (ASIC) Google yang dikembangkan secara kustom dan digunakan untuk mempercepat beban kerja machine learning. Untuk mengetahui informasi selengkapnya tentang hardware TPU, lihat Arsitektur TPU. Cloud TPU adalah layanan web yang menyediakan TPU sebagai resource komputasi skalabel di Google Cloud.

TPU melatih model Anda secara lebih efisien menggunakan hardware yang dirancang untuk melakukan operasi matriks besar yang sering ditemukan dalam algoritma machine learning. TPU memiliki memori bandwidth tinggi (HBM) di chip yang memungkinkan Anda menggunakan model dan ukuran batch yang lebih besar. TPU dapat dihubungkan dalam grup yang disebut slice yang menskalakan beban kerja Anda dengan sedikit atau tanpa perubahan kode.

Kode yang berjalan di TPU harus dikompilasi oleh compiler akselerator linear algebra (XLA). XLA adalah compiler tepat waktu yang mengambil grafik yang dimunculkan oleh aplikasi framework ML dan mengompilasi komponen aljabar linear, kerugian, dan gradien grafik menjadi kode mesin TPU. Bagian program lainnya berjalan di mesin host TPU. Compiler XLA adalah bagian dari image VM TPU yang berjalan di mesin host TPU.

Kapan harus menggunakan TPU

Cloud TPU dioptimalkan untuk workload tertentu. Dalam beberapa situasi, Anda mungkin ingin menggunakan GPU atau CPU di instance Compute Engine untuk menjalankan workload machine learning. Secara umum, Anda dapat memutuskan hardware mana yang paling cocok untuk beban kerja Anda berdasarkan panduan berikut:

CPU

  • Pembuatan prototipe cepat yang memerlukan fleksibilitas maksimum
  • Model sederhana yang tidak memerlukan waktu lama untuk dilatih
  • Model kecil dengan ukuran batch kecil dan efektif
  • Model yang berisi banyak operasi TensorFlow kustom yang ditulis dalam C++
  • Model yang dibatasi oleh I/O yang tersedia atau bandwidth jaringan sistem host

GPU

  • Model dengan sejumlah besar operasi PyTorch/JAX kustom yang harus berjalan setidaknya sebagian di CPU
  • Model dengan operasi TensorFlow yang tidak tersedia di Cloud TPU (lihat daftar operasi TensorFlow yang tersedia)
  • Model sedang hingga besar dengan ukuran batch efektif yang lebih besar

TPU

  • Model yang didominasi oleh komputasi matriks
  • Model tanpa operasi PyTorch/JAX kustom di dalam loop pelatihan utama
  • Model yang dilatih selama berminggu-minggu atau berbulan-bulan
  • Model besar dengan ukuran batch efektif yang besar
  • Model dengan penyematan ultra-besar yang umum dalam beban kerja rekomendasi dan pengurutan lanjutan

Cloud TPU tidak cocok untuk workload berikut:

  • Program aljabar linear yang sering memerlukan cabang atau berisi banyak operasi aljabar menurut elemen
  • Workload yang memerlukan aritmetika presisi tinggi
  • Beban kerja jaringan saraf yang berisi operasi kustom dalam loop pelatihan utama

TPU di Google Cloud

Anda dapat menggunakan TPU melalui VM Cloud TPU, Google Kubernetes Engine, dan Vertex AI. Tabel berikut mencantumkan resource untuk setiap layanan Google Cloud.

LayananGoogle Cloud Resource
Cloud TPU Memulai VM Cloud TPU
Google Kubernetes Engine

Tentang TPU di GKE

Menjalankan Ray di GKE dengan TPU
Vertex AI

Pelatihan di Vertex AI dengan TPU

Menggunakan TPU untuk prediksi online di Vertex AI

Praktik terbaik untuk pengembangan model

Program yang komputasinya didominasi oleh operasi non-matriks seperti menambahkan, membentuk ulang, atau menyambungkan, kemungkinan tidak akan mencapai penggunaan MXU yang tinggi. Berikut adalah beberapa panduan untuk membantu Anda memilih dan mem-build model yang sesuai untuk Cloud TPU.

Tata Letak

Compiler XLA melakukan transformasi kode, termasuk menyusun matriks menjadi blok yang lebih kecil, untuk menjalankan komputasi secara efisien di unit matriks (MXU). Struktur hardware MXU, array sistolik 128x128, dan desain subsistem memori TPU, yang lebih memilih dimensi yang merupakan kelipatan 8, digunakan oleh compiler XLA untuk efisiensi tiling. Akibatnya, tata letak tertentu lebih kondusif untuk ubin, sementara tata letak lainnya memerlukan perubahan bentuk sebelum dapat disusun ubin. Operasi pembentukan ulang sering kali dibatasi memori di Cloud TPU.

Bentuk

Compiler XLA mengompilasi grafik ML tepat waktu untuk batch pertama. Jika batch berikutnya memiliki bentuk yang berbeda, model tidak akan berfungsi. (Mengompilasi ulang grafik setiap kali bentuk berubah terlalu lambat.) Oleh karena itu, model apa pun yang memiliki tensor dengan bentuk dinamis tidak cocok untuk TPU.

Padding

Program Cloud TPU berperforma tinggi adalah program yang komputasi padatnya dapat disusun menjadi 128x128 bagian. Jika komputasi matriks tidak dapat mengisi seluruh MXU, compiler akan mengisi tensor dengan nol. Ada dua kelemahan padding:

  • Tensor yang ditambahkan dengan nol tidak memanfaatkan core TPU secara maksimal.
  • Padding akan meningkatkan jumlah penyimpanan memori on-chip yang diperlukan untuk tensor dan dapat menyebabkan error kehabisan memori dalam kasus ekstrem.

Meskipun padding dilakukan secara otomatis oleh compiler XLA jika diperlukan, seseorang dapat menentukan jumlah padding yang dilakukan melalui alat op_profile. Anda dapat menghindari padding dengan memilih dimensi tensor yang sangat cocok untuk TPU.

Dimensi

Memilih dimensi tensor yang sesuai akan sangat membantu dalam mengekstrak performa maksimum dari hardware TPU, terutama MXU. Compiler XLA mencoba menggunakan ukuran batch atau dimensi fitur untuk menggunakan MXU secara maksimal. Oleh karena itu, salah satu dari keduanya harus kelipatan 128. Jika tidak, compiler akan menambahkan salah satunya menjadi 128. Idealnya, ukuran batch serta dimensi fitur harus merupakan kelipatan 8, yang memungkinkan ekstraksi performa tinggi dari subsistem memori.

Mulai menggunakan Cloud TPU

Meminta bantuan

Hubungi dukungan Cloud TPU. Jika Anda memiliki project Google Cloud yang aktif, bersiaplah untuk memberikan informasi berikut:

  • Google Cloud Project ID Anda
  • Nama TPU Anda, jika ada
  • Informasi lain yang ingin Anda berikan

Apa langkah selanjutnya?

Ingin mempelajari Cloud TPU lebih lanjut? Referensi berikut dapat membantu: