Achieving High Performance in True On-Policy Reinforcement Learning for Mixture-of-Experts LLMs
Abstract
Reinforcement learning (RL) post-training for large language models (LLMs) typically generates rollouts with one engine and trains with another. The engines run different kernels on different batch shapes and assign different log-probabilities to the same tokens, so the behavior likelihood the trainer reconstructs differs from the likelihood each token was sampled under, an unmodeled error in the PPO/GRPO likelihood ratios. We make the generation engine compute the training forward's exact bits: on a 128-expert mixture-of-experts model with expert-parallel sharding, every sampled-token log-probability is bitwise identical to the training forward over full-learning-rate GRPO steps. The stack keeps the engine's native kernels, CUDA graphs, and in-switch collectives, constraining their algorithm selection and matching scoring-side arithmetic. At production batch sizes, bitwise exactness costs 2.5\% of generation throughput against the same engine with determinism off, and about 10\% of the full RL iteration.