Google ve Meta, PyTorch'u TPU'larda yerel olarak çalıştırmak için TorchTPU'yu yayımladı
Ekibe göre, 2025'in ortasında Google'da Meta'nın PyTorch Ekibi ile başlatılan açık kaynaklı backend, en büyük Google Cloud müşterilerinden bazıları tarafından üretimde zaten kullanılıyor.
Yapay zekâ illüstrasyonuGoogle ve Meta, PyTorch'un Google'ın TPU çiplerinde yerel olarak çalışmasını sağlayan açık kaynaklı bir backend olan TorchTPU'yu tanıttı. Projenin ekibi, projenin 2025'in ortasında Google'da Meta'daki PyTorch Ekibi ile iş birliği içinde başladığını ve etkin geliştirme aşamasında olduğunu söyledi.
Ekibe göre TorchTPU, en büyük Google Cloud müşterilerinden bazıları tarafından çıkarım ve eğitim iş yükleri için üretimde zaten kullanılıyor. Ekip, TPU'ların tarihsel olarak TensorFlow ve JAX ekosistemlerine bağlı olduğunu belirtti.
Ekip, TorchTPU'yu önceki girişimlerden bir kopuş olarak tanımladı. Google, 2020'de TPU için Lazy Tensors tabanlı, PyTorch'un ilk derlenmiş sürümü olan PyTorch/XLA'yı yayımlamıştı. Ekip, bu entegrasyonun kaynak kodda müdahaleci değişiklikler gerektirdiğini, ciddi derleme süresi gecikmesinden muzdarip olduğunu, PyTorch'un eager yürütme modunu desteklemediğini ve yeni dağıtık bilgi işlem özelliklerinde geride kaldığını söyledi. TorchTPU ise eager yürütmeyi en yüksek öncelik olarak ele alıyor ve çalışma zamanında hiçbir JAX bağımlılığı içermiyor.
Belirtilen hedef, bir iş yükünü TPU'ya taşımanın, koddaki hedef aygıtı GPU'dan TPU'ya değiştirmekten biraz daha fazlasını gerektirmemesi. Ekip, bir modelin şekilleri, sharding yapısı ve işlemleri GPU'lar için ayarlanmış olabileceğinden, tek başına böyle bir değişikliğin en yüksek TPU performansını sağlamayacağı uyarısında bulundu, ancak modelin yine de kutudan çıktığı gibi çalışması gerektiğini söyledi.
Ekibe göre proje şu ana kadar şunları yaptı:
- DeferAndFuse adlı yeni bir eager modu da dahil olmak üzere eager ve derlenmiş mod desteği ekledi;
- yaklaşık 450 ATen işlemi için destek ekledi;
- JAX/Pallas ve Helion DSL aracılığıyla yazılan özel kernel'leri etkinleştirdi;
- en çok indirilen 5,000 Hugging Face modelinde kutudan çıktığı gibi yürütmeyi doğruladı;
- vLLM, SGLang, TorchTitan, Ray ve TorchRec ile entegrasyonu doğruladı.
Ekip ayrıca sınırlamaları da sıraladı. Bir TPU aygıtına yalnızca tek bir kullanıcı alanı işlemi erişebildiğinden, pekiştirmeli öğrenmede birlikte konumlandırma kutudan çıktığı gibi desteklenemiyor; ekip bunu aşmanın yollarını araştırdığını söyledi. Ayrıca alttaki derleyici ve çalışma zamanı tensör şekillerinin derleme zamanında bilindiğini varsaydığından, sınırlı dinamizm desteği de ekliyor.
Hassasiyet konusunda TorchTPU, varsayılan olarak bfloat16'da çarpıp FP32'de biriktiren bir orta modu kullanır. Daha yüksek modlar, yaklaşık 3 kat veya yaklaşık 6 kat hesaplama maliyetiyle daha fazla FP32 hassasiyetini geri kazandırır ve ekip çoğu derin öğrenme iş yükü için varsayılanı öneriyor.