Squeezing 1TB Model Rollout Into a Single H200: INT4 QAT RL End-to-End Practice

发表时间: 2026-01 · Blog post by SGLang RL Team, InfiXAI Team, Ant Group Asystem & AQ Infra Team, slime Team, RadixArk Team (lmsys.org)

原文: https://lmsys.org/blog/2026-01-26-int4-qat

作者/机构：SGLang RL Team, InfiXAI Team, Ant Group Asystem & AQ Infra Team, slime Team, RadixArk Team (January 26, 2026)

速读

一句话结论
本文在强化学习训练中实现了一套端到端的 INT4 量化感知训练（QAT）方案，通过训练时伪量化与推理时 W4A16 真实量化相结合，将 1TB 级大模型的经验采样压缩至单台 H200 节点内，在消除跨节点通信瓶颈的同时，保持了与 BF16 全精度一致的训练稳定性和模型能力。

要解决什么问题
在大语言模型的强化学习阶段，经验采样（Rollout）环节需要频繁进行模型推理。当模型规模达到 1TB 级别（如 Kimi-K2-Thinking）时，全精度或常规精度的模型权重无法完整装入单台机器的显存（如 141GB 的 H200），必须依赖跨节点并行，这引入了极其高昂的跨节点通信延迟，成为拖慢整体采样效率的核心瓶颈。为了压缩显存，常规做法是采用训练后量化直接在推理时降低权重精度，但这会引发严重的机制卡点：由于模型在训练期从未见过量化噪声，直接将权重压缩至 INT4 会导致严重的信息丢失和特征分布偏移。这种训练端与推理端的精度不匹配，会使得输出概率的误差随训练步数不断震荡放大，最终导致强化学习训练崩溃或效果大幅衰减。

怎么做的
核心思路是引入量化感知训练，让模型在训练阶段就提前适应推理时的低精度截断误差，从而实现训练与推理的严格对齐。该方案由训练端的伪量化、权重更新端的真实量化以及推理端的 W4A16 加速三个关键部件构成。
在训练阶段，系统在显存中始终维护高精度的 BF16 主权重。在正向传播时执行“伪量化”：不改变物理存储格式，而是按组计算最大绝对值，将高精度浮点数动态映射并裁剪到 INT4 的 $[-7, 7]$ 整数区间，随后立即反量化回浮点数。这一过程人为注入了量化误差，迫使模型通过梯度更新来补偿这种精度损失。在反向传播时，由于量化核心的取整操作导数几乎处处为零，会阻断梯度，作者引入了直通估计器（STE）。STE 在反向传播时将取整操作的导数强制定义为 1：

$$ \frac{\partial \text{Round}(x)}{\partial x} = 1 $$


这使得梯度能够无损穿透量化层，顺利更新底层的 BF16 主权重。
在权重更新与流转阶段，系统提取收敛后的 BF16 权重进行“真实量化”。通过分组计算缩放因子，将浮点数不可逆地压缩至 INT4 域。由于 PyTorch 缺乏原生 INT4 类型，系统利用位运算将 8 个 4-bit 数值紧密打包进 1 个 INT32 整数中，完成从标准训练格式到 Marlin 推理引擎格式的转换。
在 Rollout 推理阶段，采用 W4A16 模式（INT4 权重乘以 BF16 激活值）。推理引擎通过位移和按位与操作以极低开销解包权重，并全程调用 BF16 Tensor Cores 进行矩阵乘法。这既将显存占用降低了 75%，又保证了计算精度与训练端的高度一致。

