PyTorch di Cloud TPU

PyTorch di Cloud TPU memungkinkan Anda melatih, menyesuaikan, dan menyajikan model open source di hardware Cloud TPU menggunakan semantik PyTorch standar, API terdistribusi, dan alat ekosistem yang sudah Anda gunakan, termasuk Hugging Face, vLLM, SGLang, llm-d, TorchTitan, TorchRec, dan Ray.

Di dasar stack ini terdapat TorchTPU, backend native PyTorch open source yang dikembangkan bersama tim PyTorch di Meta. TorchTPU menyediakan library, framework, dan alat untuk setiap tahap siklus proses machine learning, mulai dari eksekusi runtime inti dan kernel kustom hingga inferensi produksi, pelatihan skala besar, dan orkestrasi cluster.

TorchTPU terhubung ke Hugging Face, TorchTitan, vLLM, SGLang, Ray, dan llm-d, dengan torch.profiler, DTensor, TorchRec, dan Pallas.

Mulai

  • Menjalankan model PyTorch di TPU: Mulai dengan TorchTPU dan dokumentasi TorchTPU.
  • Menyajikan LLM dan model multimodal: Gunakan vLLM-TorchTPU atau SGLang-TorchTPU.
  • Merutekan permintaan di seluruh fleet Kubernetes: Gunakan llm-d.
  • Melakukan pra-pelatihan atau penyesuaian model dasar: Gunakan TorchTitan di TPU (atau TorchRec untuk model rekomendasi).
  • Menskalakan pipeline pemrosesan data, pelatihan, RL, dan inferensi: Gunakan Ray.
  • Melakukan orientasi, membuat profil, dan menyesuaikan model dengan asisten coding Anda: Gunakan TPU agent suite (Pratinjau).

TorchTPU: Backend native PyTorch untuk Cloud TPU

TorchTPU menyediakan backend perangkat, integrasi compiler grafik, dan API kernel kustom untuk PyTorch di Cloud TPU. Dibuat bersama tim PyTorch, TorchTPU memungkinkan Anda men-debug model secara interaktif dalam mode eager, mengompilasi kode model dengan torch.compile, dan menskalakan di seluruh slice dengan torch.distributed (tpu_dist).

Untuk menginstal TorchTPU dan menjalankan operasi tensor di Cloud TPU:

pip install torch_tpu

import torch

device = torch.device("tpu")

x = torch.randn((4, 4), device=device)
y = torch.ones((4, 4), device=device)
result = torch.matmul(x, y).cpu()

print(f"Executed on device: {x.device}")
print(f"Result shape on CPU: {result.shape}, sum: {result.sum():.4f}")

Motivasi

Meskipun Cloud TPU secara historis telah dipasangkan dengan JAX dan TensorFlow, menjalankan PyTorch di TPU sebelumnya memerlukan perekaman grafik lazy-tensor (PyTorch/XLA), yang sering kali berarti menambahkan kode untuk menandai region yang dapat dikompilasi. TorchTPU menjalankan operasi secara eager secara default menggunakan model eksekusi satu proses per perangkat standar PyTorch, sehingga workload dan library yang ada berjalan tanpa penulisan ulang struktural.

Desain

TorchTPU menghadirkan dukungan TPU native PyTorch ke PyTorch dengan tiga sasaran desain:

  • Semantik PyTorch standar: Menjalankan satu proses per perangkat dan mempertahankan tampilan tensor PyTorch serta mutasi di tempat tanpa dependensi runtime JAX.
  • Eksekusi segera secara default: Model PyTorch standar langsung berjalan dalam mode segera untuk proses debug interaktif dan pemeriksaan langkah demi langkah.
  • Pengoptimalan compiler dan kernel: Selain eksekusi segera operator demi operator, TorchTPU menyediakan fusi operator dalam mode eager, integrasi torch.compile dengan batas bentuk dinamis untuk mengurangi kompilasi ulang, dan API kernel kustom.

Untuk kompilasi grafik, torch.compile menangkap grafik komputasi PyTorch dengan TorchDynamo dan mengompilasinya ke dalam kode mesin TPU melalui OpenXLA dan StableHLO. Operasi kolektif berjalan di tingkat grup proses torch.distributed (tpu_dist).

