消除训练-推理不匹配的确定性推理跨张量并行规模
机器学习
2025-11-25 v1 计算与语言
机器学习
摘要
确定性推理正在成为大型语言模型(LLM)应用的关键因素，如LLM-as-a-judge评估、多智能体系统和强化学习(RL)。然而，现有的LLM服务框架 exhibits 非确定性行为：即使在贪婪解码下，相同输入在系统配置(如张量并行(TP)大小、批量大小)不同时也会产生不同输出。这源于浮点算术的非交换性以及跨GPU的归约顺序不一致。虽然先前工作通过批大小不变的内核解决了批大小相关的非确定性问题，但关于不同TP大小下的确定性推理仍是未解决的尤其是在RL场景中，训练引擎通常使用完全分片数据并行(即TP=1)，而推演引擎依赖于多GPU TP 以最大化推理吞吐量，导致两者之间存在天然的不匹配。这种精度不匹配问题可能导致RL训练的次优性能甚至崩溃。我们识别并分析了TP引起的不一致的根本原因，提出基于树形不变内核(TBIK)，一套TP不变的矩阵乘法和归约原语，确保无论TP大小都能实现位相同的结果。我们的关键洞察是通过统一的层次结构二叉树对齐 intra-和 inter-GPU 归约顺序。我们在Triton中实现这些内核并将其集成到vLLM和FSDP中。实验确认在不同TP大小下实现零概率发散和位级可重现性。此外，我们在具有不同并行策略的RL训练管道中实现vLLM和FSDP之间的位级相同结果。代码可在 https://github.com/nanomaoli/llm_reproducibility 查看。
引用
@article{arxiv.2511.17826,
title = {Deterministic Inference across Tensor Parallel Sizes That Eliminates Training-Inference Mismatch},
author = {Ziyang Zhang and Xinheng Ding and Jiayi Yuan and Rixin Liu and Huizi Mao and Jiarong Xing and Zirui Liu},
journal= {arXiv preprint arXiv:2511.17826},
year = {2025}
}