2 Saniyenin Altında RL Son Eğitimi için Ağırlık Aktarımı

Kimi-K2 (1T parametre) için yakın zamanda 1,3 saniyede makineler arası parametre güncellemeleri gerçekleştirdik; bu süreçte ağırlıklar 256 eğitim GPU'sundan (BF16) 128 çıkarım GPU'suna (FP8) aktarıldı.

Asenkron pekiştirmeli öğrenme ince ayarında, eğitim ve çıkarım ayrı GPU'larda çalışır. Her eğitim adımından sonra, yeni ağırlıkların çıkarım düğümlerine gönderilmesi gerekir. Birçok mevcut çerçevede trilyon parametreli modeller için bu işlem birkaç saniye, hatta dakika sürebilir.

RDMA noktadan noktaya iletişimden yararlanarak, çıkarım motorunu değiştirmeden ağırlık aktarımını son derece hızlı hale getirebiliyor ve kodun yazılmasını ve bakımını daha kolay hale getirebiliyoruz.

Çözümümüz, kaynağın doğrudan hedefin GPU belleğine yazma yaptığı tek taraflı bir temel işlev olan RDMA WRITE üzerine kurulmuştur.

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

Hedef taraf aktarım için bildirim bile almaz. Bu bize, çıkarım düğümlerinde herhangi bir kontrol mantığı olmadan eğitim düğümleri tarafından yönlendirilen düşük gecikmeli, yüksek bant genişliğine sahip, sıfır kopyalı aktarımlar sağlar.

Meta veri toplama – Denetleyici, tüm eğitim ve çıkarım GPU'larından parametre meta verilerini toplar.

Zamanlama hesaplaması – Denetleyici; hangi eğitim GPU'sunun hangi parametreyi hangi çıkarım GPU'suna ve hangi sırayla göndereceğini eşleyen statik bir ağırlık aktarım zamanlaması hesaplar.

Zamanlama dağıtımı – Denetleyici, zamanlamayı tüm eğitim GPU'larına gönderir.

Yürütme – Her eğitim adımından sonra denetleyici, eğitim GPU'larına aktarımları başlatmaları için sinyal gönderir.

Yüksek düzeyli iş akışı tanımlandıktan sonra, temel zorluk trilyon parametre ölçeğinde ağırlık aktarımlarının verimli bir şekilde nasıl yürütüleceğidir. Burada yürütme yolunun ayrıntılarını açıklıyoruz.

Eğitimdeki parametreler FSDP yerleşimlerine göre dağıtılır. full_tensor() kullanılarak bir DeviceMesh içindeki tüm GPU'lar tam parametreyi yeniden oluşturabilir, dolayısıyla tümü ağırlık aktarımı için kaynak olarak hizmet verebilir.

Birden fazla ayrık DeviceMesh bir mesh grubu oluşturur. Aynı gruptaki DeviceMesh'ler ayrık olduğundan, aktarımları birbirine müdahale etmez ve tamamen paralel çalışabilir. Sıralamayı zorunlu kılmak için mesh grupları arasına küresel bir bariyer ekliyoruz.

Ağırlık aktarım hattı, host’tan cihaza kopyaları, tensor birleştirmeyi, füzyonu ve kuantizasyonu, ayrıca MeshGroups genelinde RDMA transferlerini örtüştürür.

Her parametre tensörünün aktarımını bir görev olarak ele alıyoruz. Ağırlık aktarım süreci birden fazla donanım kaynağı türü kullanır, bu nedenle bir ağırlık aktarım görevini zaman içinde örtüşen farklı işlem hattı aşamalarına bölüyoruz:

Anahandandan cihaza memcpy — FSDP ağırlığı CPU'ya offload ederse

Parametre hazırlığı — full_tensor() ile tam ağırlığı yeniden oluşturun, projeksiyon birleştirmeyi uygulayın, gerekirse nicemleyin.

RDMA aktarımı — Uzak çıkarım GPU belleğine sıfır kopyalı yazma

Küresel bariyer — Tüm full_tensor() çağrıları tamamlandıktan sonra, Ethernet üzerinden GLOO kullanarak mesh grupları arasında senkronizasyon yapın.

Uygulamada, her işlem hattı aşaması için bir FIFO görev kuyruğu tutarız. Kuyruğun başındaki görev aşamayı tamamladığında, bir sonraki aşama kuyruğunun sonuna taşınır.

full_tensor() ve diğer GPU işlemler ek GPU bellek kullanımı ortaya çıkarır. Bellek yetersizliği (OOM) hatasından kaçınmak için, bir görevin yürütülmesini yalnızca devam eden anlık görevler yapılandırılabilir bir filigrandan daha az geçici GPU belleği kaplıyorsa başlatırız.

