Transfer Bobot untuk Pasca-Pelatihan RL dalam waktu di bawah 2 detik

Kami baru saja berhasil mencapai pembaruan parameter lintas mesin 1,3 detik untuk Kimi-K2 (1T parameter), mentransfer bobot dari 256 GPU pelatihan (BF16) ke 128 GPU inferensi (FP8).

Dalam fine-tuning reinforcement learning asinkron, pelatihan dan inferensi berjalan pada GPU terpisah. Setelah setiap langkah pelatihan, bobot baru harus dikirimkan ke node inferensi. Banyak kerangka kerja yang ada memakan waktu beberapa detik—atau bahkan menit—untuk model dengan parameter triliunan.

Dengan memanfaatkan komunikasi point-to-point RDMA, kami dapat membuat transfer bobot menjadi sangat cepat tanpa mengubah mesin inferensi, serta membuat kode lebih mudah ditulis dan dipelihara.

Solusi kami dibangun di atas RDMA WRITE, sebuah primitif satu sisi di mana sumber langsung menulis ke dalam memori GPU tujuan.

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.
    ...

Sisi tujuan bahkan tidak akan mendapatkan pemberitahuan untuk transfer tersebut. Ini memberi kami transfer latensi rendah, throughput tinggi, tanpa salinan (zero-copy) yang digerakkan oleh node pelatihan tanpa logika kontrol apa pun pada node inferensi.

Pengumpulan metadata – Kontroler mengumpulkan metadata parameter dari semua GPU pelatihan dan inferensi.

Penghitungan jadwal – Kontroler menghitung jadwal transfer bobot statis, memetakan GPU pelatihan mana yang mengirimkan parameter apa ke GPU inferensi mana, dan dalam urutan apa.

Distribusi jadwal – Kontroler mengirimkan jadwal ke semua GPU pelatihan.

Eksekusi – Setelah setiap langkah pelatihan, kontroler memberi sinyal kepada GPU pelatihan untuk memulai transfer.

Dengan alur kerja tingkat tinggi yang telah ditentukan, tantangan utamanya adalah bagaimana mengeksekusi transfer bobot secara efisien pada skala triliunan parameter. Di sini kami menjelaskan detail jalur eksekusi.

Parameter dalam pelatihan didistribusikan sesuai dengan penempatan FSDP. Menggunakan full_tensor(), semua GPU dalam DeviceMesh dapat merekonstruksi parameter penuh, sehingga semuanya dapat berfungsi sebagai sumber untuk transfer bobot.

Beberapa DeviceMesh yang terpisah membentuk grup mesh. Karena DeviceMesh dalam grup yang sama terpisah, transfernya tidak saling mengganggu dan dapat berjalan sepenuhnya secara paralel. Antar grup mesh, kami menyisipkan pembatas global untuk menegakkan urutan.

Pipeline transfer bobot tumpang tindih dengan penyalinan host-to-device, perakitan tensor, fusi dan kuantisasi, serta transfer RDMA lintas MeshGroups.

Kami memperlakukan transfer setiap tensor parameter sebagai sebuah tugas. Proses transfer bobot memanfaatkan berbagai jenis sumber daya perangkat keras, sehingga kami membagi tugas transfer bobot ke dalam tahap pipeline berbeda yang tumpang tindih dalam waktu:

Memcpy host ke perangkat — Jika FSDP memindahkan bobot ke CPU

Persiapan parameter — Rekonstruksi bobot penuh dengan full_tensor(), terapkan fusi proyeksi, kuantisasi jika diperlukan.

Transfer RDMA — Penulisan zero-copy ke memori GPU inferensi jarak jauh

Pembatas global — Setelah semua panggilan full_tensor() selesai, sinkronkan lintas grup mesh menggunakan GLOO melalui Ethernet.

Dalam implementasinya, kami mempertahankan antrean FIFO untuk tugas di setiap tahap pipeline. Kapan pun tugas di kepala antrean menyelesaikan tahap tersebut, tugas itu dipindahkan ke ekor antrean tahap berikutnya.

full_tensor() dan operasi GPU lainnya memperkenalkan penggunaan memori GPU tambahan. Untuk menghindari galat kehabisan memori (out of memory), kami memulai eksekusi suatu tugas hanya jika tugas yang sedang berjalan saat ini memakan memori GPU sementara yang lebih sedikit daripada watermark yang dapat dikonfigurasi.