Kekuatan utama

  • Dukungan operator, model, dan Hugging Face: TorchTPU mendukung operator PyTorch (ATen) standar dalam penerusan maju dan mundur, serta terus diuji terhadap model Hugging Face populer. Anda dapat men-deploy model di Transformers v5 dengan meneruskan device_map="tpu", menjalankan pipeline Diffusers dengan mode eager, torch.compile, dan paralelisme tensor, serta membuat profil workload dengan torch.profiler di TensorBoard dan XProf.
  • Penskalaan chip tunggal ke multislice: Backend tpu_dist mendukung DistributedDataParallel (DDP), FullyShardedDataParallel (FSDP2), dan DTensor. Selain itu, torch.tpu.distributed.topology_aware_mesh memetakan dimensi paralelisme ke interkoneksi TPU fisik.
  • Mode eksekusi segera yang dapat dikonfigurasi: Menyediakan tiga mode eksekusi segera:
    • Strict (default): Mengirim operator satu per satu dan secara asinkron, dengan semantik PyTorch imperatif standar.
    • Debug: Berjalan secara sinkron dengan pemeriksaan batas tensor.
    • DeferAndFuse: Menggabungkan urutan operator eager untuk mengurangi traffic High Bandwidth Memory (HBM).
  • Kompilasi dan penyimpanan ke dalam cache grafik penuh: torch.compile(backend="tpu") mengompilasi grafik menjadi file yang dapat dieksekusi TPU dan menyimpannya ke dalam cache di seluruh proses, host, dan mulai ulang tugas.
  • Kernel TPU kustom: Anda dapat memanggil kernel Pallas kustom, termasuk kernel yang berkomunikasi langsung antar-chip, sebagai fungsi PyTorch standar, menggunakan kernel bawaan dari library Tokamax, atau menulis kernel berbasis tile portabel dengan Helion.

Untuk mengetahui informasi selengkapnya, lihat dokumentasi TorchTPU, ajukan pertanyaan di forum diskusi PyTorch di TPU, atau laporkan masalah di GitHub.

Inferensi dan penyajian produksi

Untuk deployment produksi, PyTorch di Cloud TPU terintegrasi dengan mesin inferensi open source dan infrastruktur penyaluran terdistribusi native Kubernetes.

vLLM-TorchTPU dan SGLang-TorchTPU: Mesin inferensi

vLLM-TorchTPU dan SGLang-TorchTPU adalah backend TorchTPU untuk mesin inferensi vLLM dan SGLang. Kedua backend menjalankan model terbuka seperti Qwen, DeepSeek, GLM, Kimi, GPT-OSS, dan Gemma 4. Untuk mengetahui informasi selengkapnya, baca postingan blog Mengumumkan Mesin Inferensi vLLM dan SGLang Baru di TPU.

Motivasi

Menyajikan workload penalaran, multimodal, dan penggunaan alat sering kali memerlukan definisi model kustom, metode decoding spekulatif, dan skema kuantisasi. Membangun backend vLLM dan SGLang di TorchTPU memungkinkan Anda menyesuaikan dan men-debug stack inferensi langsung di PyTorch.

Desain

Kedua backend mempertahankan penjadwal vLLM dan SGLang standar, pengelola cache KV, dan server API, sekaligus menggunakan eksekusi TorchTPU dan operator Pallas kustom di bawahnya. Sebagian besar model menjalankan definisi model PyTorch standar, sementara blok yang intensif komputasi memanggil kernel Pallas di dalam grafik yang dikompilasi sehingga kernel khusus hardware dapat diupdate secara independen.

Kekuatan utama

  • Paralelisme, kuantisasi, dan decoding spekulatif: Mendukung paralelisme tensor, data, pakar, pipeline, dan konteks di seluruh host, beserta input multimodal (gambar dan video) dan embedding. Metode decoding spekulatif yang didukung mencakup prediksi multi-token (MTP), EAGLE-3, DFlash, DSpark, dan pencocokan n-gram. Format kuantisasi yang didukung mencakup cache KV 8-bit (FP8) serta checkpoint model floating point 8-bit (FP8) dan 4-bit (MXFP4 dan NVFP4).
  • Caching awalan dua tingkat, panggilan alat, dan output terstruktur: Caching awalan otomatis (termasuk untuk model perhatian hibrid dan berulang) mempertahankan awalan aktif di HBM dan memindahkan awalan yang lebih lama ke memori host untuk menghindari komputasi ulang pada permintaan multi-giliran. Panggilan alat dan output terstruktur sudah disertakan.
  • Transfer cache KV yang tumpang-tindih dengan TPU Sync: Mendukung inferensi pengisian awal dan dekode yang terpisah (memisahkan pemrosesan perintah dan pembuatan token ke pool worker yang berbeda) di seluruh host dan offload KV host menggunakan TPU Sync. TPU Sync adalah library C++ open source yang membaca buffer perangkat PyTorch langsung di thread latar belakang tanpa menyalin data melalui memori CPU host.

