Deterministic Inference across Tensor Parallel Sizes That Eliminates Training-Inference Mismatch

Deterministic inference is increasingly critical for Large language model (LLM) applications such as LLM-as-a-judge evaluation, multi-agent systems, and reinforcement learning (RL). However, existing LLM serving frameworks can produce different outputs for identical inputs when tensor parallel (TP) size or batch size changes, even under greedy decoding. This arises from the non-associativity of floating-point arithmetic and inconsistent reduction orders across GPUs. While prior work has addressed batch-size–related nondeterminism through batch-invariant kernels, determinism across different TP sizes remains an open problem, particularly in RL settings, where the training engine typically uses Fully Sharded Data Parallel (FSDP) (i.e., TP = 1) while the rollout engine relies on multi-GPU TP to maximize the inference throughput, creating a probability mismatch that can degrade or even collapse training. We identify and analyze the root causes of TP-induced inconsistency and propose Tree-Based Invariant Kernels (TBIK), a set of custom matrix multiplication and reduction kernels that guarantee bit-wise identical results across TP sizes. Our key insight is to enforce a consistent reduction order across and within GPUs. We implement TBIK in Triton and integrate it into vLLM and FSDP, achieving bit-wise deterministic inference across different TP sizes and zero probability divergence between rollout and training engines in RL pipelines. By eliminating mismatches caused by different parallelization strategies, TBIK enables true on-policy RL at scale for the first time, leading to improved model performance and faster convergence.

Large language models can give different results to the same input when they are run in different system settings, such as using different numbers of GPUs or serving different numbers of users at once. These small differences can make training unstable and hurt the final model. We found that the root cause is floating-point arithmetic, the way computers approximate and add numbers. When the same values are added in different orders, the final result can change slightly. We designed new computation and communication methods that keep this order consistent across different ways of running the model. Our experiments show that this makes model outputs consistent and helps training become more stable, converge faster, and achieve better performance.