Google ve Meta, TorchTitan’ı TPU’ya taşıdı: PyTorch modelleri için TPU optimizasyon rehberi

Google, TorchTPU duyurusunun bir parçası olarak Meta’nın büyük ölçekli üretken yapay zekâ eğitimi için geliştirdiği referans platform TorchTitan‘ın Google Cloud TPU’larına uyarlanmış çatalını (fork) ve modelleri TPU’da optimize etmek için bir “oyun kitabı” yayımladı. Google Developers Blog’daki yazı Google mühendisleri ile Meta’dan Will Constable tarafından birlikte kaleme alındı. Çatal github.com/google-pytorch/torchtitan adresinde.

Neden önemli?

Google’a göre PyTorch modellerini TPU’da eğitmek isteyen geliştiriciler bugüne kadar ya kod tabanını tamamen JAX’e taşımak ya da eski torch_xla ile uğraşmak zorundaydı. TorchTPU ile TPU artık torch.device("tpu") olarak PyTorch’un çekirdek arayüzüne entegre. Ancak yazıya göre cihaz uyumluluğu işin yarısı; geliştiricilerin üretim düzeyinde dağıtık eğitim tariflerine de ihtiyacı var. TorchTitan, Llama 3 70B/405B ve Qwen 3 8B/32B gibi yoğun modeller, Qwen 3 235B gibi uzman karışımı (MoE) modeller ve FLUX.1 12B gibi difüzyon dönüştürücüleri için hazır tarifler sunuyor. Ayrıca VeRL gibi ek eğitim (post-training) hatlarında, pekiştirmeli öğrenme sırasında vLLM-TorchTPU ile ağırlık senkronize eden eğitici olarak kullanılabiliyor.

Temel teknik bileşenler

• Takılabilir mimari: Model, yapılandırma ağacından kuruluyor; bu sayede dikkat katmanı Pallas tabanlı Splash Attention ile, son projeksiyon ve çapraz entropi kaybı birleşik bir Pallas çekirdeğiyle, MoE uzman matris çarpımları gruplanmış çekirdeklerle model tanımı değiştirilmeden yer değiştirebiliyor. Her değişim bir yapılandırma bayrağı.
• SimpleFSDP: Parametre parçalamayı DTensor üzerinden ifade ederek kolektif işlemleri derlenen grafiğin içine taşıyor; böylece XLA iletişim ile hesaplamayı üst üste bindirebiliyor. Google, bunun FLUX’ta TorchTitan’ı JAX MaxDiffusion ile yüzde 99 performans denkliğine getirdiğini söylüyor.
• Graph Trainer: İleri geçiş, kayıp ve geri geçişi tek bir grafik olarak yakalayıp tüm eğitim adımını tek XLA programı olarak çalıştırıyor. Google’a göre bu yolla Qwen3 235B’de yüzde 29’un üzerinde, optimizer da grafiğe dahil edildiğinde Qwen3-8B’de yüzde 51-52 MFU elde ediliyor.

Adım adım optimizasyon

Rehber, modelleri TPU’da hızlandırmayı yeni kod yazmaktan çok mevcut bir tarif üzerinde sıralı yapılandırma değişiklikleri olarak tarif ediyor. İlk adım çalışma modunu seçmek: Eager mod modeli ayağa kaldırmak ve hata ayıklamak için uygun ama performans için değil; DeferAndFuse eager modu kod değişikliği olmadan bir miktar kazanç sağlıyor; blok başına torch.compile ise her dönüştürücü bloğunu tek XLA programına çeviriyor. Ardından SimpleFSDP ve karma hassasiyet politikasıyla parçalama, özel çekirdekler ve büyük modeller için katmanlar üzerinde “scan” ve aktivasyonları host belleğine boşaltma gibi teknikler geliyor.

TorchTPU tanıtım yazısında verilen örneğe göre bir FLUX modelinde adım süresi bu optimizasyonlarla 11,49 saniyeden 2,78 saniyeye indi. Google ayrıca TorchTitan’ın Ironwood (TPU v7x) üzerinde 1.024 çipte Llama 3 405B ile çip başına saniyede 500 token ve yüzde 50’nin üzerinde MFU’ya ulaştığını belirtiyor. Tüm performans rakamları Google’ın kendi ölçümleri.

Kaynak: Google Developers Blog: Integrating TorchTPU into TorchTitan and Optimizing Models for TPU · GitHub: google-pytorch/torchtitan · PyTorch Blog: Introducing TorchTPU