Untuk mengetahui informasi selengkapnya tentang vLLM-TorchTPU, lihat repositori vLLM-TorchTPU, ajukan pertanyaan di saluran Slack vLLM, atau laporkan masalah vLLM-TorchTPU di GitHub.

Untuk mengetahui informasi selengkapnya tentang SGLang-TorchTPU, lihat repositori SGLang-TorchTPU atau laporkan masalah SGLang-TorchTPU di GitHub.

llm-d: Inferensi terdistribusi berbasis Kubernetes

llm-d adalah framework sandbox Cloud Native Computing Foundation (CNCF) open source untuk inferensi LLM terdistribusi di Kubernetes, yang dikembangkan bersama Red Hat, IBM, Google Cloud, dan komunitas yang terdiri dari lebih dari 500 kontributor dan 35 organisasi.

Motivasi

Meskipun mesin inferensi instance tunggal mengelola eksekusi per replika, fleet produksi harus menangani lonjakan traffic, prioritas multi-tenant, berbagi konteks di seluruh fleet, dan peluncuran tanpa periode nonaktif. llm-d menyediakan lapisan operasi dan pemilihan rute tingkat cluster untuk mengoordinasikan fleet server inferensi.

Desain

llm-d berjalan di depan server model seperti vLLM dan SGLang di Kubernetes. llm-d menyediakan hal berikut:

  • Merutekan permintaan berdasarkan status cache awalan dan beban replika.
  • Mengoordinasikan kumpulan worker prefill dan dekode yang dipisahkan.
  • Mendukung sharding MoE dan inferensi batch asinkron.
  • Mengelola offload cache KV bertingkat di seluruh memori perangkat, memori host, dan penyimpanan.

Kekuatan utama

  • Pemilihan rute sesi multi-turn: Merutekan percakapan multi-turn ke replika yang menyimpan konteks yang di-cache. Penjadwalan prioritas giliran dan dukungan Responses API stateful memungkinkan agen menggunakan kembali konteks di seluruh giliran untuk mengurangi komputasi ulang.
  • Kualitas layanan multi-tenant: Kontrol alur prioritas dan keadilan memungkinkan permintaan yang sensitif terhadap latensi dan workload agen yang bursty berbagi fleet tanpa kekurangan resource.
  • Operasi fleet dan update bertahap: Menyediakan API konfigurasi v1, ketersediaan tinggi router, dan update bertahap terkoordinasi yang mengelola replika pra-pengisian dan dekode sebagai satu unit logis.
  • Penyajian TorchTPU native di GKE: Anda dapat menyajikan Cloud TPU bersama akselerator lain di balik satu endpoint llm-d dan mengelola replika vLLM-TorchTPU dan SGLang-TorchTPU dengan alat Kubernetes standar.

Untuk mengetahui informasi selengkapnya, ajukan pertanyaan di saluran Slack llm-d, atau laporkan masalah di GitHub.

Library pelatihan, penyesuaian, dan domain

Untuk melakukan pra-pelatihan, penyesuaian, dan adaptasi model besar, TorchTPU menyediakan resep pelatihan referensi dan library khusus domain yang disesuaikan untuk hardware TPU.

TorchTitan di TPU: Referensi pelatihan berskala besar

TorchTitan di TPU memperluas TorchTitan, platform referensi modular PyTorch untuk pelatihan AI generatif, dengan penerusan sharding TPU, kernel kustom, dan resep pelatihan. Untuk mengetahui informasi selengkapnya, baca postingan blog Mengintegrasikan TorchTPU ke TorchTitan dan Mengoptimalkan Model untuk TPU.

Motivasi

