Training and rollout stop disagreeing: vime and RL-Kernel hold logprobs bitwise identical for 200 steps on AMD MI300X
Synopsis
In reinforcement-learning post-training, the training engine and the rollout engine can compute different logprobs even with the same model, weights and inputs, because kernels, parallelism and reduction order differ, which affects the importance ratio, KL diagnostics and clipping. vime aligns the training timeline while RL-Kernel aligns numerical execution, and together they put attention, feed-forward, logprob and communication paths on one numerical contract on AMD Instinct MI300X. In a Qwen3-8B GRPO experiment on 8 MI300X GPUs, the strict path ran 200 consecutive steps with mismatch_count and max_abs_diff both zero.
Interpretation
Training and rollout can now produce bitwise-identical logprobs for the same token on the AMD ROCm platform, not merely numerically close values. Training and rollout use different execution engines and operators, so even with identical model, weights and inputs, differences in block size, Split-K or Split-KV strategy, reduction order, fusion patterns and intermediate precision can change logprobs; existing work has focused mainly on NVIDIA platforms, with more limited systematic exploration on ROCm. In an end-to-end Qwen3-8B GRPO experiment on 8 AMD Instinct MI300X GPUs, the vime plus RL-Kernel strict path ran 200 consecutive steps with mismatch_count = 0 and max_abs_diff = 0.
The division of labor is that vime manages the training timeline and RL-Kernel manages numerical execution, and only together do they carry the same token batch through training and rollout along one numerical path. vime connects Megatron training, vLLM rollout and a data buffer into a complete RL post-training workflow, tracking which token batch belongs to which step, which weight version generated it, which rollout record enters an update and when new weights are synchronized to vLLM; RL-Kernel specifies which values participate in a computation, how they are partitioned and merged, which intermediate precision is used and where rounding occurs. Validation ran on 1 node with 8 MI300X 192GB GPUs, Megatron at TP4/CP2/PP1, rollout with 2 vLLM engines at TP4 each, dataset dapo-math-17k, 1 prompt times 8 samples per step, KL loss coefficient 0.001, and HIP Graph at FULL_AND_PIECEWISE preserving the production execution path.
Bitwise consistency is not achieved by pinning a single operator but by fixing the forward boundaries that determine logprobs as one numerical chain. RMSNorm, GEMM, attention, linear logprob and distributed collectives look separate, but all merge partial results, and block size, reduction trees, exp and log implementations, intermediate precision and fusion can change merge order and rounding boundaries, so pinning one kernel or collective is not enough to guarantee bitwise-identical logprobs. The strict path covers RMSNorm reduction domain and epsilon, attention position, mask and split policy, GEMM and SwiGLU K-dimension reduction and accumulation precision, linear logprob true vocabulary range and cross-rank LSE merge, and collective payload ownership, dtype and fixed rank reduction order; on ROCm it adds deterministic MFMA GEMM kernels, vocabulary reduction and HIP IPC communication, and fixes the AITER/CK attention execution schedule.
Perspective
The current end-to-end validation covers Qwen3-8B Dense, vime, vLLM, Megatron-LM and AMD Instinct MI300X, run on 1 node with 8 MI300X 192GB GPUs; next steps are to extend the work to more MoE and multimodal models and additional AMD GPU architectures while continuing to optimize the ROCm strict path.
The bitwise-consistency conclusion comes from a 200-step experiment with Qwen3-8B Dense on 8 MI300X GPUs, and extension to MoE, multimodal models and other AMD GPU architectures is planned work; the strict path does not require training and rollout to share every memory layout or scheduling policy, as long as those optimizations do not change the numerical semantics of the compared results.