效果如何
实验基于 Qwen3-30B、Qwen3-235B 以及 Kimi-K2-Thinking 模型，在 dapo-math-17k 数据集上进行强化学习训练，并在 aime-2024 基准上每 10 步进行一次严格评估。硬件环境采用 NVIDIA H 系列 GPU。对比基线包括代表全精度上限的“BF16 训练–BF16 推理”路线、代表主流低精度方案的“BF16 训练–FP8 推理”路线，以及代表常规量化的“无 QAT 训练–直接 INT4 推理”路线。
量化结果表明，在 AIME 评估得分和奖励增长曲线上，INT4 QAT 方案与 BF16 全精度基线高度重合，不仅峰值分数一致，且完全避免了常规量化路线中误差震荡放大的问题。在衡量训练与推理一致性的输出概率绝对误差指标上，INT4 QAT 的误差极低，几乎与 BF16 基线贴合，且显著优于 FP8 推理方案。在采样效率方面，INT4 的极致显存压缩使得 1TB 级模型能够完整塞进单台机器，彻底消除了双节点部署时的跨节点通信延迟，单节点 Rollout 速度实现了质的飞跃。
作者也坦诚了该方案的局限与代价：首先，由于当前的 H 系列 GPU 缺乏原生的 INT4 Tensor Cores，W4A16 实际上仍在使用 BF16 单元计算，因此在不跨节点的情况下，INT4 仅缓解了显存带宽压力，并未获得计算层面的成倍加速，单步延迟仅略优于 FP8；其次，训练端引入的伪量化操作增加了额外的计算开销，导致训练步的耗时明显长于纯 BF16 训练，这在一定程度上抵消了 Rollout 环节带来的端到端时间收益。

主要贡献

在强化学习（RL）训练中，极大参数量模型的经验采样（Rollout）阶段往往面临严重的显存与跨节点通信瓶颈。为了解决这一问题并在开源生态中提供平衡高性能与低成本的参考实现，本文基于 Kimi K2 团队的技术启发，成功落地了端到端的 INT4 感知量化训练（Quantization-Aware Training, QAT） 完整流水线。

本文的核心创新点与贡献包括：

  1. 突破显存瓶颈实现单机 Rollout：通过权重压缩与低比特量化，将约 1TB 规模的类 K2 模型压缩至单张 H200（141GB）GPU 中，彻底消除了跨节点通信的瓶颈，使 Rollout 效率相较于 W8A8（FP8 权重与激活）有显著提升。
  2. 实现媲美全精度的训练-推理一致性：在训练阶段引入伪量化（Fake Quantization）使权重适应 INT4 分布，在推理阶段采用 W4A16（INT4 权重，BF16 激活）并依赖 BF16 Tensor Cores 进行计算，从而达到了与 BF16 全精度训练几乎一致的稳定性与行为表现。
  3. 开源生态的全链路工程实现：在 slime 框架上打通了从训练、权重格式转换到 SGLang 高效推理的完整 QAT INT4 闭环，并集成了多轮 VLM/LLM 训练、Rollout Router Replay 等前期研究成果，为极大规模模型的低成本强化学习训练提供了工业级解决方案。

整体架构与设计原则

构建端到端 QAT INT4 闭环流水线。系统实现了一个从训练到推理的完整 INT4 QAT 闭环。在QAT 训练阶段，训练侧在内存中始终保持 BF16 格式的主权重（Master Weights），但在前向传播中通过伪量化（Fake Quantization）引入量化噪声。所谓“伪量化”，是指系统并未真正将 BF16 张量转换为低精度的 INT4 物理存储，而是保留浮点计算路径，通过插入 QDQ（量化-反量化，Quantize-DeQuantize） 操作来模拟低精度算术。具体而言，高精度权重首先被“离散化”到 INT4 范围，随后立即被还原。尽管物理数据类型仍为浮点型，但其数值精度已实质性降低。原始值与还原值之间的差异引入了量化误差，这在数学上等同于向网络中注入噪声，迫使模型通过梯度更新来适应这种精度损失。对于反向传播，由于量化的核心操作是取整（Rounding）（一种几乎处处导数为 0 的阶跃函数，会完全阻断梯度流并阻止底层主权重的更新），系统采用了 STE（直通估计器，Straight-Through Estimator） 来绕过量化的不可导性。STE 采用了一种“直通梯度估计”策略：在反向传播期间，它将取整操作的导数定义为 1（即视为恒等映射）。这使得梯度能够顺利穿过取整层并更新高精度的浮点权重，从而完成 QAT 训练的闭环。在权重转换阶段，系统导出收敛后的 BF16 权重并执行真实量化（Real Quantization），将其转换为适合推理引擎（如 Marlin）的 INT4 格式。在 RL Rollout 阶段，SGLang 加载 INT4 权重并运行高效的 W4A16 推理（INT4 权重 × BF16 激活）。生成的经验数据随后流回第一阶段用于下一次 RL 训练迭代，形成自洽的闭环。
端到端 QAT INT4 流程