Penskalaan pelatihan di seluruh slice TPU multi-host memerlukan parameter sharding di seluruh interkoneksi chip dan komunikasi yang tumpang-tindih dengan komputasi. Google dan Meta bersama-sama mengembangkan sharding TPU dan penerusan kompilasi di dalam TorchTitan, sehingga Anda tidak perlu merancang strategi sharding TPU kustom atau menjadwalkan komunikasi secara manual.

Desain

Model ditulis dalam PyTorch perangkat tunggal tanpa kode khusus akselerator. TorchTitan mengonfigurasi lapisan model dari file konfigurasi, sehingga Anda dapat mengganti lapisan individual dengan kernel Pallas kustom (seperti Splash Attention) tanpa mengedit definisi model atau loop pelatihan.

Implementasi SimpleFSDP TorchTitan menempatkan operasi sharding parameter langsung di dalam grafik yang dikompilasi sehingga compiler XLA dapat tumpang-tindih komunikasi perangkat dengan komputasi, sementara kompilasi grafik langkah penuh menangkap seluruh langkah pelatihan sebagai satu program XLA. Untuk operasi multislice, sharding hybrid (HSDP) melakukan sharding dalam setiap slice melalui interkoneksi antar-chip (ICI) dan mereplikasi di seluruh slice melalui jaringan pusat data (DCN). Tata letak ini membatasi traffic lintas-slice menjadi satu pengurangan gradien per langkah.

Kekuatan utama

  • Resep pra-pelatihan dan pasca-pelatihan RL: Mencakup resep referensi untuk LLM padat (Llama 3, Qwen 3, dan MuseGlimmer 30B), model mixture-of-experts (MoE) (Qwen 3 235B, divalidasi pada hingga 8.192 chip TPU7x), dan diffusion transformer (FLUX.1). Untuk pasca-pelatihan RL (seperti Group Relative Policy Optimization, atau GRPO), VeRL (Pratinjau) mendelegasikan pembaruan parameter ke TorchTitan dan mengirimkan bobot yang diperbarui ke mesin inferensi untuk menghasilkan lintasan pelatihan (rollout).
  • Setelan compiler dan paralelisme yang telah dikonfigurasi sebelumnya: Konfigurasi resep mencakup setelan compiler dan paralelisme yang telah dikonfigurasi sebelumnya sehingga Anda dapat menyesuaikan Model FLOPs Utilization (MFU) melalui file konfigurasi, bukan perubahan kode.
  • Orientasi untuk arsitektur baru: Backend pemodelan Hugging Face Transformers menjalankan checkpoint model Hugging Face Hub (termasuk model MoE tanpa paralelisme pakar) melalui pelatih TorchTitan. Untuk port model lengkap, TorchTitan menyertakan keterampilan orientasi langkah demi langkah untuk engineer dan agen coding.
  • Penggunaan kembali komponen mandiri: Anda dapat mengadopsi TorchTitan secara menyeluruh atau mengimpor modul individual, operator kustom, dan flag compiler ke dalam codebase pelatihan yang ada.

Untuk mengetahui informasi selengkapnya, lihat panduan memulai TorchTitan di TPU, ajukan pertanyaan di forum diskusi PyTorch, atau laporkan masalah di GitHub.

TorchRec: Sistem rekomendasi

TorchRec adalah library domain PyTorch yang menyediakan tabel embedding terdistribusi dan primitif sharding untuk model rekomendasi. Di Cloud TPU, TorchRec memindahkan pencarian tabel embedding sparse ke prosesor TPU SparseCore khusus sambil menjalankan operasi matriks dense di TPU TensorCore.

Untuk mengetahui informasi selengkapnya tentang TorchRec, lihat dokumentasi TorchRec, ajukan pertanyaan di forum diskusi PyTorch di TPU, atau laporkan masalah di GitHub.

Alat developer dan orkestrasi cluster

Untuk menskalakan pipeline multi-node dan mengotomatiskan orientasi model, ekosistem TorchTPU terintegrasi dengan Ray dan menyediakan ekstensi agen coding dengan akses ke dokumentasi TPU.

Ray: Pelatihan, inferensi, dan RL terdistribusi

Ray menskalakan workload Python dan AI di seluruh cluster menggunakan Ray Train, Ray Serve, Ray Data, Ray Core, dan framework RL seperti verl.

