TorchTPU: PyTorch Native di TPU, 5.000 Model Hugging Face Diuji

Google bersama Meta merilis TorchTPU, backend PyTorch native untuk TPU yang membuat beban kerja PyTorch berjalan di perangkat TPU dengan perubahan kode minimal. Klaim yang menyertainya cukup spesifik: model 30 miliar parameter mencapai 25.700 token per detik pada panjang sekuens 5.120 di satu chip TPUv7, dan pipeline ujinya mengevaluasi sekitar 5.000 model Hugging Face dari lebih 500 arsitektur. Proyek ini masih tahap awal, dan roadmap-nya mencakup dynamic shapes, performa kompilasi, cakupan operator, reinforcement learning, serta kuantisasi.

TL;DR

• TorchTPU adalah backend PyTorch native untuk TPU, dikembangkan Google bersama Meta, dan dirilis open source.
• Pendahulunya, PyTorch/XLA, dipakai sejak 2020 tapi menuntut perubahan kode untuk menandai region yang bisa dikompilasi dan tidak mendukung mode eager.
• Mode DeferAndFuse diklaim Google 50 persen atau lebih cepat dari Strict Eager sekaligus menurunkan puncak memori HBM.
• Uji integrasi mencakup sekitar 5.000 model Hugging Face dan lebih 500 arsitektur, mulai dari TPU generasi v5e.
• Migrasi dari GPU melibatkan penggantian perangkat cuda menjadi tpu dan backend nccl menjadi tpu_dist.

Apa bedanya TorchTPU dengan PyTorch/XLA?

Bedanya di titik awal desainnya. PyTorch/XLA, yang dipakai Google sejak 2020, memanfaatkan JAX, XLA, dan lazy tensor: operasi dikumpulkan dulu, lalu dikompilasi menjadi satu graf sebelum dieksekusi. Pendekatan itu efisien untuk graf besar, tapi menuntut kode menandai secara eksplisit bagian mana yang boleh dikompilasi.

TorchTPU membalik prioritasnya: mode eager dijadikan hal utama, sehingga kode PyTorch yang sudah ada bisa berjalan dengan sedikit perubahan. Google menyebut empat kelemahan pendekatan lama, yaitu kebutuhan mengubah kode sumber secara intrusif, latensi kompilasi yang parah, tidak adanya dukungan untuk eksekusi imperatif native, dan tertinggalnya dukungan primitif komputasi terdistribusi baru di ekosistem PyTorch.

Aspek	PyTorch/XLA	TorchTPU
Model eksekusi	Lazy tensor, graf dikompilasi dulu	Eager native, plus mode terkompilasi
Perubahan kode	Perlu menandai region yang bisa dikompilasi	Minimal, terutama penggantian nama perangkat
Latensi kompilasi	Disebut parah oleh Google	Menjadi salah satu fokus roadmap
Dukungan eager	Tidak ada	Ada tiga mode, termasuk hybrid
Status	Jalur lama sejak 2020	Tahap awal, terus berkembang

Bagaimana cara memindahkan kode dari GPU ke TPU?

Untuk kasus paling sederhana, penggantiannya dua baris: perangkat cuda menjadi tpu dan backend nccl menjadi tpu_dist. Google juga menyediakan API torch.accelerator.current_accelerator agar kode yang sama bisa memilih akselerator yang tersedia tanpa cabang if yang ditulis manual.

import torch
import torch.distributed as dist

def main():
    device = torch.accelerator.current_accelerator
    dist.init_process_group()
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    model = model.to(device)

Dokumentasi API akselerator bisa dibaca di PyTorch Docs, sementara detail kompilator di balik eksekusi terkompilasi ada di proyek OpenXLA. Perlu dicatat, potongan kode di atas adalah penyederhanaan dari contoh di laporan: kode asli tetap memerlukan definisi model, loader data, dan alur pelatihan.

Yang menarik, perubahan kode yang kecil bukan berarti jaminan perilaku yang sama. Operator yang belum didukung, bentuk tensor dinamis, dan pola komunikasi antar perangkat adalah tiga hal yang paling sering berubah perilakunya ketika kode pindah antar akselerator. Karena itu pengujian model nyata, bukan sekadar menjalankan unit test, tetap wajib.

Berapa performa yang dilaporkan?

