Colab TPU Üzerinde Saf JAX ile Gemma 4 E2B: Google'ın 4-Bit Dışa Aktarımı ve Kusursuz Paketleme (Repack) Karşılaştırması

Bu rehber, tek bir TPU v5e çipinde saf JAX (pure-JAX) motoru kullanarak Gemma 4 E2B modelini çalıştırma adımlarını ele almakta; Google tarafından eğitilen orijinal ağırlıklar ile aynı modelin iki farklı 4-bit derlemesini karşılaştırmaktadır. Çalışmadaki tüm metrikler Colab v5e-1 çalışma zamanında ölçülmüştür.

Google, E2B modelini 4-bit olarak gemma-4-E2B-it-qat-w4a16-ct adıyla yayınlamaktadır. Ancak bu resmi dışa aktarım (export), her ağırlığı modelin hiç eğitilmediği bir ızgaraya (grid) ikinci kez yuvarlamaktadır. Buna karşın, doğrudan eğitilmiş ızgara değerlerini saklayan repack (yeniden paketleme) yaklaşımı:

Bir sonraki token tahminlerinde orijinal modele 342,6 kat daha yakın kalır.
En yüksek olasılıklı ilk token eşleşmesinde %85,96 yerine %99,29 benzerlik sağlar.
128 token'lık bir hikayeyi orijinal modelle kelimesi kelimesine birebir aynı yazar.
Aynı çıkarım hızında çalışır ve 0,8 GB daha az indirme boyutu sunar.

Maraton (The Marathon)#

Bu çalışma, TPU arka ucunda JAX kullanımını pratikleştiren açık kaynaklı (Apache 2.0 lisanslı) AI GDE Marathon: JAX on TPU Tutorial serisinin bir parçasıdır. Colab'in sunduğu TPU ortamları ölçümlerin alındığı v5e donanımıyla aynı olduğundan, geliştiriciler kendi çiplerinde sonuçları birebir doğrulayabilir.

Neden Aynı Modelin İki Farklı 4-Bit Sürümü Var?#

Gemma 4 E2B, kuantizasyon farkındalıklı eğitim (QAT) ile eğitilmiştir; yani her 32 değerlik grup için belirlenen ölçekle 4-bit ızgarada öğrenim görmüştür. Google bu sonucu iki farklı biçimde sunar:

QAT Referansı (google/gemma-4-E2B-it-qat-q4_0-unquantized): Eğitilen ızgara değerlerini doğrudan bf16 formatında saklar.
Google Dışa Aktarımı (google/gemma-4-E2B-it-qat-w4a16-ct): int4 formatındadır; her grup en büyük ağırlık ÷ 7.5 formülüyle yeniden yuvarlanmıştır.
Repack Derlemesi (xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct): int4 formatındadır; doğrudan eğitilmiş seviyeleri ve adımları saklar.

Her iki 4-bit yapı da aynı W4A16 formatını ve tensör boyutlarını kullandığından aynı yükleyicilerle çalışır; temel fark barındırdıkları sayısal hassasiyettir.

Başlamadan Önce Gerekenler#

TPU v5e çalışma ortamına erişimi olan bir Google hesabı.
Yaklaşık 40 dakika süre (büyük bölümü 26 GB indirme ve derleme aşamasıdır).
Hugging Face erişim belirteci gerekmez; tüm kontrol noktaları ve motor kodları herkese açıktır.

Adım 1 — Notebook'u TPU Üzerinde Açma#

Colab üzerinden Çalışma Zamanı → Çalışma zamanı türünü değiştir → v5e-1 TPU seçimi yapılarak tüm hücreler çalıştırılır. Hücre 1.1 donanımı (TPU v5 lite, 16.91 GB HBM, 50.5 GB RAM) kaydeder. Hücre 1.2 ise PyTorch bağımlılığı içermeyen saf JAX motoru tpu-jaxı indirir. Motor, ağırlıkları standart XLA operasyonlarıyla açarak çarpar.

Adım 2 — Üç Kontrol Noktasını İndirme#

Hücre 2.2 tüm tensörleri indeksler. Google'ın dışa aktarımında embed_tokens ile bayt bayt aynı olan lm_head.weight tensörü fazladan saklanırken (0,805 GB), repack sürümünde bu kopya elenerek diskten ve bellekten tasarruf sağlanır.

Pratik Geliştirici Çıkarımı: Üretimde kuantize model kullanırken varsayılan export dosyalarına körü körüne güvenmek yerine, QAT ızgara matematiğini korumak ve tekrarlanan tensörleri ayıklamak hem model kalitesini hem de altyapı verimliliğini doğrudan artırır.

Bu konuyu derinlemesine öğrenmek isterseniz: Ölçeklendirme ve Optimizasyon modülüne göz atın.