Gemma-4 (2B/4B/12B) AWS Inferentia2'ye Nasıl Taşındı
Gemma-4'ün AWS Inferentia2'ye taşınması: karışık dikkat kafaları, KV paylaşımı ve NxD sharding sorunlarının çözümü.
Google'ın Gemma-4 ailesini (E2B, E4B, 12B) AWS Inferentia2 üzerinde çalıştırmak, TPU/XLA için tasarlanmış üç mimari özelliğin—katmanlar arası KV paylaşımı, gruplu sorgu dikkatinde (GQA) karışık kafa sayıları ve kayan pencere/global dikkat katmanlarının farklı KV kafa sayılarıyla iç içe geçmesi—AWS'nin vendor yığınını (optimum-neuron, neuronx-distributed, Neuron vLLM backend) kırdığını ortaya koyuyor. Vendor yığınında Gemma-4 model sınıfı bulunmuyor ve graph builder KV paylaşımını ifade edemiyor; herkese açık Neuron vLLM endpoint'i akıcı görünen ama anlamsız çıktı üretiyordu.
Çözüm, NxD'nin statik ağırlık-tampon eşlemesiyle savaşmak yerine Hugging Face forward geçişini doğrudan izlemekten (trace) geçiyor. Bu sayede KV paylaşımı, izlenen grafikte doğal bir veri bağımlılığı olarak ortaya çıkıyor; yalnızca gerçekten KV önbelleği tutan katmanlar (E2B'de 15 katman) tampon ayırıyor. E4B'de tensor paralellik (TP=2) altında GQA'nın 8 sorgu/2 KV kafası, KV kafalarını TP'ye bölerken num_key_value_groups'u sabit tutmayı gerektiriyor; aksi halde repeat_kv yanlış eşleşiyor. 12B modelinde ise bölünemeyen tek KV kafalı global katmanlar replike bırakılıp grup sayısı küçültülerek çözülüyor.
Sonuçta üç model boyutu da CPU referansıyla token-token aynı greedy çıktı üretiyor: E2B ~44 tok/s, E4B ~33-39 tok/s (TP=2), 12B ~15 tok/s (TP=2). Bu, özel donanımlarda mimari uyumsuzlukları çözmek isteyen mühendisler için somut bir sharding ve tracing reçetesi sunuyor.