Yeniden paketlenmiş QAT Gemma 4 26B-A4B, tek bir TPU v6e'de FP8 işlemenin 1.9 katı verimle sunuluyor

Ne oldu

Google'ın quantization-aware training (QAT) ile eğitilmiş Gemma 4 26B-A4B modeli artık vLLM altında tek bir Google Cloud TPU v6e çipinde çalışıyor; 17.43 GiB HBM kullanıyor, 53.888 token KV cache tutuyor ve saniyede 1.283 çıktı tokeni üretiyor. Derlemeyi belgeleyen bir dev.to yazısına göre, önceki tek çip seçeneği olan RedHat'in FP8 checkpoint'i 27.99 GiB tüketiyor, 3.456 token cache tutuyor ve saniyede 668 token sunuyor. Bu, 15.6 kat daha fazla KV cache ve 1.9 kat daha yüksek verim demek; ayrıca 3.880 kayıtlık bir sınıflandırma test setinde iki sürüm birbirinin bir puanı içinde puan alıyor.

26B'deki boşluk

Gemma 4 26B-A4B, 25.8 milyar parametreli, katman başına 128 expert'e sahip ve token başına 8 expert'in aktif olduğu bir mixture-of-experts modeli. bf16'da 48.07 GiB gerektiriyor, tek bir v6e çipi ise 28.74 GiB kullanılabilir HBM sunuyor; yani nicelenmemiş halde sığamıyor. Google, her Gemma 4 boyutu için 4-bit QAT varyantları eğitti, ancak dev.to yazısı vLLM kullanıcıları için 26B'de bir boşluk tespit ediyor: model kartı, llama.cpp için hedeflenmiş, vLLM'in yükleyemediği bir GGUF Q4_0 export'u; "nicelenmemiş" bir bf16 QAT checkpoint'i; ve vLLM için tasarlanmış format olan Compressed-Tensors W4A16 checkpoint'leri listeliyor — ki bunlar yalnızca E2B, E4B, 12B ve 31B boyutları için yayımlandı. Geriye kalan tek yol, 0.75 GiB boşlukla sığan üçüncü taraf FP8 derlemesiydi — bu da yaklaşık bir buçuk 2.048 tokenlik istek için cache demek.

Yeniden nicelendirme değil, yeniden paketleme

Yazar, boşluğu Google'ın "nicelenmemiş" QAT export'unu yeniden paketleyerek doldurdu. Tensörleri bf16 olarak saklanıyor, ancak inceleme, giriş boyunca 32'lik her ağırlık grubunun zaten −8 ile 7 arasındaki seviyelerle 16 seviyeli bir Q4_0 ızgarası üzerinde olduğunu gösterdi; bu, quantization-aware training'in kalıntısı. Bu nedenle 4-bit checkpoint üretmek, kayıplı bir dönüşüm değil, yalnızca bir container değişikliğidir.

Tek tuzak grup başına scale değeri. Ders kitabı kuralı olan "grubun maksimum mutlak ağırlığını 8'e böl", tepe değerin ±8 seviyesinde olduğunu varsayar. Tepe daha düşükteyse kural yanlış adım türetir ve tüm grubu orijinal değerleri içermeyen bir ızgaraya yeniden yuvarlar — grup başına yaklaşık %5'lik medyan hata, buna karşın tüm şekil doğrulamaları hâlâ geçer. Yeniden paketleme scripti bunun yerine maksimumu 1'den 8'e kadar her m değerine bölmeyi dener, tüm grubu yeniden üreten ilk adımı korur, sonra bu adımı 32 değer üzerinde en küçük kareler yöntemiyle iyileştirir.

Ortaya çıkan checkpoint, Hugging Face'te xbill9 hesabı altında yayımlandı. Attention, yoğun (dense) MLP ve 3.840 expert'in tamamı compressed-tensors formatında, 32 grup boyutuyla int4'e nicelendirildi; router, embedding'ler, normlar ve vision tower değiştirilmeden kopyalandı. Doğrulamada ızgaradan çıkan sıfır grup bulundu; expert değerlerinin %92,6'sı kaynakla bit düzeyinde aynı, kalan farklar yalnızca scale'i bf16'da saklanmasından kaynaklanıyor — vLLM'in TPU W4A16 katmanlarının bunu yüklediği hassasiyet bu — ve en kötü durumda göreli fark 1.1e-2.

vLLM'in TPU backend'inde değişen ne oldu

TPU'da sunum, vLLM'in JAX tabanlı TPU backend'i olan tpu-inference içinde iki parça gerektirdi: linear katmanlar için W4A16 ekleyen, halihazırda onaylanmış bir pull request (#3653) ve expert'ler için yeni bir WNA16FusedMoEMethod (#3660). Expert yöntemi, paketlenmiş int4 ağırlıklarını ve 32 genişliğindeki scale'leri kendisi yükler; çünkü varsayılan yol, çıkış kanalı başına tek scale ile yeniden nicelendirme yapar ve eğitilmiş ızgarayı yok eder. Ağırlıkları, her 32 genişliğindeki grup için bir scale kabul eden ve çarpma öncesi her ağırlık parçasını hızlı çip içi bellekte ters nicelendiren GMM backend'ine yönlendirir. Aktivasyonlar bf16 olarak kalıyor; dolayısıyla nicelendirme yalnızca ağırlik tarafında (weight-only) sürüyor.

Checkpoint ayrıca taşınabilir: yerleşim, vLLM'in int4 mixture-of-experts checkpoint'leri için zaten beklediği biçimle eşleştiğinden, aynı dosya yamalanmamış vLLM 0.30.0 üzerinde bir NVIDIA L4'te de yükleniyor.

Neden önemli

Öne çıkan konu serbest kalan bellek. Ağırlık depolamasının 27.99 GiB'den 17.43 GiB'ye düşmesi doğrudan KV cache'e dönüşüyor — 15.6 kat fazlası — ve sunum yaparken bağlayıcı kısıt, context uzunluğu ve batch boyutu açısından cache'tir. 1.9 kat verim ve test setinde FP8'in bir puanı içinde doğrulukla, 26 milyar parametrelik bir model tek bir orta segment hızlandırıcıda pratik hale geliyor. Bu derleme ayrıca, eğitilmiş ızgaraya saygı gösterildiği sürece QAT export'larının neredeyse kayıpsız biçimde container'lar arasında taşınabileceğini gösteriyor ve ortaya çıkan artifact, TPU donanımının ötesinde standart vLLM üzerinde sunulabiliyor. Her script, log ve kayıt başına çıktı, eşlik eden depoya işlendi; böylece iddialar uçtan uca yeniden üretilebilir.