Motivasi

Pelatihan multislice dan pipeline multi-level (seperti alur kerja RL yang menggabungkan prapemrosesan data, pelatihan, dan pembuatan peluncuran) memerlukan penyediaan slice otomatis, penemuan topologi, dan koordinasi worker yang fault-tolerant. Utilitas peluncuran mandiri seperti torchrun tidak menyediakan slice TPU atau mengoordinasikan pipeline aktor multi-level, sehingga Ray mengelola penjadwalan dan penempatan di seluruh cluster untuk workload TorchTPU.

Desain

Saat Anda mengirimkan tugas yang menentukan pembuatan TPU, topologi, dan jumlah worker, Ray akan mencadangkan slice TPU lengkap sehingga semua worker dalam slice dijadwalkan bersama, menemukan alamat interkoneksi chip, dan memulai satu worker PyTorch per perangkat TPU. Backend PyTorch Ray Train menginisialisasi perangkat tpu dan grup proses tpu_dist secara otomatis sebelum memulai tugas.

Kekuatan utama

  • Penempatan yang sesuai topologi: Mencocokkan penempatan worker di seluruh host dengan generasi TPU fisik dan tata letak slice.
  • Pemulihan kegagalan otomatis: Menjalankan tugas pelatihan terdistribusi di beberapa slice TPU di GKE dengan pemulihan otomatis saat worker atau slice gagal.
  • Pengambilan data langsung ke TPU: Metode iter_torch_batches Ray Data mengambil batch sebagai tensor PyTorch langsung ke memori perangkat TPU.
  • Orkestrasi inferensi multi-host: Mengelola deployment inferensi multi-host untuk vLLM-TorchTPU dan SGLang-TorchTPU.

TPU agent suite: (Pratinjau) Agen coding untuk TPU

TPU agent suite adalah ekstensi open source untuk asisten coding—termasuk Claude Code dan Google Antigravity, dengan lebih banyak agen coding yang akan menyusul—yang berfokus pada inferensi transformer (LLM padat dan model MoE) di TPU v6e dan TPU7x.

Motivasi

Asisten coding tujuan umum sering kali tidak memiliki konteks terkini tentang topologi hardware TPU, flag compiler XLA, telemetri XProf, dan batasan kernel Pallas. TPU agent suite memberi asisten Anda akses ke dokumentasi TPU resmi dan alur kerja terstruktur. Hal ini membantu mencegah saran flag yang tidak valid dan mengotomatiskan pembuatan profil dan tolok ukur iteratif.

Desain

Suite ini menambahkan pusat informasi TPU dan enam sub-agen khusus tugas ke asisten coding yang sudah ada. Agen pengarah mendelegasikan setiap perintah ke sub-agen yang sesuai tanpa menjalankan perintah itu sendiri. Setiap eksekusi perintah di VM TPU atau cluster GKE Anda memerlukan persetujuan Anda, dan loop penyesuaian menjalankan sejumlah upaya yang tercatat sebelum menampilkan hasil.

Kekuatan utama

  • Orientasi model dan verifikasi akurasi: Agen Orientasi Model TPU menyelesaikan operator yang tidak didukung, memilih jalur inferensi (vLLM-TorchTPU atau PyTorch standar), dan membandingkan output lapisan demi lapisan dengan model referensi CPU untuk memverifikasi akurasi numerik.
  • Diagnostik error dan bottleneck: Agen TPU Diagnostics menganalisis log error, memeriksa profil XProf dan trace compiler XLA, serta mengisolasi bottleneck runtime.
  • Pengoptimalan kernel, compiler, dan sharding: Agen TPU Kernel Development membantu kernel Pallas kustom, agen TPU Performance Optimization membandingkan kombinasi flag compiler XLA, dan agen TPU Sharding merencanakan tata letak paralelisme multislice.
  • Pencarian dokumentasi dengan kutipan: Agen TPU Knowledge Base menjawab pertanyaan teknis dengan kutipan eksplisit ke dokumentasi dan tutorial TPU resmi, serta menyatakan jika suatu topik tidak dibahas.

Untuk mengetahui informasi selengkapnya, lihat dokumentasi TPU agent suite, bagikan masukan dan ajukan pertanyaan di GitHub Discussions, atau laporkan masalah di GitHub.

Langkah berikutnya