ScaleBITS: Scalable Bitwidth Search for Hardware-Aligned Mixed-Precision LLMs¶
会议: NeurIPS 2026
arXiv: 2602.17698
OpenReview: https://openreview.net/forum?id=TpYhEftiHb
领域: 模型压缩
关键词: 大语言模型量化、混合精度分配、贪心位宽搜索、双向通道重排、硬件对齐分块
一句话总结¶
ScaleBITS 通过双向通道重排将敏感权重聚拢并映射为硬件对齐的计算块，提出基于量化参考点一阶梯度的可扩展批量贪心搜索算法，在不引入额外推理延迟的前提下实现了 2–3 bit 极限低比特大模型的高性能混合精度分配。
研究背景与动机¶
随着大语言模型（LLM）参数规模的急剧扩大，训练后仅权重（Weight-Only）量化已成为降低显存占用与推理成本的核心手段。在 4-bit（如 INT4、FP4）精度下，现有量化算法已能较好保持模型性能；然而一旦将平均位宽推向 4-bit 以下（如 2-bit、3-bit 极限低比特区间），均匀精度量化（Uniform Precision）便会遭遇严重的性能崩溃。这一瓶颈的根源在于大模型权重对量化误差的敏感度分布高度不均——极小比例（0.1%–1%）的关键权重主导了整体性能退化，而绝大多数权重对激进量化具备极高容忍度。
现有的应对方案呈现出两种极端的权衡困境。一方面，以 SpQR、SqueezeLLM、PB-LLM 为代表的细粒度/非结构化混合精度方案，将少量敏感权重保留在 FP16，虽然精度较高，但非规则的离群点稀疏存储和索引开销严重破坏了硬件访存局部性，导致硬件利用率低下；中间粒度的方案（如 OWQ、SlimLLM）虽按列分组保持了规则性，却无法灵活适配局部敏感度，且搜索空间被局限在单一层内的相邻位宽平衡。另一方面，现有的位宽搜索策略（基于可微搜索、强化学习或遗传算法）在十亿级参数量上计算复杂度过高，难以全局优化。贪心搜索虽具备 \((1-1/e)\) 近似保证，但在全模型数百万个权积分块上直接枚举评估边际损失需要 \(O(N^2)\) 次全模型前向，计算上完全不可行。此外，传统方法在全精度模型处求一阶或二阶导数估计敏感度，忽视了激进量化后损失曲面发生的剧烈偏移，导致敏感度排序严重失真。
针对上述硬件执行割裂与搜索可扩展性缺失的双重矛盾，本文的核心切入点是：利用硬件矩阵乘算子的分块（Tiling）执行规律，将敏感度聚类约束在规则计算块内部；同时利用量化后非局部最优点的特性，用轻量的一阶梯度动态近似贪心搜索的相对边际增益。核心 idea：通过双向通道重排将分散的高敏感权重聚拢于矩阵对齐分块内，并在递进量化点上利用一阶梯度与批量更新逼近全局贪心搜索，实现无硬件运行时开销的超低比特混合精度分配。
方法详解¶
整体框架¶
ScaleBITS 的整体流水线包含三个紧密协作的阶段：首先利用全局递进量化状态评估权重的行、列敏感度，执行双向通道重排（Bi-directional Channel Reordering），在保持网络数学等价性的前提下将高敏感权重聚拢到矩阵左上角；随后，将重排后的权重矩阵划分为与硬件 GEMM Kernel 瓦片（Tile）严格对齐的规则分块（如 \(64 \times 128\)），消除块内控制流发散；最后，建立全局内存预算约束下的离散优化目标，通过基于一阶敏感度近似与两阶段批量更新的可扩展贪心搜索算法（Scalable Greedy Search），自动化求解各分块的最优位宽配置。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["预训练浮点模型权重大矩阵"] --> B["阶段 1：双向通道重排<br/>行列敏感度聚合 + 耦合层等价重排"]
B --> C["阶段 2：硬件对齐分块划分<br/>匹配 GEMM Tile (如 64×128)"]
C --> D["阶段 3：可扩展贪心搜索<br/>量化参考点一阶梯度 + 批式位宽更新"]
D --> E["端到端混合精度量化大模型<br/>融合 Triton 算子无开销推理"]关键设计¶
1. 量化参考点动态一阶敏感度：克服全精度曲面偏移的动态度量
传统混合精度方法通常在全精度权重 \(\bm{w}\) 附近进行泰勒展开，并假设梯度接近零而依赖高代价的二阶曲率（Hessian 或 Fisher 信息矩阵）。然而在极限低比特（2–3 bit）场景下，模型已大幅偏离原始浮点解空间，静态的二阶估计会导致层级与块级敏感度排序严重失效。ScaleBITS 转而在当前的量化参考点 \(\bm{w}^Q\) 附近展开损失函数。由于量化后模型并非局部极小点，一阶梯度项占据主导，使得块分量 \(\bm{w}_i\) 的敏感度可高保真地简化为： $\(s_i := |L(\bm{w}^Q + \Delta \bm{w}_i) - L(\bm{w}^Q)| \approx |\bm{g}^{(\bm{w}^Q_i)\top} \Delta \bm{w}_i|\)$ 对于线性层矩阵内部元素 \(W_{ij}^Q\)，其梯度可解析分解为输出行梯度与输入列激活的外积 \(g_{ij}^{(\bm{W}^Q)} = g_i^{(\bm{y})} x_j\)，进而导出元素级敏感度公式： $\(s_{ij} = |g_i^{(\bm{y})}| \cdot |x_j| \cdot |\Delta W_{ij}|\)$ 该式揭示了敏感权重在行（输出通道梯度）和列（输入通道激活幅值）两个维度上的双向富集机理。
2. 双向通道重排与硬件对齐分块：将离散敏感权重聚拢入规整算子
为了消除非结构化量化带来的索引稀疏性和计算分支，ScaleBITS 将权重划分形式约束为规整的二维子矩阵块 \(\bm{W}_i \in \mathbb{R}^{m \times n}\)（实验采用 \(64 \times 128\)），保证每个块在 Tensor Core 瓦片加载时以统一精度进行解量化与计算。由于敏感权重天然离散分布于各个通道中，直接分块会导致敏感性被稀释。ScaleBITS 采用 \(\ell_1\) 范数沿行与列分别聚合元素敏感度，对权重矩阵的输入通道（列）与输出通道（行）实施双向重排，将绝大多数极敏感参数聚拢至矩阵左上角区域。为了严格保持网络的数学函数等价性，重排在网络耦合路径上统一协调：残差主干（Residual Stream）通道采用跨层全局聚合的联合排列，并同步重排嵌入层输出；MLP 模块内部的 gate、up 与 down 投影遵循局部闭环重排；注意力模块则在保留 RoPE 约束的前提下在每个注意力头内部独立重排 Value 与 Output 投影。所有排列在量化前离线固化，推理期间零额外重排开销。
3. 可扩展批量贪心搜索：高保真逼近全局单调次模优化的快速算法
将位宽配置建模为全局比特预算 \(B\) 下的整数规划问题：\(\min_{\bm{b}} L(Q(\bm{w}, \bm{b}))\)，s.t. \(\frac{1}{N}\sum_i b_i \le B\)。在次模性与单调性假设下，经典贪心搜索每步仅提升 1 个块的位宽，在百万级分块的 LLM 上具有 \(O(N^2)\) 的灾难性复杂度。ScaleBITS 提出了带有两阶段批量更新与拒绝接受机制的可扩展贪心算法（Scalable Greedy Search）。在位宽低于预算时，算法执行纯扩张更新，按上升边际收益 \((s_{\text{up}})_i = \bm{g}^{(\bm{w}_i^{Q_{b_i}})\top} (\bm{w}_i - \bm{w}_i^{Q_{b_i}})\) 选出 top-\(k\) 个敏感块批量增加 1 bit；在预算达到后，执行平衡置换更新，将 top-\(k/2\) 敏感块提升 1 bit，同时将基于步长幅值衰减指标 \((s_{\text{down}})_i = \frac{1}{2^{b_i}} \|\bm{g}^{(\bm{w}_i^{Q_{b_i}})} \odot \bm{w}_i^{Q_{b_i}}\|_1\) 选出的 bottom-\(k/2\) 不敏感块降低 1 bit。引入验证集损失校验机制：若批量更新导致前向损失上升，则回滚更新并将步长折半（\(k \leftarrow k/2\)），兼顾搜索全局收敛速度与数值稳定性，将搜索时间从数万小时压缩至 1 小时内。
损失函数 / 训练策略¶
ScaleBITS 属于训练后量化（PTQ）体系，不涉及长周期反向传播或参数重训。敏感度评估与批量贪心搜索采用来自 RedPajama 的 128 条长度为 4096 的校准序列，通过前向与反向传播获取梯度向量 \(\bm{g}\)。搜索空间设定为 \(\mathcal{B} = \{1, 2, \dots, 8\}\) bit，并采用温启动策略：初始化为底部位宽 \(\bm{b}^{(0)} = \lfloor B \rfloor\)，避免模型在全 0-bit 或全 1-bit 下梯度崩溃。量化算子采用标准的分组近邻舍入（RTN，group size 128），并结合定制开发的 Triton 融合解量化 GEMM 算子进行无分支硬件加速。
实验关键数据¶
主实验¶
在 LLaMA-2-7B、LLaMA-3-8B 及 70B 等模型上，对 2–3 bit 极限低比特量化开展系统评测。Wiki2 指标采用 WikiText-2 困惑度（Perplexity，越低越好），0-Shot 为 6 个常识问答任务平均准确率（%，越高越好），MMLU 为 5-shot 准确率（%，越高越好）。基准涵盖均匀精度标量量化（RTN、GPTQ）、非均匀量化（GuidedQuant）以及先进混合精度方法（SlimLLM）。
| 模型与配置 | 目标位宽 | 方法分类 | WikiText-2 PPL (↓) | 0-Shot 准确率 (↑) | MMLU (↑) |
|---|---|---|---|---|---|
| LLaMA-3-70B (FP16) | 16.0 | 浮点上限 | 2.59 | 76.88% | 78.74% |
| LLaMA-3-70B | 3.1 | RTN-g128 (均匀) | 19.35 | 66.33% | 54.56% |
| LLaMA-3-70B | 3.1 | GPTQ-g128 (均匀) | 4.75 | 73.72% | 74.51% |
| LLaMA-3-70B | 3.0 | GuidedQuant (非均匀) | 3.90 | 75.51% | 76.86% |
| LLaMA-3-70B | 3.1 | ScaleBITS + RTN (本文) | 3.69 | 75.00% | 76.88% |
| LLaMA-3-70B | 2.1 | RTN-g128 (均匀) | 2e4 | 35.61% | 23.02% |
| LLaMA-3-70B | 2.1 | GPTQ-g128 (均匀) | 28.80 | 46.60% | 26.63% |
| LLaMA-3-70B | 2.0 | GuidedQuant (非均匀) | 10.21 | 64.39% | 45.55% |
| LLaMA-3-70B | 2.1 | ScaleBITS + RTN (本文) | 7.80 | 64.72% | 58.82% |
| LLaMA-3-8B | 2.1 | SlimLLM-g128 (混合) | 66.06 | 37.13% | 25.51% |
| LLaMA-3-8B | 2.1 | ScaleBITS + RTN (本文) | 11.80 | 58.53% | 39.36% |
消融实验¶
针对敏感度梯度更新方式、双向重排及搜索效率组件开展系统消融（以 LLaMA-3-8B 量化与单张 H100 GPU 搜索/推理为基准）：
| 配置 / 变体 | 核心改动说明 | 搜索用时 / 迭代轮数 | WikiText-2 PPL (3.1 bit) | 8192×8192 算子延迟 (BS=16) |
|---|---|---|---|---|
| ScaleBITS 完整方案 | 自适应量化梯度 + 双向重排 + 批量贪心 | 0.5–1.0 h / 16–36 轮 | 6.63 | 58.0 μs |
| 固定初始梯度 (w/o Adaptive) | 仅第一轮计算梯度，后续迭代冻结敏感度 | 0.4 h / 20 轮 | 7.95 | 58.0 μs |
| 禁用通道重排 (w/o Reorder) | 权重沿原始物理通道直接分块量化 | 0.8 h / 28 轮 | 9.42 | 58.0 μs |
| 单向仅列重排 (Input Only) | 仅依输入激活重排，忽略行梯度聚类 | 0.7 h / 25 轮 | 7.81 | 58.0 μs |
| 经典全量贪心 (Classic Greedy) | 单步单块迭代，精确 forward 评估增益 | \(\approx 10^{10}\) h / \(\approx 3\times 10^6\) 轮 | （计算不可行） | 58.0 μs |
| 均匀 INT4 基准 (Uniform INT4) | 全矩阵统一量化为 4-bit 瓦片 | 0.05 h (无搜索) | 5.82 | 57.8 μs |
| 浮点 BF16 基准 (CUTLASS) | 未压缩原始全精度基准 | 0 h | 5.54 | 86.0 μs |
关键发现¶
- 双向重排打破硬件结构性惩罚：消融证明双向通道重排使分块敏感度方差急剧拉大，使得仅占全局极小面积的左上角核心块能够“借用”低敏区域腾出的位宽配额，PPL 相比无重排直降近 3 点，验证了 2D 聚类在规则分块中的核心地位。
- 量化参考点梯度的抗失真优势：随着位宽压缩至 2-bit，损失地形严重重塑。动态更新梯度的方案显著优于静态全精度梯度，避免了在错误参数方向上过度分配位宽。
- 混合精度零延迟损耗落地：基于 Triton 编写的硬件对齐 GEMM 算子测试显示，在平均 4-bit 混合精度（40% INT2 + 40% INT4 + 20% INT8）下，推理时延为 58.0 μs，与纯均匀 INT4 算子（57.8 μs）几乎完全一致，相比 BF16（86.0 μs）实现 1.48 倍吞吐加速。
亮点与洞察¶
- 硬件友好的二维敏感度聚类：突破了以往研究仅在输入通道（列）进行激活重排的思维定势，敏锐指出矩阵梯度的外积性质决定了权重敏感度在行、列双向均呈现高度偏置，使硬件规则分块兼得细粒度灵活性与计算规整度。
- 量化曲面一阶占优的精妙简化：巧妙指出模型量化后脱离了局部极小值驻点，从而使极其昂贵的二阶 Hessian 矩阵退化为次要项，单次一阶反向传播即可精准预测边际损失排序，将搜索时间从不可行降至 1 小时。
- 连续细粒度的 Pareto 前沿覆盖：解除了传统方案仅能支持离散 2/3/4 bit 或层内固定比例的机械约束，使工程师可以面向任意任意连续显存预算（如 2.37 bit）一键产出理论最优分配模型。
局限与展望¶
- 算子适配依赖定制 Triton 实现：虽然分块对齐了硬件 Tile，但现存深度学习框架默认的量化算子库（如 TensorRT-LLM、vLLM）大多仅支持均匀整型位宽，混合位宽 Block-GEMM 仍需专属底层算子支撑方可产品化落地。
- 仅聚焦权重端未涉及激活值混合量化：当前方法主要面向 Weight-Only 场景（W2A16 / W3A16），在大 Batch 显存受激活主导时，尚未支持激活值维度的混合精度联合搜索。
- 极限 2-bit 下缺少轻量参数微调补偿：当前完全基于免训练 PTQ 机制，若能在搜索完成后配合几百步 LoRA 或轻量重构微调，有望进一步拉升 2-bit 极限下的零样本推理泛化性。
相关工作与启发¶
- vs SlimLLM / OWQ: 相比于 SlimLLM 仅在层内按列划分固定比例相邻位宽，ScaleBITS 实现了跨层、跨模块的全局矩阵块级自由分配，且通过双向重排消除了对非结构化列稀疏性的依赖。
- vs GuidedQuant / GPTQ: 传统二阶敏感度优化致力于在固定位宽内优化舍入网格，ScaleBITS 证明了在超低比特下，合理的“位宽资源重分配”带来的收益远超固定位宽下的精细量化技巧，即简单的 RTN 搭配混合精度即可击败精密的均匀量化器。
- vs SqueezeLLM / SpQR: 克服了其依靠坐标稀疏列表保存离群点造成的访存开销与内存断裂，用硬件对齐分块与 Tensor Core 算子融合完全抹平了运行时推理惩罚。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [揭示了行-列双向通道敏感度结构并将其与 GPU GEMM Tile 对齐，构思新颖巧妙]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 LLaMA-2/3、Gemma-2 多尺度基准模型，并在真实 H100 硬件上交付了融合算子测试]
- 写作质量: ⭐⭐⭐⭐⭐ [数学推导严谨流畅，从单调次模理论到工程化批量近似的逻辑链条极其自洽]
- 价值: ⭐⭐⭐⭐⭐ [攻克了 2–3 bit 极限低比特大模型部署的核心精度瓶颈，极具产业落地工程价值]