Google və Meta, Google TPU-ları üçün nativ PyTorch backend-i olan TorchTPU-nu təqdim edir

TorchTPU, Google və Meta arasındakı açıq mənbəli əməkdaşlıqdır və Google TPU-ları üçün nativ PyTorch backend-i təqdim edir. O, eager icranın optimallaşdırılması, torch.compile inteqrasiyası və nativ aparat sürətləndirməsi vasitəsilə sənayedə rəqabətqabiliyyətli inferensiya və təlim performansını hədəfləyir və modellərin əlavə düzəliş olmadan işləməsini və TPU-ya köçürülməsini mümkün edir.

Google TPU-ları böyük ML modellərini geniş miqyasda təlim etmək və yerləşdirmək üçün istifadə olunub. TPU-lar tarixi olaraq TensorFlow və JAX ilə bağlı olub, lakin bu proses 2025-ci ilin ortalarında Google-da Meta-nın PyTorch komandası ilə başladı. 2020-ci ildə təqdim edilən və Lazy Tensors-a əsaslanan əvvəlki PyTorch/XLA inteqrasiyası mənbə kodunda ciddi müdaxilə tələb edən dəyişikliklərə ehtiyac duyurdu, kompilyasiya müddətində ağır gecikmələrə məruz qalırdı, eager icra dəstəyi yox idi və paylanmış primitivlərdə geri qalırdı.

Dizayn prinsipləri və istifadə asanlığı

TorchTPU, JAX-ın uyğunlaşdırılması deyil, PyTorch üçün prinsipial TPU backend-i kimi qurulub. Onun prinsiplərinə PyTorch üzvlüyü (runtime daxilində JAX asılılığının olmaması, güclü eager rejimi, Torch Dynamo-ya əsaslanan torch.compile, həmçinin DDP, FSDP və DTensor dəstəyi), ekosistemlə problemsiz inteqrasiya və icma ilə əməkdaşlıq daxildir. İstifadəçilərə verilən diqqət portlama asanlığıdır: model.to("cuda") ifadəsini model.to("tpu") ilə əvəz etdikdə model əlavə düzəliş olmadan işləməlidir, lakin pik performans üçün əlavə tuning lazım ola bilər. Layihə həmçinin forma dəyişiklikləri nəticəsində XLA-nın yenidən kompilyasiyasının qarşısını almaq üçün məhdud dinamizmi hədəfləyir və əvvəlcədən kompilyasiya edilə bilən dinamik TPU kernel-lərini araşdırır. TorchTPU wheel paketi quraşdırılmış TPU maşınında PyTorch-un idxalı yalnız PyTorch-un autoload mexanizmi vasitəsilə kiçik dəyişikliklər tələb edir.

Eager icra və kompilyasiya

TorchTPU, PyTorch dispatcherini PrivateUse1 olaraq qeydiyyatdan keçmiş C++ backend-i ilə genişləndirir və TPU tensor saxlamasını standart at::Tensor mexanizmi vasitəsilə təyin edir. O, üç eager rejimi təqdim edir: Strict (standart) hər dəfə bir Aten əməliyyatını asinxron göndərir; Debug (TORCH_TPU_LAUNCH_BLOCKING=1) əlavə sərhəd yoxlamaları ilə sinxron rejimdədir; DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1) isə bir neçə əməliyyatı təxirə salır və birləşdirir, onun Strict eager-dən 50% və ya daha çox sürətli olduğu və pik HBM yaddaşını azaltdığı bildirilir. .to("cpu") çağırışı kimi materializasiya nöqtəsində TorchTPU liveness təhlili tətbiq edir, əməliyyat qrafı üçün keş açarı qurur və keşdə saxlanan icra faylını göndərir və ya asinxron kompilyasiya edir. Repeated Ops Heuristic təkrarlanan ardıcıllıqları, məsələn transformer blok yığımlarını aşkar edir və materializasiyanı erkən işə salır. Üç səviyyəli kompilyasiya keşi proses-lokal yaddaşı, maşın-lokal paylaşılan yaddaşı və hostlar ilə yenidən başlatmalar arasında yenidən istifadə üçün daimi uzaq fayl sistemi keşini əhatə edir.

Dəstəklənən iş yükləri və dəqiqlik

TorchTPU təxminən 450 Aten əməliyyatını dəstəkləyir; buna əsas memarlıqlar üçün tələb olunan Core ATen əməliyyatlarının irəli və geri keçidləri, həmçinin JAX/Pallas və ya Helion DSL vasitəsilə xüsusi kernel-lər daxildir. O, Hugging Face-in ən çox endirilən 5 000 modelində əlavə düzəliş olmadan yoxlanılıb; buna LLM-lər, audio və nitq, görüntü və mətn kodlayıcıları daxildir. DTensor-ə əsaslanan sharding, vLLM və SGLang ilə çox hostlu inferensiya, TorchTitan ilə fine-tuning və təlim, həmçinin Ray və TorchRec ilə inteqrasiya da yoxlanılıb. TPU matris vurma vahidləri FP32 yığılma ilə bfloat16 girişlərdən istifadə edir, kontraksiya etməyən əməliyyatlar isə vektor vahidində FP32-də işləyir. TorchTPU PyTorch dəqiqlik API-lərini bu aparata uyğunlaşdırır: medium (standart) FP32 yığılma ilə BF16 vurma istifadə edir; high təxminən 3 qat hesablama xərci ilə 3 keçidli alqoritmdən istifadə edir; highest isə təxminən 6 qat əlavə yüklə 6 keçid istifadə edir. high və ya highest seçildikdə performans xəbərdarlığı işə düşür.

Mənbələr: PyTorch · Google Cloud Blog · Hugging Face Blog