Avatar IV artık TPU üzerinde

HeyGen, 18 milyardan fazla parametreye sahip Avatar IV video üretim modelini Google Cloud'un Trillium (v6e) TPU'larına taşıdı. Taşıma işlemi torchax ve XLA araçlarıyla yapıldı; model sekiz çipli bir ağ üzerinde FSDP ve Ulysses dizi paralelliği kullanılarak çalıştırıldı.

Gerçek zamanlı akışta 1.86 kat hızlanma

Ekip, gerçek zamanlı akış için 1.86 kat hızlanma elde etti. Bunu sağlamak için all-to-all toplama işlemleri ardışık hale getirildi, seyrek dikkat blok boyutları maske dolgusunu ortadan kaldıracak şekilde hizalandı ve softmax bağımlılıkları önceden hesaplanan Cauchy-Schwarz üst sınırıyla aşıldı.

Kalite kontrolü

Bu özel Pallas çekirdeği ve derleyici optimizasyonları, ancak iki aşamalı sıkı kalite kontrolünden geçtikten sonra devreye alındı. Amaç, piksel çıktılarının bayt düzeyinde aynı ya da matematiksel olarak eşdeğer kalmasını garanti etmekti.