Co-addestramento online del modello draft per la decodifica speculativa nel post-addestramento RL su larga scala e con contesti lunghi

Abstract degli autori

La decodifica speculativa accelera la generazione dei rollout, che rappresenta la voce di costo dominante nel post-addestramento con apprendimento per rinforzo (RL). Il co-addestramento online può aumentare ulteriormente l’accuratezza del modello draft, producendo accelerazioni maggiori. Tuttavia, estendere questo approccio al co-addestramento di modelli di grandi dimensioni con contesti lunghi presenta due ostacoli: (1) le implementazioni standard del parallelismo del contesto (CP) causale non supportano l’attenzione sui rami e (2) le rappresentazioni del modello bersaglio sono distribuite tra le fasi del parallelismo di pipeline (PP). Affrontiamo entrambi gli ostacoli con un sistema end-to-end per il co-addestramento online del modello draft su larga scala. Per il CP, estendiamo l’attenzione ad anello a zigzag con sequenze compattate e carico bilanciato, combinando l’attenzione sui rami locale a ciascun rank con l’attenzione causale sulla sequenza principale. Per il PP, TapChannel trasporta le rappresentazioni intermedie del modello bersaglio tra le fasi attraverso un percorso separato, senza modificare la pianificazione della pipeline. Gli esperimenti mostrano che i modelli draft co-addestrati seguono da vicino la policy baseline, offrendo al contempo accelerazioni sostanziali nella generazione dei rollout e nell’intero processo, con modelli di dimensioni fino a 122B. Il nostro progetto per il CP raggiunge una buona scalabilità forte a 256K token, con un notevole risparmio di memoria rispetto ai lavori precedenti, mentre il trasporto per il PP comporta un sovraccarico modesto. Il codice è disponibile all’indirizzo https://github.com/NVIDIA-NeMo/RL/issues/3698.