技术方案核心 - 成功落地了从训练到推理的完整INT4量化感知训练全流程方案 该方案采用“训练端伪量化 + 推理端真实量化”的组合 实现了与BF16全精度训练相媲美的稳定性与训推一致性 [2][3] - 通过INT4极致压缩 将1TB级别的超大模型采样任务容纳于单机H200的141GB显存内 消除了跨机通信瓶颈 显著提高了Rollout效率 [2] - 方案在开源生态中打通了全流程 为社区提供了兼顾高性能与低成本的可落地参考 [3] 技术方案细节 - 总体流程实现了从训练到推理的完整QAT INT4闭环 核心收益包括突破显存瓶颈、实现训推一致以及单机效率倍增 [6] - 在QAT训练阶段 训练侧维护BF16主权重 前向传播通过伪量化引入量化噪声 反向传播则利用STE技术确保梯度能够回传以更新主权重 [8][9] - 在权重转换阶段 将训练收敛的BF16权重导出并执行真实量化 转换为INT4格式 供推理引擎使用 [10] - 在推理阶段 由SGLang加载INT4权重并执行高效的W4A16推理 生成的经验数据回流用于下一轮RL训练 构成迭代闭环 [10] 核心策略选择 - 量化格式参考了Kimi-K2-Thinking 选用了INT4方案 主要考虑其在现有硬件上的支持更广泛 且有成熟的Marlin Kernel实现 [11] - 在1×32量化Scale粒度下 INT4动态范围充足、精度稳定 其性能与生态链路均已高度优化 [11] - 训练方法采用了Fake Quantization配合STE的经典组合 通过维护BF16主权重并在前向计算中模拟量化噪声 最大程度保证了低精度训练的收敛性与稳定性 [11] 训练侧改造与验证 - 在Megatron-LM中实现了Fake Quantization与STE机制 核心是在前向传播中将权重动态映射到INT4精度范围参与计算 但在反向传播时让梯度直接穿透量化层 [16][17] - 通过消融实验验证了QAT方案的必要性 实验证明训练端的Fake Quantization与推理端的Real Quantization必须协同开启 才能有效抑制训推不一致和分布偏移 [18][21] - 在“关闭QAT训练，直接进行INT4 Rollout”的传统PTQ模式下 误差随着训练步数呈现震荡上升的趋势 [21] 权重更新与格式适配 - 工程落地时遇到了“格式鸿沟” QAT训练产出标准格式权重 而SGLang推理引擎的Marlin Kernel要求权重必须经过特定的打包与重排处理 [25] - 设计了一套逆向的`restore_weights_before_loading`保护机制 确保模型能在标准权重格式与Marlin权重格式之间平滑切换 [27] - 实现了动态权重管理机制 在模型权重加载结束后自动触发流程 将标准权重即时转换为高度优化的Marlin权重格式 [28] 推理阶段优化 - SGLang在推理阶段复用了优化成熟的W4A16量化方案 使用紧凑的INT4格式 将两个4-bit权重打包进一个字节 相比BF16节省了75%的内存 [32] - 引入了动态的`moe_align_block_size` 根据当前Token数量和Expert分布自动选择block_size 提升显存带宽利用率 [35] - 将gating部分融合成一个高性能的kernel 避免了反复启动kernel和读写中间结果 [35] 训练与评估效果 - 在Qwen3-235B-A22B与Kimi-K2-Thinking模型上的训练实验表明 “BF16训 - INT4推”配置下的Raw-Reward能保持稳健增长 且增长趋势与“BF16训 - BF16推”及“BF16训 - FP8推”基本一致 [38][41] - 在aime-2024基准测试集上的评估显示 “BF16训 - INT4推”方案的评估分数呈现稳健上升态势 其性能提升的斜率与最终峰值 均与全精度及FP8方案保持了较高的重合度 [42][46] - 在Qwen3-30B与Qwen3-235B模型上的验证显示 INT4 QAT策略下的训推差异与BF16基准呈现出惊人的重合度 且显著低于FP8模式 [51] 性能表现分析 - 在Qwen3-235B的Rollout性能对比中 INT4与FP8均较BF16基线实现了显著加速 但由于NVIDIA H系列GPU没有原生的INT4 Tensor Core INT4在单步推理耗时上仅表现出微弱优势 与FP8基本处于同一性能梯队 [54] - 对于Kimi-K2-Thinking模型 在双节点场景下INT4与FP8性能相似 整体受限于跨节点通信带宽 但在单节点场景下 INT4通过将模型体积减半 成功将1TB级模型加载至单机显存 消除了跨机通信开销 大幅缩减了Rollout耗时 [56][61] - INT4 QAT方案在当前硬件环境下的核心收益在于通过压缩显存 解锁了高效的单机部署Rollout方案 [56] 未来工作方向 - 计划解决训练侧因引入QAT Fake Quantization计算带来的较大额外性能开销问题 旨在实现全链路的加速 [62] - 随着NVIDIA Blackwell架构的普及 将积极探索FP4精度在RL训练与推理中的应用可行性 [62]