Angka paling konkret datang dari model 30 miliar parameter bernama Muse Glimmer yang diimplementasikan dengan PyTorch standar. Pada satu chip TPUv7, throughput inferensi mencapai 25.700 token per detik pada panjang sekuens 5.120, dan 22.400 token per detik pada panjang 10.240.

Untuk pelatihan, empat chip TPUv7 menghasilkan 5.600 token per detik pada panjang 5.120 dan 4.600 token per detik pada 10.240. Google juga menguji Qwen3 MoE 235B untuk pelatihan dengan skala dari 1.024 sampai 8.192 chip TPUv7, dan melaporkan performa per dolar yang kompetitif untuk beban MoE berskala besar saat inferensi memakai vLLM.

Skenario	Perangkat	Throughput
Inferensi, sekuens 5.120	1x TPUv7	25.700 token per detik
Inferensi, sekuens 10.240	1x TPUv7	22.400 token per detik
Pelatihan, sekuens 5.120	4x TPUv7	5.600 token per detik
Pelatihan, sekuens 10.240	4x TPUv7	4.600 token per detik

Angka-angka itu berasal dari pengukuran Google, bukan pengujian independen. Pembaca yang terbiasa dengan klaim vendor sebaiknya memperlakukannya sebagai batas atas yang dicapai pada konfigurasi tertentu, bukan angka yang otomatis berlaku untuk model lain dengan bentuk komputasi berbeda.

Apa itu DeferAndFuse dan kenapa penting?

DeferAndFuse adalah mode eksekusi hybrid yang menunda sebagian operasi untuk digabungkan, lalu menjalankannya sebagai bagian terkompilasi. TorchTPU menyediakan tiga mode eager: strict, debug, dan DeferAndFuse.

Google mengklaim mode ini bisa 50 persen atau lebih cepat daripada Strict Eager sekaligus menurunkan puncak pemakaian memori HBM. Klaim itu masuk akal secara mekanika: menggabungkan operasi kecil mengurangi jumlah perpindahan data, dan pada akselerator, perpindahan data sering lebih mahal daripada komputasinya.

Yang belum jelas dari laporan adalah seberapa konsisten keuntungan itu di berbagai arsitektur model. Mode hybrid biasanya paling menguntungkan pada model dengan banyak operasi kecil berurutan, dan bisa kehilangan manfaat pada model dengan operasi besar yang sudah efisien. Karena itu, memilih mode sebaiknya dilakukan lewat pengukuran pada model sendiri, bukan mengikuti angka umum.

Apa yang belum selesai?

Roadmap yang disebut mencakup dynamic shapes, performa kompilasi, cakupan operator, reinforcement learning, dan kuantisasi. Daftar itu memberi gambaran jelas soal batas hari ini: model dengan bentuk tensor yang berubah-ubah, kebutuhan kuantisasi untuk inferensi murah, dan beban pelatihan penguatan belum menjadi jalur utama.

Statusnya juga dinyatakan masih tahap awal dan terus berkembang. Bagi tim produksi, itu berarti TorchTPU lebih tepat dipakai untuk eksperimen dan evaluasi biaya daripada untuk memindahkan seluruh pelatihan. Bagi tim riset yang ingin menguji model di perangkat selain GPU, justru inilah waktu paling murah untuk mencoba.

Konteks tambahan: TPU tidak hanya tersedia lewat layanan awan Google, tapi juga dijual untuk penerapan di premises bagi pelanggan enterprise. Artinya keputusan memakai TPU tidak selalu berarti terikat ke satu penyedia awan, meski rincian komersialnya di luar cakupan laporan ini.

Apakah ini berarti TPU siap menggantikan GPU?

Belum. Yang berubah adalah biaya masuk: selama ini memindahkan kode PyTorch ke TPU menuntut penyesuaian yang tidak sepele, dan TorchTPU menurunkan hambatan itu. Tapi hambatan teknis hanyalah satu dari beberapa faktor.

Faktor lain adalah ketersediaan operator untuk model terbaru, kematangan tooling profiling, dan pengalaman tim. Ekosistem CUDA sudah matang lebih dari satu dekade, termasuk pustaka pihak ketiga yang tidak selalu punya padanan. Untuk model yang seluruh komponennya didukung, perpindahan bisa menarik secara biaya; untuk model dengan operator eksotis, biaya penyesuaian bisa menghabiskan penghematan.

