Prime Intellect ha rilasciato prime-rl versione 0.6.0, un framework opensource che si concentra su reinforcement learning su modelli Mixture-of-Experts (MoE) a scala trilione. Prime-rl supporta workload complessi di agenti, come software engineering su orizzonti di lunga durata. Questa versione introduce nuove ottimizzazioni per l’allenamento su grandi modelli come GLM-5.
Il team di ricerca ha allenato GLM-5 su attività SWE fino a una lunghezza di sequenza di 131.000 token. Tempi di esecuzione di ogni step sono rimasti sotto i 5 minuti. La dimensione del batch era di 256 rollouts. L'addestramento è andato a buon fine utilizzando soltanto 28 H200 nodes.
TL;DR
prime-rl 0.6.0 permette di allenare modelli MoE da trilione di parametri su workload RL con agenti.
GLM-5 è stato allenato su SWE a una lunghezza massima di 131k token, con step sottocinque minuti e utilizzo di soli 28 H200 nodes.
L'RL asincrono separa trainer e inferenza per ottimizzarli in autonomia.
L'inferenza include FP8, Wide EP, P/D disaggregation, KV offloading, e router replay.
Il training utilizza 3-D parallelism (FSDP, EP, CP) più FP8 a blocchi.
Che cos’è prime-rl 0.6.0?
Prime-rl è un framework opensource per reinforcement learning asincrono. Permette di postallenare modelli open-source su compiti agenti. La versione 0.6.0 estende le sue capacità per gestire scale da trilione di parametri.
Il modello esemplare presentato nell’annuncio è zai-org/GLM-5.1. Le ottimizzazioni si applicano anche ad altri modelli di grandi dimensioni. Esempi notevoli sono moonshotai/Kimi-K2.7-Code e nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16.
Comando Avvio GLM-5.1
Per avviare un esecuzione full di GLM-5.1 su un cluster Slurm si utilizza:
uv run rl @ examples/glm5_llmd/rl.toml --output-dir /shared/outputs/glm5-llmd
Questo comando avvia un training complesso e distribuito per modelli di grandi dimensioni.
RL asincrono: ruolo del training
Le attività con agenti presentano outlier di lunga durata. Alcuni rollouts di coding possono andare avanti per molte ore. Attendere a conclusione di questi prima dell'aggiornamento politica sarebbe problematico.
L’RL asincrono evita questo problema. Sistema trainer e inferenza vengono disaggiorati e quindi evolvono autonomamente. L’inferenza politica si aggiorna appena l’optimiser completa un passo.
Punto di sincronizzazione unico: l'aggiornamento politico. Prime-rl spinge i nuovi pesi non appena sono disponibili. I rollouts attivi mantengono il loro cache di prefisso attivo. Così, un singolo rollout può utilizzare token provenienti da diverse versioni di politica.
I nuovi rollouts invece procedono in maniera diversa. Essi reinizializzano la propria cache KV, persino se il prefisso è lo stesso. Un “KV-cache salt” costringe a questo comportamento. Richieste basate su politiche troppo vecchie vengono scartate. Il valore maxoffpolicy_steps controlla questo limite.
Ottimizzazioni per l'inferenza
La fase di inferenza rappresenta spesso il collo di bottiglia in sistemi RL. Prime-rl ottimizza per throughput, mantenendo latenza limitata.
FP8
La bassa precisione FP8 velocizza prefill e decode. Prime-rl utilizza FP8 con DeepEP e DeepGEMM per accelerare le operazioni di inferenza.
Wider Expert Parallelism (EP)
Il Wide EP distribuisce esperti tra ≥32 GPU. Si accompagna a una grande grandezza di parallelismo dei dati, ad esempio 32. Ogni GPU mantiene esperti distanti, agendo come endpoint. Sincronizzazione avviene per livello, tramite operazioni di dispatch e combine.
Disaggregazione Prefill/Decode
Coppie di modello-ambiente possono presentare un rapporto 4:1 prefill/decode. Worker condivisi possono sovrascaricare la latenza end-to-end. Ciò riduce la capacità di PipelineRL. La separazione P/D divide worker di prefill e decode. Gli output lunghi smettono di limitare i worker di decode.
Gestione KV Cache
Concurrencies alte richiedono spazio di KV Cache esteso. Prime-rl supporta offloading su CPU e disco. La vLLM prevede una pool per worker. Mooncake Store invece condivide RAM e disco in tutti i nodi.
Routing delle richieste
Prime-rl include come default una versione forkata di vllm-router. Supporta anche il router NVIDIA Dynamo come soluzione a drop-in. Router valutano i worker per cache di KV, profondità della coda e carico in tempo reale.
Router Replay (R3)
Discrepanza tra trainer e inferenza uccide di nascosto il training. Il router capture le decisioni di routing dell'inferenza. Le riproduce direttamente nel trainer. Ciò taglia di circa un ordine di grandezza il mismatch KL. I parametri esperto hanno una forma [numlayers, topk, seq_len]. Questo payload può arrivare a centinaia di GB. In grandi scale il tasso di dati raggiunge svariati Gbps. Prime-rl li tratta come payload opaco, gestito mediante operazioni PyTorch ottimizzate.
Ottimizzazioni per Training
Il trainer si basa su torchtitan, una codebase di training nativo PyTorch. Si appoggia a parallelism 3-D: FSDP, EP e CP. L’esempio GLM-5 usa tutti e tre.
Strategia | Cosa si parcella | Uso principale | Dettaglio chiave
- FSDP (FSDP2): Parametri, gradienti, stati ottimizzatore
- Baseline memory amortizzazione
- Scheda per layer via fully_shard
- Expert Parallelism (EP): Esperti per layer
- Riduce attivo layer memory
- Distribuzione all2all; utilizzo DeepEP o native PyTorch
- Context Parallelism (CP): Dimensione sequenza
- Memoria attivazione per lunga context
- Uly