Google’ın Tunix sistemi, büyük ölçekli ajan tabanlı pekiştirmeli öğrenmenin (RL) TPU'ları verimli kullanmasını engelleyen darboğazı ortadan kaldırıyor. Etkileşim verisi üretme işini politika güncelleme işinden ayırarak Tunix, TPU kullanımını tek haneli yüzdelerden tam kapasiteye yakın seviyelere çıkarıyor ve hesaplama israfını büyük ölçüde azaltıyor.

Ajan tabanlı RL'deki darboğaz

Ajan tabanlı RL, daha tanıdık olan "sonraki token" dil modeli eğitiminden farklıdır. Bir ajanın API çağrıları göndermesi, kod yürütmesi veya simüle edilmiş bir ortamda adım atması, ardından sonuca tepki vermesi gerekir. Bu nedenle eğitim döngüsü senkrondur: model bir eylem üretir, ortam çalışır, sonuç döner ve ancak ondan sonra model bir gradyan güncellemesi alır. Tek bir ortam adımı birkaç saniye sürdüğünde, maliyetli TPU donanımı boşta bekler ve bildirilen kullanım %10'un altına düşebilir. Bu verimsizlik doğrudan daha yüksek bulut faturalarına ve daha yavaş araştırma döngülerine dönüşür.

Tunix'in ayrıştırılmış mimarisi

Tunix, iki aşamayı —yörünge (trajectory) üretimi ve politika optimizasyonu— ayrı donanım havuzlarına çekerek sorunu çözer.

  • Asenkron aktörler ucuz CPU'lar veya GPU'lar üzerinde çalışır. Her aktör, atanan ortamla sürekli etkileşime girer, eylemleri ve gözlemleri kaydeder ve sonuçta oluşan yörüngeleri paylaşılan bir depoya aktarır.
  • Sürekli öğreniciler özel TPU Pod'larını işgal eder. Öğrenici, merkezi tampondan (buffer) gruplar (batches) çeker ve herhangi bir aktörün bir rollout işlemini bitirmesini beklemeden gradyan güncellemelerini gerçekleştirir.
  • Yüksek verimli tampon (buffer) ortada yer alır ve yörüngeler için bir hazırlık alanı görevi görür. Öğrenici, tamponun veri sağlama hızı kadar hızlı okuyabildiği için TPU asla duraksamaz.

Net etki, TPU'ların neredeyse her zaman meşgul kaldığı ve kullanımın %100'e yaklaştığı bir eğitim hattıdır.

Teknik engeller ve Tunix bunları nasıl aşıyor

Değişken uzunluklu bölümler ve XLA yeniden derleme

JAX'ın XLA derleyicisi sabit tensör şekilleri için optimize edilmiştir. Ancak ajan tabanlı görevler, normalde maliyetli yeniden derlemeleri tetikleyecek farklı uzunluklarda diziler üretir. Tunix, daha kısa dizileri bir araya getirir ve benzer uzunluktaki bölümleri gruplara (buckets) ayırarak şekillerin XLA'nın derlenmiş çekirdekleri (kernels) yeniden kullanabileceği kadar uzun süre sabit kalmasını sağlar. Sonuç, performansı felç edebilecek derleyici yükü olmadan istikrarlı bir verimliliktir.

Devasa modelleri birçok TPU çipine ölçeklendirme

70 milyardan fazla parametreye sahip ajanları eğitmek, ağırlıkların ve verilerin birden fazla TPU düğümüne yayılmasını gerektirir. Tunix, hem model parametrelerini hem de aktivasyonları bölmek (shard) için JAX'ın ShardMap primitifini kullanır; bu da öğrenicinin tüm modeli bellekte tutarken verileri yüksek hızda beslemesine olanak tanır. Bu bölme (sharding) stratejisi, daha önce tek bir TPU podu için ulaşılamaz olan modelleri eğitme imkanı sağlar.

Ayrıştırılmış boru hatlarından kaynaklanan bayat gradyanlar

Aktörler öğreniciden daha hızlı çalıştığında, sağladıkları veriler mevcut politikaya göre "bayat" (stale) hale gelebilir. Tunix bu sapmayı iki mekanizma ile hafifletir: önem örneklemesi (importance-sampling) eski örnekleri geçerliliklerini yansıtacak şekilde yeniden ağırlıklandırır ve yapılandırılabilir bir bayatlık eşiği, önceden belirlenmiş yaşı aşan yörüngeleri atar. Bunlar birlikte, boru hattı asenkron çalışırken bile öğrenmeyi kararlı tutar.

Kullanıcıların dikkat etmesi gerekenler

  • Gecikme denetimi – Ayrıştırmanın faydası ortamın yanıt süresine bağlıdır. Ekipler uçtan uca gecikmeyi ölçmeli ve aktör havuzlarının tamponu iyi dolu tutacak şekilde boyutlandırıldığından emin olmalıdır.
  • İşçi havuzu tasarımı – Ucuz CPU'lar veya GPU'lar birçok aktöre ev sahipliği yapabilir, ancak bunları aşırı yüklemek ağ veya depolama üzerinde kaynak rekabetine (contention) neden olabilir. Tamponun veri alma hızıyla eşleşen dengeli bir havuz esastır.
  • Tampon dayanıklılığı – Merkezi depo, yeni bir darboğaz haline gelmeden yüksek yazma ve okuma hızlarını yönetebilmelidir. Düşük kuyruk gecikmesine (tail latency) ve yeterli bant genişliğine sahip bir depolama sistemi seçmek, mimarinin tartışmaya kapalı bir parçasıdır.

Potansiyel dezavantajlar

Ayrıştırılmış mimari daha fazla hareketli parça getirir: ayrı donanım filoları, kalıcı bir tampon ve bayatlık sınırlarını uygulamak için koordinasyon mantığı.

Özet

Tunix, ajan tabanlı RL'deki baskın maliyetin modelin kendisi değil, senkron etkileşim döngülerinin neden olduğu boşta bekleme süresi olduğunu göstermektedir. Rollout işini ucuz donanıma devrederek ve sürekli öğrenen bir TPU podunu yüksek verimli bir tamponla besleyerek Google, %10'un altındaki bir kullanım sorununu tam kapasiteye yakın bir iş akışına dönüştürmüştür.