量化格式与训练策略的关键选择。在量化格式方面，系统遵循了 Kimi-K2-Thinking 的实践，选择了 INT4 (W4A16)。与 FP4 相比，INT4 在现有硬件（Pre-Blackwell 架构）上具有更广泛的支持，且生态系统中已经存在成熟且高性能的 Marlin 内核。实验表明，在 $1 \times 32$ 的缩放粒度（Scale Granularity）下，INT4 提供了足够的动态范围和稳定的精度，同时其性能和工具链已经得到了充分优化。作为工业界“足够好”的量化标准，INT4 在性能、风险和维护成本之间取得了合理的平衡。在训练策略方面，系统采用了经典的伪量化 + STE组合，通过维护 BF16 主权重、在前向传播中模拟量化噪声，并在反向传播中让梯度直通，最大化了低精度训练的收敛性和稳定性。

核心方法与实现细节

训练侧：在 Megatron-LM 中改造伪量化

实现伪量化与直通估计器（STE）。该阶段的核心目标是在训练过程中动态模拟量化误差，迫使模型“学习”适应低精度表示。因此系统采用了伪量化：权重在存储和更新时保持为 BF16 格式，但在前向传播中被临时映射到 INT4 的精度范围内。在具体工程实现上，核心逻辑被封装在 megatron/core/extensions/transformer_engine.py 的 _FakeInt4QuantizationSTE 类中。该类基于每个 Group 的最大绝对值执行动态量化，模拟 INT4 的 [-7, 7] 数值范围和截断（Clipping）行为，但底层计算依然在 BF16 下进行，仅注入量化误差。针对至关重要的反向传播，系统引入了 STE，确保梯度能够原封不动地穿过量化层，直接用于更新主权重，从而保持训练过程的连续性。
训练侧伪量化与 STE

伪量化消融实验与一致性分析。为了验证 QAT 的必要性并研究训练-推理精度不匹配带来的影响，系统设计了两种非对称场景的消融实验：一是开启 QAT INT4 训练但使用 BF16 进行 Rollout；二是关闭 QAT 训练直接使用 INT4 进行 Rollout。系统使用对数概率绝对差（Logprob Abs Diff）来衡量训练与推理的不一致性。左图展示了“QAT INT4 训练 + BF16 Rollout”的结果，即使使用了高精度的 BF16 推理，误差依然显著偏高。这是因为 QAT 已经通过“补偿”机制使权重适应了 INT4 的量化噪声，如果在推理时移除量化，这种补偿就会变成一种扰动，导致分布偏移（Distribution Shift）。右图展示了“无 QAT 训练 + 直接 INT4 Rollout”的结果（对应传统的训练后量化 PTQ 设置），由于模型在训练期间从未见过量化噪声，将权重压缩到 INT4 会导致严重的信息丢失，并使特征分布偏离训练时的状态，导致误差随训练步数震荡并放大。这些结果有力地证明了训练侧的伪量化与推理侧的真实量化必须联合开启。只有当训练中模拟的噪声与推理中的真实量化严格对齐时，才能抑制训练-推理的不匹配，避免分布偏移，并将误差保持在基线水平附近，从而真正闭环端到端的低精度 RL 训练。
Rollout BF16 训练侧 QAT INT4 效果对比
Rollout INT4 weight-only 训练侧 QAT INT4 效果对比

权重更新阶段

