Marlin

这是 Marlin（Mixed Auto-Regressive Linear kernel 的缩写，也是一种地球上速度最快的鱼的名字），一个极度优化的 FP16xINT4 矩阵乘法内核，专为 LLM 推理设计。它能够在批次大小达到 16-32 个 token 时提供接近理想（4倍）的加速效果（相比之下，先前具有类似加速效果的工作仅支持 1-2 个 token）。这使得 Marlin 非常适合大规模服务、推测解码或高级多推理方案，如 CoT-Majority。

核心技术：

大多数现代 GPU 的 FLOP 与字节比约为 100-200。因此，只要我们对每 4 位量化权重执行少于 25-50 次（张量核心）乘加运算，理论上就有可能维持相对于 FP16 权重接近理想的 4 倍加速。这意味着，权重量化带来的全部性能优势，原则上应能扩展到比现有内核当前支持的批次大小大 4-8 倍的范围。

然而，在实践中实现这一点极具挑战性，因为我们本质上需要同时充分利用所有可用的 GPU 资源（全局内存、L2 缓存、共享内存、张量核心、向量核心）。Marlin 通过多种技术和优化实现了这一点，简要概述如下：

• 计算组织：我们以这样的方式组织计算，确保所有激活值基本上始终从 L2 缓存中获取，并在寄存器中进一步重用多次，以防止从共享内存的重复加载成为瓶颈。
• 异步全局权重加载：我们异步执行全局权重加载，使其与所有计算操作以及激活值加载重叠，并采用允许立即逐出的缓存策略，避免将不再重用的值不必要地污染 L2 缓存。
• 共享内存双缓冲：由于激活值数据量相对较大，共享内存的占用相当显著。我们通过双缓冲技术来加载共享内存，使其与计算和全局加载重叠。
• 指令流水线饱和：我们精心安排反量化和张量核心指令的顺序，确保两个 GPU 流水线都得到良好饱和，不会相互阻塞。
• 离线数据重排：通常，量化权重和分组缩放因子都会离线重排成一种在执行期间具有理想访问模式的布局，例如，允许直接将权重反量化到张量核心的组织格式中。
• 多线程块协同计算：我们让线程块中的多个线程束计算同一输出分块的部分结果，以实现更高的线程束数量，最大化计算和延迟隐藏，而无需增加输出分块大小（增加分块大小会使在实际矩阵上进行良好分区变得困难）。
• 高效内存访问：所有加载都使用最大向量长度以实现峰值效率。我们还执行多种布局转换，以确保所有共享内存的读写操作都是无冲突的（特别是对于矩阵加载指令），并且全局归约以最小的内存开销进行。
• 静态内存偏移：我们设置并展开循环，使大多数内存偏移是静态的，从而最小化运行时的索引计算。
• 条带化分区方案：我们实现了一种“条带化”分区方案，每个 SM 处理的分块段可能（部分）跨越多个列“切片”。这可以在大多数矩阵形状上实现良好的 SM 利用率，同时最小化所需的全局归约步骤。
• 高效全局归约：全局归约直接在输出缓冲区中进行（临时将 FP32 累加器降级为 FP16），该缓冲区保持在 L2 缓存中；归约操作也经过优化，以避免任何不必要的读写。
• 汇编级优化：我们使用 NSight-Compute 对内核的 PTX 汇编进行了广泛分析。CUDA 代码中还包含了一些看似冗余或略微次优的结构，但它们能编译成更快的 PTX。

性能基准测试：

我们首先在一个可以在 NVIDIA A10 GPU 上理想分区的大型矩阵上，比较 Marlin 与其他流行的 4 位推理内核的性能。这允许所有内核达到其最佳性能。所有内核均在分组大小 128 下执行（但请注意，缩放因子的格式并非 100% 相同）。

现有内核在批次大小为 1 时实现了相对接近最优的 3.87 倍加速（注意分组缩放因子有 0.125 位的存储开销），但随着输入数量的增加，其性能迅速下降。相比之下，Marlin 在所有批次大小下都提供了基本理想的加速，在批次大小约为 16-32 时仍能实现最大可能的 3.87 倍加速。