Çeşitli tasarım seçimleri, sistemimizin yaygın açık kaynaklı çözümlere göre çalıştırılmasını önemli ölçüde daha hızlı ve bakımını daha kolay hale getirir.

Rank0 tabanlı ağırlık aktarımı, ağırlıkları çıkarım düğümüne dağıtmadan önce toplar.

Noktadan noktaya iletişim, ağırlıkları her iki eğitim düğümünden çıkarım düğümünün ağ arayüzlerine aktarır.

Yaygın bir örüntü, tüm parametreleri sıfır numaralı rank GPU'lardan geçirmektir: eğitim rank-0 üzerinde toplama, çıkarım rank-0'a gönderme, ardından tekrar dağıtma. Bu durum, tek bir GPU'nun PCIe bant genişliği ve NIC (örn. 400 Gbps ≈ 50 GB/s) ile sınırlı, hızla bir dar boğaza dönüşür.

Buna karşılık, noktadan noktaya kurulumumuz, her eğitim GPU'sunun doğrudan her çıkarım GPU'suna gönderme yapmasına olanak tanıyarak tek bir bağlantı yerine tüm ağ yapısını doyurur.

RPC tabanlı aktarım, çıkarım tarafında bir arabellek ve worker çağrıları kullanırken, tek taraflı aktarım ağırlıkları çıkarım worker’ının katılımı olmadan doğrudan yazar.

Bazı sistemler, her tensör için çıkarım motorunun update_weight() metodunu çağırmaya güvenir. Bu, çıkarım kodunda müdahaleci değişiklikler, ayrıca RPC'lerden, serileştirmeden ve kontrol düzlemi koordinasyonundan kaynaklanan ek yük anlamına gelir.

RDMA WRITE temel işleviyle, ekstra kopyalar olmaksızın çıkarım GPU belleğindeki ağırlıkları sessizce güncelleriz. Hiçbir kontrol düzlemi mesajı ve CPU kontrol mantığı dahil olmaz. Çıkarım motorunda hiçbir değişiklik yapılması gerekmez.

Ağırlık aktarım süreci dört tür donanım kaynağından yararlanabilir: (1) Anahandandan cihaza veri taşıma (2) Projeksiyon birleştirme ve nicemleme için GPU hesaplaması (3) Veri düzlemi için RDMA ağı (4) Kontrol düzlemi için Ethernet.

Tasarımımız, ağırlık aktarım görevlerini işlem hattı aşamalarına bölerek farklı donanım kaynakları arasında kolayca örtüşme sağlar.

Bazı uygulamalar her eğitim adımında bir aktarım zamanlamasını yeniden hesaplar, sürekli olarak meta verileri toplar ve talimatları dağıtır. Bu, gereksiz kontrol düzlemi gecikmesi ekler.

Zamanlamamız başlatma sırasında bir kez hesaplanır. Her eğitim yinelemesi planı başarıyla yeniden oynatır: denetleyici bir "devam" sinyali verir ve GPU'lar önceden atanmış rotalarını izler. Yürütme öngörülebilir ve hafiftir.

Tüm ağırlık güncelleme sürecini tek bir monolitik fonksiyona bağlamak caziptir: meta veri toplama, ad eşleştirme, düğüm içi toplama, projeksiyon birleştirme, nicemleme, alt dilimleme iletişim dünyası, düğümler arası ağ aktarımı. Doğru programlamak zordur ve optimize etmek daha de zordur.

Mühendisliğimizde bu adımları ayrı bileşenler olarak ayırıyoruz. Her bileşen izole bir şekilde birim testine tabi tutulabilir, mantığı yürütülebilir ve optimize edilebilir.

Hızlı ve güvenilir ağırlık aktarımı, büyük ölçekli pekiştirmeli öğrenme ince ayarı için kritik bir yapı taşıdır. RDMA WRITE temel işlevini, statik aktarım zamanlamasını ve işlem hattı tabanlı yürütmeyi birleştirerek Kimi-K2 üzerinde trilyon parametreli güncellemeleri yalnızca 1,3 saniyeye indirdik. Yaklaşımın mantığını yürütmek basittir, bakımı kolaydır ve geleneksel tasarımların dar boğazlarından kaçınır.

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE: tek taraflı aktarımlar

Yüksek düzeyli iş akışı

Ağırlık aktarımı yürütme

DeviceMesh ve Mesh Grupları

Görev işlem hattı

GPU bellek kullanımı kontrolü

Neden hızlı ve basit?

Noktadan noktaya iletişim

Tek taraflı veri aktarımı

İşlem hattı oluşturma (Pipelining)

Statik Zamanlama

Temiz ayrım

Sonuç

Alıntı