权重流转与动态格式适配。为了复用 SGLang 中现有的推理侧优化，系统直接采用了其内置的 Marlin kernel 作为 INT4 的推理后端。但在实际操作中存在显著的“格式鸿沟”：QAT 训练输出的权重是标准格式（类似于 Hugging Face 格式），而 SGLang 的 Marlin 内核要求权重必须经过特殊的打包（Packed）和重排（Permuted）才能被高效读取。由于 RL 训练需要频繁更新权重，必须解决格式兼容性问题。为此，系统设计了逆向的 restore_weights_before_loading 安全机制。该机制利用缓存的 _original_shapes 元数据，在任何权重更新发生之前，将内存中的 Marlin 权重格式还原（Resize）回原始形状。这防止了由于形状不匹配导致的运行时错误，并实现了标准权重格式与 Marlin 格式之间的平滑切换。系统还增加了一个系统级的 post_process_weights API，允许控制面根据训练进度显式触发此过程。为了解决加载后的格式适配问题，系统在 compressed_tensors_moe.py 中实现了动态权重管理机制。在权重加载完成后，系统自动运行 process_weights_after_loading，调用如 gptq_marlin_moe_repack 和 marlin_moe_permute_scales 等算子，在内存中将标准权重转换为高度优化的 Marlin 格式，从而最大化推理时的显存访问和计算效率。
SGLang 侧权重处理流水线

权重更新时的真实量化操作。这是核心的真实量化步骤。与训练时的伪量化不同，此步骤通过 int4_block_quantize 不可逆地压缩了精度：根据配置的 Group Size 计算每个 Group 的 Scale，并将高精度浮点数映射到 INT4 的整数域 [-7, 7] 中。为了最大化显存效率，系统随后执行了比特打包（Bit Packing）。由于 PyTorch 缺乏原生的 INT4 数据类型，系统利用位操作技巧实现了 pack_int4_to_int32，将 8 个 INT4 值紧密打包到一个 INT32 整数中（即 $8 \times 4 \text{ bits} = 32 \text{ bits}$）。最后，这些打包后的权重连同 Scales 一起传递给推理引擎，完成了从“训练格式”到“推理格式”的转换。
权重更新

推理阶段

极简打包与近零开销的解包机制。在 RL Rollout 期间，系统直接复用了 SGLang 成熟的 W4A16 量化方案。SGLang 通过将两个 4-bit 值打包到一个字节中，以紧凑的 INT4 格式存储权重，与 BF16 相比节省了 75% 的显存。在推理时，Triton 内核使用位操作（>> 4 和 & 0xF）进行高效解包。得益于计算与 IO 的重叠，这种解包操作几乎是零开销的。
SGLang W4A16 推理

针对 MoE 算子的深度融合优化。在显存优化方面，SGLang 引入了动态的 moe_align_block_size，它根据当前的 Token 数量和专家（Expert）分布选择合适的 block_size，对同一专家的 Token 进行分组和对齐，以提高带宽利用率。在计算融合方面，除了集成高性能的 Marlin INT4 实现外，SGLang 还将门控（Gating）部分融合到一个单一的高性能内核中，避免了重复的内核启动和中间数据的读写。这种 INT4 推理方案兼容 GPTQ 和 AWQ 等主流格式，并同时支持对称和非对称模式。

实验环境

  • 数据集：训练阶段使用 dapo-math-17k 数据集；评估阶段使用 aime-2024 基准测试。
  • 模型架构与规模：验证模型包括 Qwen3-30B-A3B、Qwen3-235B-A22B 以及参数量约 1TB 规模的 Kimi-K2-Thinking 模型。
  • 硬件配置：NVIDIA H 系列 GPU（Pre-Blackwell 架构）。特别指出，单节点实验使用了一台配备单张 H200（141GB 显存）的机器来容纳 1TB 级别的模型。
  • 软件配置：基于 slime 框架、SGLang 推理引擎、Megatron-LM 进行底层改造，依赖 Triton 和 Marlin INT4 高性能内核。

实验结果

训练侧 Raw-Reward 收益表现。系统在 dapo-math-17k 数据集上对比了 Qwen3-235B-A22B 和 Kimi-K2-Thinking 在不同配置下的训练表现。结果显示，与“BF16 训练–BF16 推理”和“BF16 训练–FP8 推理”相比，“BF16 训练–INT4 推理”配置依然实现了稳定的 Raw-Reward 增长，且其增长趋势与前两者高度一致，验证了该方案在训练侧的有效性。（详见 Fig 8, Fig 9）

