Bir Kıdemli ML Mühendisinin JAX Günlükleri: İlk İzlenimler ve PyTorch Karşılaştırması

Bir Kıdemli ML Mühendisinin JAX Günlükleri: İlk İzlenimler ve PyTorch Karşılaştırması

Yıllardır PyTorch kullanan bir kıdemli makine öğrenmesi mühendisi olarak JAX dünyasına adım attım. İşte performans, fonksiyonel programlama ve ilk izlenimlerim.

Yapay zeka ve makine öğrenmesi dünyası baş döndürücü bir hızla değişiyor. Sektörde uzun yıllar harcamış, TensorFlow’un yükselişini ve ardından PyTorch’un tahtı ele geçirişini yakından takip etmiş bir kıdemli mühendis olarak, yeni bir araca geçiş yapmak her zaman büyük bir soru işareti yaratır. Ancak son zamanlarda adını sıkça duyduğumuz, özellikle büyük dil modellerinin (LLM) arkasındaki devasa hesaplama güçlerinde tercih edilen JAX, artık göz ardı edilemeyecek bir noktaya ulaştı. Ben de bu trendi yakından incelemek ve kendi projelerimde ne kadar efektif olabileceğini görmek adına JAX öğrenmeye karar verdim. Bu makalede, bir kıdemli mühendis gözünden JAX ile geçirdiğim ilk haftaların teknik analizini, yaşadığım zihniyet değişimini ve gerçek dünya senaryolarındaki performans gözlemlerimi aktaracağım.

JAX Nedir ve Neden Birdenbire Sektörün Göz Bebeği Oldu?

JAX, Google tarafından geliştirilen ve temelde yüksek performanslı sayısal hesaplamalar yapmayı sağlayan bir kütüphanedir. Ancak onu sadece bir makine öğrenmesi kütüphanesi olarak tanımlamak haksızlık olur. JAX, bildiğimiz NumPy API’sini birebir taklit ederken, arka planda bu işlemleri GPU (Grafik İşlem Birimi) ve TPU (Tensor İşlem Birimi) üzerinde çalıştırabilen, otomatik türev alabilen ve kodları XLA (Accelerated Linear Algebra – Hızlandırılmış Doğrusal Cebir) derleyicisi ile optimize edebilen bir ekosistemdir.

Özellikle derin öğrenme dünyasında PyTorch uzun süredir mutlak hakimiyete sahip. Buna rağmen, büyük araştırma laboratuvarları ve Google DeepMind gibi devler neden hızla JAX’e yöneliyor? Bunun en temel sebebi, JAX’in sunduğu benzersiz dönüşüm (transformation) yetenekleridir. JAX bizlere sadece tensor işlemleri sunmaz; yazdığımız sıradan Python fonksiyonlarını alıp bunları paralel hale getirebilen, türevini çıkarabilen ve donanıma göre optimize edebilen fonksiyonel bir altyapı sağlar. Dolayısıyla, donanım sınırlarını zorlayan devasa modeller eğitirken JAX, PyTorch’a kıyasla çok daha düşük seviyeli kontrol ve muazzam bir hız avantajı sunmaktadır.

Bununla birlikte, JAX öğrenmek sadece yeni bir Framework (yazılım çerçevesi) öğrenmek anlamına gelmiyor. JAX, yazılım geliştirme tarzınızı ve algoritmaları ele alış biçiminizi kökten değiştiren bir felsefeye sahiptir. Bu felsefeyi anlamak, JAX ile başarılı projeler geliştirmenin ilk ve en önemli kuralıdır.

PyTorch’tan JAX’e Geçişte Karşılaştığım Zihniyet Değişimi

PyTorch kullananlar iyi bilir; orada her şey nesne yönelimli programlama (OOP) ve dinamik grafikler üzerine kuruludur. Bir model tanımlamak için nn.Module sınıfından türetilmiş bir sınıf yazarız, bu sınıfın içinde durumları (state) yani ağırlıkları (weights) ve bias değerlerini saklarız. Ancak JAX dünyasına adım attığınızda bu alışkanlıklarınızı kapıda bırakmanız gerekiyor. Çünkü JAX tamamen fonksiyonel programlama (functional programming) paradigmasını benimser.