得益于其条带化分区方案，Marlin 在真实（较小）的矩阵和各种 GPU 上也表现出强劲的性能。下面的结果证明了这一点，我们在批次大小为 16 的情况下，对流行开源模型的 Transformer 块中所有线性层的总运行时间进行了基准测试。

最后，我们还研究了在锁定 GPU 基础时钟频率的情况下，长时间内可以维持的性能。有趣的是，我们发现降低时钟频率会显著损害先前内核的相对加速比，但对 Marlin 几乎最优的性能（相对于较低的时钟设置）没有影响。

环境要求：

• CUDA >= 11.8（特别是 nvcc 编译器的版本应与 torch 匹配）
• NVIDIA GPU，计算能力 >= 8.0（Ampere 或 Ada 架构，Marlin 尚未针对 Hopper 架构进行优化）
• torch>=2.0.0
• numpy
  运行量化脚本还需要：
• transformers
• datasets
• sentencepiece

使用方法：

如果满足所有要求，可以通过在此仓库的根目录下调用以下命令来安装 Marlin：

pip install .

安装后，使用 Marlin 内核最简单的方式是通过 marlin.Layer，这是一个代表 Marlin 量化层的 torch 模块。它允许通过 marlin.Layer.pack(linear, scales) 将一个“伪量化”（反量化值以 FP16 存储）的 torch.Linear 层转换为压缩的 Marlin 格式。或者，如果权重和缩放因子已经过适当的预处理（参见 marlin.Layer.pack(...)），也可以通过 marlin.mul(..) 直接调用内核。内核本身位于独立的 marlin/marlin_cuda_kernel.cu 文件中，该文件除了基础 CUDA 外不包含任何依赖项，因此应该很容易集成到其他低级框架中。

可以通过 python test.py 运行正确性测试，通过 python bench.py 运行基准测试。请注意，为了复现我们的“可持续性能”基准测试，需要使用以下命令将 GPU 时钟锁定在各自的基础值：

sudo nvidia-smi --lock-gpu-clocks=BASE_GPU_CLOCK --lock-memory-clocks=BASE_MEM_CLOCK

此外，如果启用了 ECC（例如在 A10 上），则最大可实现的内存带宽将比官方规格表低 10-15%，因为每个内存请求都会包含校验和开销。可以通过以下命令禁用 ECC：

sudo nvidia-smi -e 0

我们在 A10 基准测试中就是这样做的。

GPTQ 示例：

在 gptq 子文件夹中，我们还提供了 GPTQ 算法的一个略微改进的版本，具有更好的分组网格裁剪和非均匀校准样本长度，可以生成与 Marlin 兼容的 Llama2 模型的 4 位版本。此外，还有一个脚本可以在流行的 LLM 评估工具 中使用 Marlin 内核评估此类压缩模型。下面的脚本使用 lm-eval-harness==0.4.0 进行了测试，可能不适用于更新或更旧的版本。以下是相应的示例命令（必须安装 marlin、transformers 和 datasets 包）：

% 压缩 Llama2 模型并以 Marlin 格式导出模型。
python llama2.py LLAMA2_CHECKPOINT --wbits 4 --save checkpoint.pt
% 对未压缩模型执行困惑度评估。
python llama2.py LLAMA2_CHECKPOINT
% 在评估工具中使用 Marlin 内核评估压缩模型。
python eval.py --model hf --model_args pretrained=LLAMA2_CHECKPOINT --tasks mmlu \
  --marlin_checkpoint checkpoint.marlin.g128
% 评估全精度基线模型。
python eval.py --model hf --model_args pretrained=LLAMA2_CHECKPOINT --tasks mmlu 

我们测量了 4 位（分组大小=128）Marlin 模型的以下 WikiText 和 Red-Pajama 困惑度，以及 MMLU 零样本准确率：

请注意，此 GPTQ 示例目前主要旨在演示如何生成准确的 Marlin 模型，并作为内核正确性的端到端验证（而不是作为一个灵活的压缩工具）。

引用：

如果您觉得这项工作有用，请考虑引用：

@article{frantar2024marlin,
  title={MARLIN: Mixed-Precision Auto-Regressive Parallel Inference on Large Language Models},
  author={Frantar, Elias and Castro, Roberto L and Chen, Jiale and Hoefler, Torsten and Alistarh, Dan},
  journal={arXiv preprint arXiv:2408.11743},
  year={2024}
}