评估侧 AIME 性能表现。为了更严谨地评估模型能力，系统每隔 10 个训练步在 aime-2024 基准上进行一次评估。实验表明，“BF16 训练–INT4 推理”方案不仅在评估分数上呈现稳定的上升趋势，而且在斜率和峰值分数上与“BF16 训练–BF16 推理”及“BF16 训练–FP8 推理”紧密重合。这种高度对齐表明，低比特量化并未损害模型的核心表示能力，在大幅节省计算资源的同时，保留（或匹配）了全精度的泛化性能。（详见 Fig 10, Fig 11）

训练-推理差距（Train-Infer Gap）验证。系统在 Qwen3-30B 和 Qwen3-235B 上验证了输出的一致性（通过训练侧与推理侧输出的对数概率绝对差来衡量，值越低一致性越强）。结果显示，INT4 的误差曲线几乎与 BF16 基线重合，且显著低于 FP8 配置。这证实了 INT4 QAT 能够有效避免“BF16 训练–FP8 推理”模式下的精度损失，实现与全精度无异的训练-推理行为。研究团队推测一致性极高的原因有两点：一是训练侧的伪量化将权重约束在 INT4 范围内，有效抑制了并行矩阵乘法中非确定性累加顺序引起的浮点舍入误差（即“大数加小数”的精度损失）；二是推理采用 W4A16，全程依赖 BF16 Tensor Cores，保持了与训练高度对齐的计算精度。（详见 Fig 12, Fig 13）

Rollout 加速效果分析。在 Qwen3-235B 的 Rollout 性能对比中，INT4 和 FP8 相较于 BF16 基线均有显著加速，但 INT4 与 FP8 之间的差距不大。这是因为当前 NVIDIA H 系列 GPU 缺乏原生的 INT4 Tensor Cores，W4A16 本质上仍使用 BF16 Tensor Cores 进行计算，虽然大幅降低了显存带宽压力，但无法像 W8A8 那样获得原生 FP8 Tensor Cores 的计算提升。对于 Kimi-K2-Thinking 模型的 Rollout，在双节点场景下，由于跨节点带宽成为瓶颈，FP8 和 INT4 表现相近；但在单节点场景下，INT4 展现了其显存压缩的核心价值：通过将模型体积减半，约 1TB 规模的模型能够完全加载到单台机器的显存中，彻底消除了昂贵的跨节点通信，极大地缩短了 Rollout 时间。这有力地证明了在当前硬件下，INT4 QAT 的主要收益在于通过显存压缩实现高效的单节点 Rollout。（详见 Fig 14, Fig 15）

结论与未来工作

本文通过在开源框架中复现 Kimi 团队提出的 INT4 QAT 方案，验证了其在极大参数量模型 RL 训练中的有效性：不仅在精度上完美复现了媲美 BF16 基线的优势，还在效率上显著提升了 Rollout 吞吐量，证明了低比特量化在强化学习中的巨大价值。

针对未来的工作方向，团队提出了以下展望：

  1. 训练侧效率优化：目前，引入 QAT 伪量化在训练期间增加了额外的计算开销，导致训练速度明显慢于纯 BF16 训练，这在一定程度上抵消了 Rollout 加速带来的端到端收益。团队计划提出新的优化方案来解决这一训练侧瓶颈，进一步加速全流水线。
  2. 推理侧 FP4 探索：随着 NVIDIA Blackwell 架构的普及，团队将积极探索 FP4 精度在 RL 训练和推理中的应用，以进一步挖掘硬件潜力。

补充细节

致谢：本项目的成功落地得益于多个团队的联合贡献，包括 SGLang RL Team（Ji Li, Yefei Chen, Xi Chen, BBuf, Chenyang Zhao）、InfiXAI Team（Mingfa Feng, Congkai Xie, Shuo Cai）、Ant Group Asystem & AQ Infra Team（Yanan Gao, Zhiling Ye, Yuan Wang, Xingliang Shi）、slime Team（Zilin Zhu, Lei Li, Haisha Zhao）以及 RadixArk Team。同时，团队特别感谢 Verda Cloud 提供的计算资源赞助。相关功能和技术配方已同步至 slime 和 Miles 社区。