Beberapa pilihan desain membuat sistem kami berjalan jauh lebih cepat dan lebih mudah dipelihara daripada solusi sumber terbuka umum.

Transfer bobot berbasis Rank0 mengumpulkan bobot sebelum mendistribusikannya ke node inferensi.

Komunikasi point-to-point mentransfer bobot dari kedua node pelatihan ke antarmuka jaringan node inferensi.

Pola umum adalah menyalurkan semua parameter melalui GPU peringkat-0: mengumpulkan pada peringkat pelatihan-0, mengirim ke peringkat inferensi-0, lalu menyebarkan kembali. Ini dengan cepat menjadi titik hambatan, dibatasi oleh bandwidth PCIe satu GPU dan NIC (misalnya, 400 Gbps ≈ 50 GB/s).

Sebaliknya, penyiapan point-to-point kami memungkinkan setiap GPU pelatihan untuk mengirim langsung ke setiap GPU inferensi, menjenuhkan seluruh fabrik jaringan alih-alih satu tautan tunggal.

Transfer berbasis RPC menggunakan buffer di sisi inferensi dan pemanggilan worker, sedangkan transfer satu sisi menulis bobot secara langsung tanpa partisipasi worker inferensi.

Beberapa sistem mengandalkan pemanggilan metode update_weight() pada mesin inferensi untuk setiap tensor. Itu berarti perubahan intrusif pada kode inferensi, ditambah overhead dari RPC, serialisasi, dan koordinasi bidang kontrol.

Dengan primitif RDMA WRITE, kami memperbarui bobot secara diam-diam pada memori GPU inferensi, tanpa salinan tambahan. Tidak ada pesan bidang kontrol dan tidak ada logika kontrol CPU yang terlibat. Tidak diperlukan modifikasi pada mesin inferensi.

Proses transfer bobot dapat memanfaatkan empat jenis sumber daya perangkat keras: (1) Perpindahan data host-perangkat (2) Komputasi GPU untuk fusi proyeksi dan kuantisasi (3) Jaringan RDMA untuk bidang data (4) Ethernet untuk bidang kontrol.

Desain kami membagi tugas transfer bobot ke dalam tahap-tahap pipeline, memungkinkan tumpang tindih yang mudah di berbagai sumber daya perangkat keras.

Beberapa implementasi menghitung ulang jadwal transfer pada setiap langkah pelatihan, yang berulang kali mengumpulkan metadata dan mendistribusikan instruksi. Ini menambah latensi bidang kontrol yang tidak perlu.

Jadwal kami dihitung sekali saat inisialisasi. Setiap iterasi pelatihan cukup memutar ulang rencana: kontroler mengeluarkan sinyal "mulai", dan GPU mengikuti rute yang telah ditentukan sebelumnya. Eksekusi dapat diprediksi dan ringan.

Sangat menggiurkan untuk mencampuradukkan seluruh proses pembaruan bobot dalam satu fungsi monolitik: mengumpulkan metadata, pencocokan nama, pengumpulan intra-node, fusi proyeksi, kuantisasi, dunia komunikasi sub-pemotongan, transfer jaringan antar-node. Ini sulit diprogram dengan benar, dan bahkan lebih sulit dioptimalkan.

Dalam rekayasa kami, kami memisahkan langkah-langkah ini sebagai komponen individu. Setiap komponen dapat diuji unitnya, dinalarikan, dan dioptimalkan secara terisolasi.

Transfer bobot yang cepat dan andal adalah blok penyusun penting untuk fine-tuning RL skala besar. Dengan menggabungkan primitif RDMA WRITE, jadwal transfer statis, dan eksekusi terpipa, kami mereduksi pembaruan parameter triliunan menjadi hanya 1,3 detik pada Kimi-K2. Pendekatan ini mudah dinalarikan, mudah dipelihara, dan menghindari kemacetan dari desain tradisional.

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: transfer satu sisi

Alur kerja tingkat tinggi

Eksekusi transfer bobot

DeviceMesh dan Grup Mesh

Pipeline tugas

Kontrol penggunaan memori GPU

Mengapa ini cepat dan sederhana

Komunikasi point-to-point

Transfer data satu sisi

Pipelining

Jadwal Statis

Pemisahan yang bersih

Kesimpulan

Kutipan