Fonksiyonel programlamanın JAX’teki en somut yansıması, “Pure Functions” (saf fonksiyonlar) kavramıdır. JAX üzerindeki dönüşüm araçlarının (JIT, Grad vb.) doğru çalışabilmesi için yazdığınız fonksiyonların dış dünyaya bağımlı olmaması ve yan etki (side effect) üretmemesi gerekir. Yani bir fonksiyon, aynı girdilerle çağrıldığında her zaman aynı çıktıyı vermeli ve fonksiyon dışındaki hiçbir değişkeni veya durumu değiştirmemelidir. Örneğin, PyTorch’ta alışık olduğumuz model içi parametre güncellemeleri JAX’te doğrudan yapılamaz. Parametreler fonksiyona açıkça (explicitly) parametre olarak geçilmeli ve güncellenmiş parametreler fonksiyondan geri döndürülmelidir.

Bu durum ilk başlarda bana oldukça zahmetli ve gereksiz bir kod kalabalığı gibi geldi. Özellikle büyük modellerde binlerce parametreyi sürekli fonksiyonlar arasında taşımak korkutucu görünüyordu. Fakat bu yapının sunduğu determinizm ve paralel programlama kolaylığı, sistem ölçeklendikçe ne kadar büyük bir nimet olduğunu kanıtlıyor. Değiştirilemez (immutable) veri yapıları sayesinde, kodun neresinde hata oluştuğunu takip etmek ve eşzamanlı işlemleri yönetmek çok daha güvenli hale geliyor.

JAX’in Dört Atlısı: JIT, Grad, Vmap ve Pmap Nasıl Çalışır?

JAX’i güçlü kılan ve onu standart kütüphanelerden ayıran dört temel işlevsel dönüşüm mekanizması vardır. Bu mekanizmaları anlamak, JAX’in gücünden tam anlamıyla yararlanabilmek için şarttır. Gelin bu dört kavramı yakından inceleyelim:

• JIT (Just-In-Time Compilation): Yazdığınız Python kodunu alır ve XLA derleyicisi yardımıyla optimize edilmiş makine koduna dönüştürür. Bu sayede Python’ın getirdiği yavaşlık ortadan kalkar ve kod doğrudan donanım üzerinde en yüksek hızda çalışır.
• Grad (Otomatik Türev): Matematiksel fonksiyonların türevini otomatik olarak hesaplar. PyTorch’taki backward() mekanizmasının aksine, JAX’te türev alma işlemi doğrudan bir fonksiyon dönüşümüdür. Bir fonksiyonun türevini alan başka bir fonksiyon üretebilirsiniz.
• Vmap (Vectorized Map): Tek bir veri örneği için yazdığınız fonksiyonu, hiçbir döngü (loop) kullanmadan otomatik olarak bir veri yığını (batch) üzerinde çalışacak şekilde vektörize eder. Bu, veri işlemeyi inanılmaz derecede kolaylaştırır.
• Pmap (Parallel Map): Yazdığınız kodun birden fazla GPU veya TPU üzerinde paralel olarak, minimum çabayla çalıştırılmasını sağlar. Büyük modelleri birden fazla cihaza dağıtmak için kullanılır.

Bu dönüşümlerin nasıl harika bir uyum içinde çalıştığını görmek için basit bir kod örneği inceleyelim. Aşağıdaki örnekte, basit bir aktivasyon fonksiyonunun türevini alıp, bunu JIT ile nasıl hızlandırdığımızı görebilirsiniz:

import jax
import jax.numpy as jnp

# Basit bir aktivasyon fonksiyonu tanımlayalım (Saf Fonksiyon)
def activation_fn(x):
    return jnp.tanh(x)

# Fonksiyonun türevini alan yeni bir fonksiyon üretiyoruz
grad_activation = jax.grad(activation_fn)

# Şimdi bu türev fonksiyonunu JIT ile derleyerek hızlandırıyoruz
fast_grad_activation = jax.jit(grad_activation)

# Test edelim
value = 1.0
print("Türev Değeri:", fast_grad_activation(value))

Yukarıdaki kod bloğunda dikkatinizi çekmek istediğim nokta, jax.grad fonksiyonunun parametre olarak doğrudan başka bir fonksiyonu almasıdır. Bu yaklaşım, matematiksel formülleri koda dökerken inanılmaz bir esneklik ve okunabilirlik sağlar.

Adım Adım İlk JAX Modelimizi Yazalım

