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

TL;DR

The paper addresses nondeterminism in LLM inference arising from floating-point non-associativity and tensor-parallel size variability, which causes training–inference mismatch in RL. It introduces Tree-Based Invariant Kernels (TBIK) that align intra- and inter-GPU reductions using a fixed binary-tree topology, providing theoretical guarantees and practical implementations in Triton, vLLM, and FSDP. Empirical results show zero probability divergence and bitwise reproducibility across TP sizes and hardware, and the approach successfully bridges the probability gap between vLLM and FSDP in RL pipelines. This work advances reproducible, on-policy RL with large language models and lays groundwork for further optimizations and extensions to low-bit quantized kernels.

Abstract

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 exhibit non-deterministic behavior: identical inputs can yield different outputs when system configurations (e.g., tensor parallel (TP) size, batch size) vary, 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 (i.e., TP = 1) while the rollout engine relies on multi-GPU TP to maximize the inference throughput, creating a natural mismatch between the two. This precision mismatch problem may lead to suboptimal performance or even collapse for RL training. We identify and analyze the root causes of TP-induced inconsistency and propose Tree-Based Invariant Kernels (TBIK), a set of TP-invariant matrix multiplication and reduction primitives that guarantee bit-wise identical results regardless of TP size. Our key insight is to align intra- and inter-GPU reduction orders through a unified hierarchical binary tree structure. We implement these kernels in Triton and integrate them into vLLM and FSDP. Experiments confirm zero probability divergence and bit-wise reproducibility for deterministic inference across different TP sizes. Also, we achieve bit-wise identical results between vLLM and FSDP in RL training pipelines with different parallel strategy. Code is available at https://github.com/nanomaoli/llm_reproducibility.

Figures (12)

• Figure 1: The illustration of numerical variance in LLM-based RL training. Different frameworks and tensor-parallel settings lead to noticeable probability discrepancies for the same model, making it difficult to achieve stable and truly on-policy reinforcement learning.
• Figure 2: Illustration of tensor-parallel weight partitioning in the Transformer model architecture (e.g., Qwen3 Dense) under the vLLM tensor-parallel setting. In this configuration, the QKV proj, gate proj, up proj, and lm head layers are column-parallel, while the o proj and down proj layers are row-parallel. For non-matmul operations such as RMSNorm and RoPE, parameters and computations are not split across GPUs, and these layers are omitted from the figure for clarity.
• Figure 3: Implementations in Batch Invariant Operations. Left: RMSNorm. Right: MatMul.
• Figure 4: Accuracy standard deviation of Qwen3-8B on AIME24 dataset under different tensor parallelism (TP = 1/2/4/8) settings. While BIO enhances output determinism relative to vanilla BF16 inference, the accuracy variance can still reach over 4%.
• Figure 5: illustrations of column-parallel and row-parallel Matmul.