Yang paling mungkin terjadi dalam waktu dekat adalah percampuran: pelatihan besar di satu jenis akselerator, inferensi di jenis lain, dan pemilihan berdasarkan perbandingan harga per token yang diukur sendiri. TorchTPU memperluas pilihan itu, bukan menghapus salah satunya.

Kapan tim sebaiknya mulai menguji TorchTPU?

Paling masuk akal ketika biaya inferensi sudah menjadi pos pengeluaran yang terlihat, bukan ketika pelatihan masih dalam tahap eksplorasi. Alasannya, jalur yang paling matang pada rilis ini adalah menjalankan model yang sudah dilatih, sementara beban pelatihan besar masih menghadapi batas cakupan operator dan bentuk tensor dinamis.

Kondisi kedua adalah ketersediaan akses ke perangkat. Tanpa akses nyata ke TPU, pengujian hanya akan menghasilkan catatan bahwa kode berhasil diimpor. Uji yang berguna harus mengukur throughput pada panjang sekuens yang benar-benar dipakai produksi, bukan pada satu panjang yang kebetulan paling menguntungkan.

Kondisi ketiga adalah kesiapan tim untuk membaca profil. Perpindahan akselerator hampir selalu mengubah di mana waktu habis: kompilasi, perpindahan memori, atau komunikasi antar chip. Tim yang tidak punya kebiasaan profiling akan menyimpulkan bahwa TPU lambat, padahal masalahnya ada pada satu operator yang jatuh ke jalur lambat.

Bagaimana mengukur manfaatnya sebelum memutuskan pindah?

Mulai dari satu model, satu panjang sekuens, dan satu ukuran batch, lalu bandingkan token per detik serta biaya per juta token pada konfigurasi yang setara. Perbandingan harus menyertakan waktu kompilasi, karena pada mode terkompilasi biaya itu muncul di awal dan bisa mendominasi untuk beban berdurasi pendek.

Kedua, ukur puncak memori. Klaim penurunan puncak HBM pada mode DeferAndFuse hanya berguna kalau memang itu yang membatasi jumlah model yang bisa dijalankan bersamaan. Kalau batasnya di tempat lain, penurunan memori tidak mengubah biaya.

Ketiga, catat kegagalan. Setiap operator yang belum didukung adalah pekerjaan tambahan yang harus diperkirakan, bukan kejutan yang ditemukan saat rilis. Daftar kegagalan pada satu model biasanya cukup untuk memperkirakan biaya portofolio model berikutnya.

Keempat, jangan membandingkan angka vendor dengan angka sendiri di konfigurasi berbeda. Angka 25.700 token per detik itu berguna sebagai penanda bahwa perangkatnya sanggup, bukan sebagai target yang harus dikejar di setiap model.

Untuk tim di Indonesia, relevansi jangka pendeknya lebih pada pilihan arsitektur daripada pada pemakaian langsung. Selama bertahun-tahun, kode yang ditulis dengan asumsi CUDA cenderung menempel pada satu jenis perangkat, dan biaya keluar dari asumsi itu mahal. Adanya backend yang membuat panggilan perangkat tidak lagi berarti satu vendor membuka opsi negosiasi harga ketika kebutuhan komputasi naik.

FAQ

Siapa yang mengembangkan TorchTPU?

Google, bekerja sama dengan Meta, dan dirilis sebagai proyek open source. Laporan tidak menyebut nomor versi rilis atau lisensi spesifik di luar statusnya sebagai proyek terbuka.

Apakah kode PyTorch yang ada harus ditulis ulang?

Tidak sepenuhnya. Perubahan minimalnya adalah mengganti perangkat cuda menjadi tpu dan backend nccl menjadi tpu_dist, atau memakai API pemilih akselerator. Komponen yang belum didukung tetap perlu penyesuaian.

Berapa model yang diuji dalam pipeline integrasi?

Sekitar 5.000 model Hugging Face dari lebih 500 arsitektur, diuji dalam mode eager dan terkompilasi pada beberapa generasi TPU mulai dari v5e.

Apakah angka 25.700 token per detik berlaku umum?

Tidak. Angka itu hasil pengukuran pada model 30 miliar parameter dengan panjang sekuens 5.120 di satu chip TPUv7, dan tidak otomatis berlaku untuk model atau konfigurasi lain.

Apa saja yang masih dalam roadmap?

Dynamic shapes, performa kompilasi, cakupan operator, reinforcement learning, dan kuantisasi. Tim menyatakan proyek masih tahap awal dan akan terus berubah.