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

2025/11/21 by Zhang, Ziyang, Ding, Xinheng, Yuan, Jiayi +4 · 1 citation
Computer Science · Materials Science · Mathematics · #Computation and Language (cs.CL) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Machine Learning in Materials Science #Tensor decomposition and applications #Topic Modeling

paper · doi:10.48550/arxiv.2511.17826

openalex publication_date 2025/11/21 · openalex created_date 2025/11/27 · openalex updated_date 2026/07/28

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.

Citations

• Trust Region Masking for Long-Horizon LLM Reinforcement Learning
• Defeating the Training-Inference Mismatch via FP16
• Group Sequence Policy Optimization
• Automating Expert-Level Medical Reasoning Evaluation of Large Language Models
• LMR-BENCH: Evaluating LLM Agent's Ability on Reproducing Language Modeling Research
• Understanding and Mitigating Numerical Sources of Nondeterminism in LLM Inference
• An Outlook on the Opportunities and Challenges of Multi-Agent AI Systems
• Qwen3 Technical Report
• The Science of Evaluating Foundation Models
• LLMs-as-Judges: A Comprehensive Survey on LLM-based Evaluation Methods
• Flex Attention: A Programming Model for Generating Optimized Attention Kernels
• On the Resilience of LLM-Based Multi-Agent Collaboration with Faulty Agents
• The Llama 3 Herd of Models
• KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches
• Judging the Judges: Evaluating Alignment and Vulnerabilities in LLMs-as-Judges
• BioXP-0.5B: Explainable Medical-AI via RL-GRPO
• Large Language Model based Multi-Agents: A Survey of Progress and Challenges
• SGLang: Efficient Execution of Structured Language Model Programs
• Efficient Memory Management for Large Language Model Serving with PagedAttention
• A Survey on Evaluation of Large Language Models
• Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena
• AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration
• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel
• GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Training Verifiers to Solve Math Word Problems
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

Cited by

• Stage-Replay Divergence Follows the KV Cache: Fixed-Prefix Precision Controls and Bidirectional Cache Transplantation

Related

• TP-Aware Dequantization
• MARLIN: Mixed-Precision Auto-Regressive Parallel Inference on Large Language Models
• Generalized Parallel Scaling with Interdependent Generations
• Tender: Accelerating Large Language Models via Tensor Decomposition and Runtime Requantization
• Shift Parallelism: Low-Latency, High-Throughput LLM Inference for Dynamic Workloads