Şimdi teorik bilgileri bir kenara bırakıp, JAX ile basit bir doğrusal regresyon (linear regression) modeli geliştirelim. Bu süreçte JAX’in parametre yönetimini ve eğitim döngüsünü nasıl ele aldığını pratik olarak göreceğiz. PyTorch’taki gibi hazır sınıflar yerine, her şeyi sıfırdan ve fonksiyonel olarak inşa edeceğiz.

Öncelikle rastgele veri üretmemiz gerekiyor. JAX, rastgele sayı üretimi (PRNG) konusunda da PyTorch ve NumPy’dan ayrılır. JAX’te rastgelelik de saf fonksiyon prensibine uymalıdır. Bu yüzden her rastgele işlem için açıkça bir anahtar (key) belirtmemiz gerekir. Bu durum, çalışmalarımızın tamamen tekrarlanabilir (reproducible) olmasını garantiler.

# Rastgele sayı üretimi için anahtar oluşturma
key = jax.random.PRNGKey(0)
key, subkey1, subkey2 = jax.random.split(key, 3)

# Sentetik veri üretimi (y = 3x + 2 + gürültü)
X = jax.random.normal(subkey1, (100, 1))
noise = jax.random.normal(subkey2, (100, 1)) * 0.1
Y = 3.0 * X + 2.0 + noise

# Model parametrelerini başlatalım (Ağırlık ve Sapma)
weights = jax.random.normal(subkey1, (1, 1))
bias = jnp.zeros((1,))
params = {'w': weights, 'b': bias}

Verilerimizi ve başlangıç parametrelerimizi hazırladıktan sonra, tahmin yürütecek, kaybı (loss) hesaplayacak ve parametreleri güncelleyecek fonksiyonlarımızı yazabiliriz. Unutmayın, tüm bu fonksiyonların saf fonksiyon olması gerekiyor.

# İleri besleme (Forward Pass) fonksiyonu
def predict(params, x):
    return jnp.dot(x, params['w']) + params['b']

# Kayıp (Loss) fonksiyonu - Ortalama Kare Hata
def loss_fn(params, x, y):
    preds = predict(params, x)
    return jnp.mean((preds - y) ** 2)

# Güncelleme (Update) fonksiyonu
@jax.jit
def update(params, x, y, lr=0.1):
    # Kayıp fonksiyonunun parametrelere göre gradyanını alıyoruz
    grads = jax.grad(loss_fn)(params, x, y)
    
    # Parametreleri güncelliyoruz (Gradyan İnişi)
    new_params = {
        'w': params['w'] - lr * grads['w'],
        'b': params['b'] - lr * grads['b']
    }
    return new_params

Artık eğitim döngümüzü başlatabiliriz. Yazdığımız update fonksiyonu JIT ile derlendiği için son derece hızlı çalışacaktır. Birkaç adımda modelimizin hedef parametrelere (w=3, b=2) nasıl yaklaştığını görebiliriz:

# Eğitim döngüsü
for epoch in range(100):
    params = update(params, X, Y)
    if epoch % 20 == 0:
        current_loss = loss_fn(params, X, Y)
        print(f"Epoch {epoch} - Kayıp: {current_loss:.4f} | w: {params['w'][0,0]:.2f}, b: {params['b'][0]:.2f}")

Bu basit örnekte bile JAX’in zarafetini görmek mümkün. Sınıflarla, karmaşık nesne durumlarıyla uğraşmadan, tamamen matematiksel ifadeleri koda dökerek modelimizi eğittik. Özellikle büyük projelerde bu sadelik, hata yapma olasılığını ciddi oranda azaltmaktadır.

Gerçek Dünya Senaryosu: Büyük Veri Setlerinde Performans Analizi

Peki, JAX gerçek dünyada bize ne kazandırıyor? Türkiye’deki büyük bir e-ticaret platformunun kullanıcı etkileşim verilerini işleyen bir öneri sistemi (recommendation system) senaryosunu ele alalım. Günlük milyonlarca satır etkileşim verisi üzerinde matris ayrıştırma (matrix factorization) veya büyük boyutlu vektör çarpımları yapmamız gerektiğini düşünelim.

Bu tür senaryolarda standart NumPy veya optimize edilmemiş PyTorch kodları bellek (memory) ve zaman darboğazına yol açar. JAX’in vmap ve jit özellikleri tam olarak bu noktada devreye giriyor. Yaptığım testlerde, büyük boyutlu matris çarpımları ve gradyan hesaplamalarında JAX’in sunduğu performans artışını aşağıdaki gibi gözlemledim:

