Google, TPU hızlandırıcıları için yerel bir PyTorch backend'i olan TorchTPU'yu açık kaynak yaptı

Google, TPU donanımı için TorchTPU'yu açık kaynak yaptı

Google, Tensor Processing Unit'ları için PyTorch backend'i olan TorchTPU'yu, GitHub üzerindeki google-pytorch organizasyonu bünyesinde açık kaynak bir proje olarak yayımladı. 20 Ekim'de Hacker News'in ana sayfasında öne çıkan sürüm, PyPI üzerinden 0.2.0 sürümü olarak kurulabilir durumda ve Apache 2.0 lisansı taşıyor.

Projenin README'sine göre TorchTPU, PyTorch'un Python API'sinin altında yer alan operatör uygulamaları olan özel ATen kernel'ları sunuyor; böylece PyTorch modelleri ve tekil işlemler TPU cihazlarında çalıştırılabiliyor. Bir derleme önbelleği, backend'in ürettiği XLA hesaplamalarının optimize edilmiş biçimlerini saklıyor ve proje, graph-mode yürütme için torch.compile() ile entegre oluyor.

Kutuda neler var

Kernel'ların ötesinde, depo hem veri paralelliği hem de tensor paralelliği kullanan dağıtık eğitim desteğini, ayrıca Llama, Qwen, ResNet, DLRM, minGPT ve UNet'i kapsayan işlenmiş örnekleri listeliyor. Hata ayıklama ve kıyaslama araçları da dahil edilmiş.

Dokümantasyon, aynı zamanda bir dizi etkileşimli defter işlevi görüyor. Kullanıcı kılavuzu, Jupyter benzeri bir araç olan Marimo defterlerinin statik bir render'ı olarak Read the Docs üzerinde barındırılıyor ve README, bu defterlerin tarayıcıda marimo edit ile düzenlenip çalıştırılabileceğini belirtiyor. Bunları çalıştırmak için bir TPU VM gerekiyor ve en basit seçenek olarak bir v6e Google Compute Engine instance'ı gösteriliyor.

Kurulum ve derleme

Proje, Python 3.11 ve üzerini destekliyor; bu, PyTorch'un sürüm uyumluluk matrisiyle örtüşüyor. Kurulum pip üzerinden yapılıyor ancak README, PyTorch'un CPU wheel index'ini ek bir index URL'si olarak geçmeniz gerektiği konusunda uyarıyor; aksi halde pip, torch'un varsayılan CUDA build'ini çekiyor ve proje kullanıcıları bunun yerine CPU build'ine yönlendiriyor.

Kaynaktan derleme, Bazelisk wrapper'ı aracılığıyla kurulan Bazel üzerinden yürüyor ve README birkaç dikkat gerektiren noktayı belgeliyor. Release wheel'ları her zaman optimize modda derleniyor, gerekli bir wheel yapılandırma bayrağı olmadan derlemeler doğrudan başarısız oluyor ve Remote Build Execution kimlik bilgilerine sahip olmayan geliştiricilerin bu ayarları geri sökmek için ek bir bayrak eklemesi gerekiyor. Katkıda bulunanlar memnuniyetle karşılanıyor; ancak değişikliklerin devralınmadan önce Google'ın Contributor License Agreement'ını imzalamaları ve code review sürecinden geçmeleri şart.

Neden önemli

PyTorch, makine öğrenmesi araştırmacıları ve mühendislerinin büyük bir bölümünün varsayılan çalışma ortamı; TPU'lar ise tarihsel olarak Google'ın kendi framework yığınından en rahat erişilebilen donanımlar olageldi. ATen katmanında PyTorch'a entegre olan ve torch.compile() ile çalışan, amaca yönelik tasarlanmış açık kaynak bir backend, mevcut PyTorch eğitim ve çıkarım iş yüklerini Google'ın hızlandırıcılarına taşımanın sürtünmesini, onları baştan yazma ihtiyacını ortadan kaldırarak düşürüyor.

Apache 2.0 lisansı, Llama ve Qwen gibi güncel nesil modeller için herkese açık örnekler ve net bir katkı süreci, Google'ın bunu yalnızca dahili bir araç olarak değil, dış geliştiricilerin bu yol üzerine inşa etmesini istediğine işaret ediyor. 0.2.0 sürümündeki proje henüz erken aşamada; ancak hızlandırıcı ekosistemleri arasındaki rekabetin kızıştığı bir dönemde geliyor ve framework erişilebilirliği, iş yüklerinin fiilen nerede çalışacağını belirleyen kollardan biri.