Google ve Meta, Google TPU'ları için Yerel PyTorch Arka Ucu TorchTPU'yu Tanıttı

TorchTPU, Google ve Meta arasında yerel bir PyTorch arka ucu sunan açık kaynaklı bir iş birliğidir. İstekli yürütme optimizasyonları, torch.compile entegrasyonu ve yerel donanım hızlandırması sayesinde, kutudan çıktığı gibi model çalıştırmayı ve TPU'ya geçişi sağlarken sektörde rekabetçi çıkarım ve eğitim performansını hedefler.

Google TPU'ları, büyük ML modellerini devasa ölçekte eğitmek ve dağıtmak için kullanılmıştır. TPU'lar tarihsel olarak TensorFlow ve JAX ile ilişkilendirilse de, bu çaba 2025'in ortasında Google'da Meta'daki PyTorch Ekibi ile başladı. 2020'de tanıtılan ve Lazy Tensors'a dayanan önceki PyTorch/XLA entegrasyonu, kaynak kodda derin müdahaleler gerektiriyordu, ciddi derleme gecikmeleri yaşıyordu, istekli yürütme desteği yoktu ve dağıtık primitiflerde geride kalıyordu.

Tasarım ilkeleri ve kullanılabilirlik

TorchTPU, JAX uyarlaması yerine PyTorch için ilkeli bir TPU arka ucu olarak oluşturuldu. İlkeleri arasında PyTorch vatandaşlığı (çalışma zamanında JAX bağımlılığı yok, sağlam istekli mod, Torch Dynamo tabanlı torch.compile ve DDP, FSDP ve DTensor desteği), sorunsuz ekosistem entegrasyonu ve topluluk katılımı yer alır. Kullanıcı odaklılık, taşıma kolaylığı anlamına gelir: model.to("cuda") ifadesini model.to("tpu") ile değiştirmek, bir modelin kutudan çıktığı gibi çalışmasını sağlamalıdır; ancak en yüksek performans için ek ayar gerekebilir. Proje ayrıca, şekil değişikliklerinden kaynaklanan XLA yeniden derlemesini önlemek için sınırlı dinamizmi hedefler ve önceden derlenebilen dinamik TPU çekirdeklerini araştırır. TorchTPU wheel paketinin yüklü olduğu bir TPU makinesinde PyTorch'u içe aktarmak, PyTorch'ın otomatik yükleme mekanizması aracılığıyla yalnızca küçük değişiklikler gerektirir.

İstekli yürütme ve derleme

TorchTPU, PrivateUse1 olarak kaydedilmiş bir C++ arka ucu ile PyTorch dağıtıcısını genişletir ve TPU tensör depolamasını standart at::Tensor mekanizması üzerinden tanımlar. Üç istekli mod sunar: Strict (varsayılan) her Aten işlemini zaman uyumsuz olarak tek tek gönderir; Debug (TORCH_TPU_LAUNCH_BLOCKING=1) ek sınır kontrolleriyle eşzamanlıdır; DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1) ise birden fazla işlemi erteler ve birleştirir; bu modun Strict istekli moddan %50 veya daha fazla hızlı olduğu ve tepe HBM belleğini azalttığı bildirilir. .to("cpu") çağrısı gibi bir somutlaştırma noktasında, TorchTPU canlılık analizini uygular, işlem grafiği için bir önbellek anahtarı oluşturur ve önbelleğe alınmış çalıştırılabilir dosyayı gönderir veya zaman uyumsuz olarak derler. Tekrarlayan İşlem Sezgisi, transformer blok yığınları gibi tekrarlayan dizileri algılar ve somutlaştırmayı erkenden tetikler. Üç katmanlı derleme önbelleği, süreç yerel belleği, makine yerel paylaşımlı belleği ve kalıcı uzak dosya sistemi önbelleğini kapsayarak konaklar ve yeniden başlatmalar arasında yeniden kullanım sağlar.

Desteklenen iş yükleri ve hassasiyet

TorchTPU, ana mimariler için gereken Core ATen işlemlerinin ileri ve geri geçişleri dahil yaklaşık 450 Aten işlemini destekler; ayrıca JAX/Pallas veya Helion DSL aracılığıyla özel çekirdekleri de destekler. En çok indirilen 5.000 Hugging Face modelinde kutudan çıktığı gibi doğrulandı; bunlar LLM'leri, ses ve konuşma, görüntü ve metin kodlayıcılarını kapsar. DTensor tabanlı parçalama, vLLM ve SGLang ile çok konaklı çıkarım, TorchTitan ile ince ayar ve eğitim, ayrıca Ray ve TorchRec ile entegrasyon da doğrulandı. TPU matris çarpım birimleri FP32 birikimi ile bfloat16 girdiler kullanırken, daraltma yapmayan işlemler vektör biriminde FP32'de çalışır. TorchTPU, PyTorch hassasiyet API'lerini bu donanıma eşler: medium (varsayılan) FP32 birikimi ile BF16 çarpımını kullanır; high yaklaşık 3 kat hesaplama maliyetiyle 3 geçişli bir algoritma kullanır; highest yaklaşık 6 kat ek maliyetle 6 geçiş kullanır. High veya highest seçmek bir performans uyarısını tetikler.

Kaynaklar: PyTorch · Google Cloud Blog · Hugging Face Blog