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