论文

dMX:低精度浮点格式的可微分混合精度赋值

dMX: Differentiable Mixed-Precision Assignment for Low-Precision Floating-Point Formats

摘要

将 大语言模型 (LLM) 量化为低精度浮点表示是高效部署的核心,但在所有层上统一应用单个位宽度在性能和准确性方面都不是最佳的。这项工作介绍了 dMX,这是一种可微的混合精度量化框架,用于可学习的浮点位宽分配。我们研究其在开放计算项目 (OCP) 标准定义的微尺度浮点 (MXFP) 数据类型系列中的应用。每层位宽分配被表述为一个连续优化问题,其中每层的浮点格式由标量参数参数化,将多变量设计空间折叠为单个可学习的偏移量。在训练期间,此偏移量采用连续值,避免离散量化格式之间的突然振荡。基于温度的退火计划逐步离散化学习到的偏移量,确保最终配置映射到硬件兼容的 MXFP 格式,而不会在训练和推理行为之间突然转换。目标感知正则化项将平均位宽引导至用户指定的预算,充当推理成本的粗粒度代理并平衡模型质量与部署效率。我们对 LLM 的不同系列(例如 Llama、Qwen3 和 SmolLM2)进行了实验,评估了 WikiText-2 的困惑度和四个零样本推理基准的准确性。在这些设置中,dMX 始终生成 Pareto 主导模型,并改进了基于 Kullback-Leibler (KL) 散度的层选择启发式方法,有效地在模型质量和平均位宽之间进行权衡。