Tablodaki verilerden de anlaşılacağı üzere, JAX ile JIT derlemesi yapıldığında donanım kaynakları maksimum verimlilikle kullanılır. XLA derleyicisi, ardışık matris işlemlerini birleştirerek (kernel fusion) gereksiz bellek transferlerini önler. Bu durum, özellikle bulut sağlayıcılarda (AWS, Google Cloud vb.) yüksek GPU maliyetleriyle çalışan şirketler için doğrudan operasyonel maliyetlerin düşmesi anlamına gelmektedir.

Kıdemli Bir Mühendis Gözünden JAX’in Eksileri ve Zorlukları

Şimdiye kadar JAX’in hep harika yönlerinden bahsettik. Ancak hiçbir teknoloji mükemmel değildir ve kıdemli bir mühendis olarak JAX’in de bazı can sıkıcı eksiklikleri olduğunu belirtmek zorundayım. Eğer üretim ortamında (production) çalışan sistemler geliştiriyorsanız, bu zorlukları bilerek yola çıkmalısınız.

İlk olarak, JAX’in hata ayıklama (debugging) süreci PyTorch kadar kolay değildir. JIT derleyicisi devreye girdiğinde, Python kodunuz optimize edilmiş bir grafiğe dönüştürülür. Bu nedenle, fonksiyonun ortasına standart bir print() veya pdb.set_trace() koyup değişkenlerin değerini anlık olarak inceleyemezsiniz. JAX’in kendine has jax.debug.print gibi araçlarını kullanmanız gerekir ki bu da öğrenme eğrisini dikleştiren bir unsurdur.

İkinci büyük sorun ise ekosistem olgunluğudur. PyTorch; Hugging Face, PyTorch Lightning, Torchvision gibi devasa ve olgun bir kütüphane ekosistemine sahiptir. JAX tarafında ise işler biraz daha dağınıktır. Google, model tanımlama işini Flax veya Haiku gibi üst düzey kütüphanelere bırakmıştır. Ancak bu kütüphaneler henüz PyTorch kadar standartlaşmış ve topluluk desteğine kavuşmuş değildir. Bir projeye başlarken hangi yan kütüphaneyi seçeceğinize karar vermek bile başlı başına bir araştırma konusu olabilmektedir.

Sıkça Sorulan Sorular (FAQ)

JAX, PyTorch’un yerini tamamen alabilir mi?

Kısa vadede hayır. PyTorch, özellikle endüstride ve hazır modellerin kullanımında (Hugging Face vb.) liderliğini koruyor. Ancak büyük ölçekli model eğitimi, akademik araştırmalar ve özel donanım optimizasyonlarında JAX’in pazar payı hızla artıyor. İki kütüphane de farklı ihtiyaçlara hizmet ediyor.

JAX öğrenmek kariyerim için faydalı mı?

Kesinlikle evet. Özellikle üretken yapay zeka (Generative AI) ve büyük dil modelleri üzerine çalışan şirketler, yüksek performanslı hesaplama yetenekleri nedeniyle JAX bilen mühendislere öncelik vermeye başladı. Ayrıca fonksiyonel programlama bakış açısı kazanmak genel yazılım kalitenizi de artıracaktır.

JAX kullanırken hangi kütüphaneleri tercih etmeliyim?

Eğer derin öğrenme modelleri geliştirecekseniz, Google’ın aktif olarak desteklediği ve geliştirdiği Flax kütüphanesini tercih etmeniz en güvenli yoldur. Optimizasyon işlemleri için ise Optax kütüphanesi sektör standardı haline gelmiştir.

JAX sadece TPU’larda mı çalışır?

Hayır, JAX hem CPU hem de NVIDIA ve AMD GPU’larında mükemmel bir şekilde çalışır. Ancak Google tarafından geliştirildiği için TPU entegrasyonu ve performansı diğer kütüphanelere kıyasla çok daha optimize durumdadır.

Özetlemek gerekirse, JAX öğrenmek benim için heyecan verici ve ufuk açıcı bir deneyim oldu. PyTorch’un sunduğu konfor alanından çıkıp JAX’in saf matematiksel dünyasına adım atmak, kod kalitemi ve donanım algımı bir üst seviyeye taşıdı. Eğer siz de sınırları zorlamayı seviyorsanız, JAX’e kesinlikle bir şans vermelisiniz.

#Teknoloji #MakineÖğrenmesi #YapayZeka #JAX #Python