Co-Training Draf Daring untuk Speculative Decoding dalam RL Post-Training Skala Besar dan Konteks Panjang

Sistem end-to-end yang memungkinkan co-training draf daring lintas tahap context-parallel dan pipeline-parallel untuk model hingga 122B parameter

Zili Wang; Zhaopeng Qiu; Yuekai Zhang; Shuang Yu; Junjie Lai· 2026· DOI 10.48550/arXiv.2609.07108

Masalah inti

RL post-training pada large language model didominasi oleh biaya rollout generation. Speculative decoding adalah teknik yang mapan untuk mempercepat generasi ini dengan menggunakan model draf yang lebih kecil untuk mengusulkan token yang diverifikasi secara paralel oleh model target yang lebih besar. Meskipun pelatihan draf offline umum dilakukan, co-training daring—di mana draf terus diperbarui bersama kebijakan target—dapat lebih meningkatkan akurasi draf dan menghasilkan percepatan yang lebih besar. Namun, penskalaan co-training daring ke model besar dengan konteks panjang menghadirkan dua hambatan mendasar: (1) branch attention tidak didukung oleh implementasi causal context-parallel (CP) standar, dan (2) fitur target tersebar lintas tahap pipeline-parallel (PP). Karya ini mengatasi kedua tantangan tersebut dengan sistem end-to-end yang dirancang untuk co-training draf daring skala besar. Penulis menunjukkan bahwa draf yang di-co-train mengikuti baseline kebijakan dengan ketat sekaligus memberikan percepatan rollout dan end-to-end yang substansial pada skala model hingga 122B parameter. Desain CP mereka mencapai strong scaling pada 256K token dengan penghematan memori yang sign

Inovasi

Penulis melakukan eksperimen pada skala model hingga 122B parameter dan panjang konteks hingga 256K token. Temuan utama meliputi: - **Percepatan rollout dan end-to-end:** Draf yang di-co-train memberikan percepatan substansial baik pada rollout generation maupun RL post-training end-to-end. Draf mengikuti baseline kebijakan dengan ketat, memastikan bahwa percepatan tidak mengorbankan kualitas generasi. - **Penskalaan context parallelism:** Desain CP yang diusulkan mencapai strong scaling pada 256K token. Dibandingkan karya sebelumnya, desain ini menghasilkan penghematan memori yang signifikan, memungkinkan pelatihan konteks yang lebih panjang tanpa melampaui batas memori. - **Overhead pipeline parallelism:** Transport TapChannel hanya menimbulkan overhead yang moderat, sehingga praktis untuk penerapan skala besar. - **Skala model:** Sistem divalidasi pada model hingga 122B parameter, menunjukkan penerapannya pada large language model mutakhir. Hasil kuantitatif tidak dirinci sepenuhnya dalam abstrak, tetapi penulis menekankan bahwa draf yang di-co-train mengikuti baseline kebijakan dengan ketat sekaligus memberikan percepatan yang substansial. Kode tersedia di repositori GitHub y

RL post-training pada large language model didominasi oleh biaya rollout generation. Speculative decoding adalah teknik yang mapan untuk mempercepat generasi ini dengan menggunakan model draf yang lebih kecil untuk mengusulkan token yang diverifikasi secara paralel oleh model target yang lebih besar. Meskipun pelatihan draf offline umum dilakukan, co-training daring—di mana draf terus diperbarui bersama kebijakan target—dapat lebih meningkatkan akurasi draf dan menghasilkan percepatan yang lebih besar. Namun, penskalaan co-training daring ke model besar dengan konteks panjang menghadirkan dua hambatan mendasar: (1) branch attention tidak didukung oleh implementasi causal context-parallel (CP) standar, dan (2) fitur target tersebar lintas tahap pipeline-parallel (PP). Karya ini mengatasi kedua tantangan tersebut dengan sistem end-to-end yang dirancang untuk co-training draf daring skala besar. Penulis menunjukkan bahwa draf yang di-co-train mengikuti baseline kebijakan dengan ketat sekaligus memberikan percepatan rollout dan end-to-end yang substansial pada skala model hingga 122B parameter. Desain CP mereka mencapai strong scaling pada 256K token dengan penghematan memori yang signifikan dibandingkan karya sebelumnya, dan transport PP mereka hanya menimbulkan overhead yang moderat.

Sistem yang diusulkan mengatasi kedua hambatan melalui strategi paralelisasi khusus.

Mengapa penting

Karya ini mengatasi hambatan kritis dalam RL post-training: biaya rollout generation. Dengan memungkinkan co-training draf daring pada skala besar, sistem ini memungkinkan perbaikan berkelanjutan pada model draf, yang mengarah pada tingkat penerimaan yang lebih tinggi dan percepatan yang lebih besar. Dua kontribusi teknis—branch attention pada CP dan TapChannel pada PP—sangat penting untuk penskalaan ke model besar dan konteks panjang. Desain CP memperluas packed, load-balanced zigzag ring attention untuk mendukung branch attention, yang diperlukan untuk speculative decoding. Ini merupakan perluasan yang tidak trivial karena branch attention memperkenalkan kompleksitas tambahan dalam load balancing dan kausalitas. Penghematan memori pada 256K token sangat penting, karena pelatihan konteks panjang semakin umum. Transport PP melalui TapChannel adalah solusi cerdas untuk masalah distribusi fitur lintas tahap pipeline. Dengan menggunakan jalur komunikasi terpisah, penulis menghindari gangguan pada jadwal pipeline, yang jika tidak akan menyebabkan inefisiensi. Overhead yang moderat menunjukkan bahwa pendekatan ini praktis untuk penerapan di dunia nyata. Salah satu batasan adalah abstrak tidak memberikan perbandingan kuantitatif terperinci dengan karya sebelumnya, seperti angka percepatan yang tepat atau persentase penghematan memori. Penelitian selanjutnya dapat mengeksplorasi trade-off antara ukuran model draf dan tingkat penerimaan, serta dampak frekuensi co-training. Selain itu, kinerja sistem pada model yang lebih besar (misalnya, 500B+ parameter) dan konteks yang lebih panjang (misalnya, 1M token) masih perlu diuji. Secara keseluruhan, makalah ini menyajikan langkah signifikan menuju RL post-training yang efisien untuk large language model, dengan kode open-source untuk memfasilitasi penelitian lebih lanjut.

Siapa yang sebaiknya membaca

Praktisi dan peneliti ilmu komputer

Membuka konten member…