Marlin

这是 Marlin，一个混合自回归线性（Mixed Auto-Regressive Linear）内核（也是地球上速度最快的鱼类之一的名字），这是一个针对大语言模型（LLM）推理的极度优化的 FP16xINT4 矩阵乘法（matmul）内核，能够在批量大小达到 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 流水线都得到充分饱和且不会相互成为瓶颈。
• 通常，量化权重和组缩放因子都会离线重新排列成一种布局，以便在执行期间提供理想的访问模式，例如允许直接将权重反量化为张量核心组织形式。
• 我们在一个线程块中使用多个 warp（线程束）来计算同一输出瓦片的局部结果，以实现更高的 warp 数量，最大化计算和延迟隐藏，而不增加输出瓦片的大小（否则会使现实矩阵的良好分区变得困难）。
• 所有加载使用最大向量长度以达到峰值效率，我们还执行几种布局转换，以确保所有共享内存读写无冲突，特别是对于矩阵加载指令，并且全局归约以最小的内存开销发生。
• 我们设置并展开循环，使得大部分内存偏移量是静态的，从而最小化运行时索引计算。
• 我们实现了一种“条纹”分区方案，其中每个 SM（流式多处理器）处理的瓦片段可能（部分）跨越多个列“切片”。这导致在大多数矩阵形状上具有良好的 SM 利用率，同时最小化所需的全局归约步骤。
• 全局归约直接在输出缓冲区中发生（暂时将 FP32 累加器降维到 FP16），该缓冲区保留在 L2 缓存中；归约操作通常也经过优化以避免任何不必要的读写。
• 总体而言，内核的 PTX（并行线程执行）汇编代码在 NSight-Compute 中进行了广泛分析，CUDA 代码包含几个更多冗余或稍次优的结构，但它们编译后能生成更快的 PTX。

基准测试：

我们首先在 NVIDIA A10 GPU 上可以理想分区的大矩阵上，比较 Marlin 与其他流行的 4 位推理内核的性能。 这使得所有内核都能达到几乎最佳的潜在性能。 所有内核均在组大小（groupsize）128 下执行（不过，我们注意到缩放格式并不完全相同）。

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

由于其条纹分区方案，Marlin 在实际（较小）矩阵和各种 GPU 上也带来了强劲的性能。 以下结果证明了这一点，我们在批量大小为 16 的情况下，对流行开源模型的 Transformer 块中所有线性层的整体运行时间进行了基准测试。

最后，我们还研究了在锁定的基础 GPU 时钟下，长时间运行能维持何种性能。 有趣的是，我们发现降低时钟速度会显著损害前期内核的相对加速比，但对 Marlin 的近乎最优性能（相对于较低的时钟设置）没有影响。

要求：

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

用法：

如果满足所有要求，可以通过调用以下命令在该仓库的根文件夹中安装 Marlin：

pip install .

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

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

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

此外，如果启用了 ECC (错误检查和纠正)（例如在 A10 上），则由于每个内存请求都将包含校验和开销（checksum overheads），实际可达到的最大内存带宽（memory bandwidth）将比官方规格表低 10-15%。这可以通过以下方式禁用：

sudo nvidia-smi -e 0

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

GPTQ 示例：

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

% Compress Llama2 model and export model in Marlin format.
python llama2.py LLAMA2_CHECKPOINT --wbits 4 --save checkpoint.pt
% Perform perplexity evaluation of uncompressed model.
python llama2.py LLAMA2_CHECKPOINT
% Evaluate compressed model (with Marlin kernels) in the eval harness.
python eval.py --model hf --model_args pretrained=LLAMA2_CHECKPOINT --tasks mmlu \
  --marlin_checkpoint checkpoint.marlin.g128
% Evaluate full precision baseline.
python eval.py --model hf --model_args pretrained=LLAMA2_CHECKPOINT --tasks mmlu 

我们测量了 4 位（group=128）Marlin 模型的以下 WikiText 和 Red-Pajama 困惑度（perplexities），以及 MMLU 零样本准确率（zero-shot accuracy）：

Llama2	Wiki2 (FP16)	Wiki2 (INT4)	RedPaj (FP16)	RedPaj (INT4)	MMLU (FP16)	MMLU (INT4)
7B	5.12	5.27	6.14	6.30	41.80	40.07
13B	4.57	4.67	5.67	5.79	52.10	51.13
70B	3.12	3.21	4.74	4.81	65.43	64.81

我们注意到，此 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}
}

Marlin 快速上手指南

Marlin 是一款专为大语言模型（LLM）推理设计的混合精度自回归线性内核（Mixed Auto-Regressive Linear kernel）。它通过高度优化的 FP16xINT4 矩阵乘法，在 NVIDIA GPU 上提供接近理论极限的 4 倍加速比，特别适用于批量推理和推测性解码场景。

环境准备

在使用 Marlin 之前，请确保您的开发环境满足以下硬件和软件要求：

• GPU 要求: NVIDIA 显卡，计算能力（Compute Capability）>= 8.0（如 Ampere 或 Ada 架构，例如 A10、RTX 30/40 系列）。
  • 注意：目前尚未针对 Hopper 架构优化。
• CUDA 版本: >= 11.8（需与 PyTorch 安装的 CUDA 版本一致）。
• Python 依赖:
  • torch>=2.0.0
  • numpy
• 量化脚本依赖 (如需运行 GPTQ 压缩示例):
  • transformers
  • datasets
  • sentencepiece

安装步骤

1. 克隆仓库

   git clone https://github.com/IST-DASLab/Marlin.git
   cd Marlin
   

2. 安装基础依赖 推荐使用国内镜像源以加快下载速度：

   pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch>=2.0.0 numpy
   

3. 安装 Marlin 核心包 在项目根目录下执行：

   pip install .
   

4. (可选) 安装量化工具依赖 如果您需要运行 gptq 子目录中的模型压缩脚本：

   pip install -i https://pypi.tuna.tsinghua.edu.cn/simple transformers datasets sentencepiece
   

基本使用

Marlin 提供了两种主要的使用方式：通过 marlin.Layer 模块封装层，或直接调用底层算子。

1. 使用 marlin.Layer 封装量化层

这是最便捷的方式，可将标准的 torch.Linear 层转换为压缩后的 Marlin 格式。

import marlin
import torch.nn as nn

# 假设 linear 是已加载的模型层，scales 是对应的缩放参数
# 将“伪量化”层打包为 Marlin 格式
linear_marlin = marlin.Layer()
linear_marlin.pack(linear, scales)

2. 直接调用算子

如果权重和缩放参数已经过预处理，可直接调用 marlin.mul 进行计算。

# 前提：weights 和 scales 已完成必要的预处理
output = marlin.mul(weights, scales, input_activations)

3. 测试与基准验证

安装完成后，您可以运行内置脚本来验证正确性或测试性能：

• 运行正确性测试:

  python test.py
  

• 运行性能基准测试:

  python bench.py
  

注意: 若要复现官方文档中的“可持续性能”基准测试结果，需要将 GPU 频率锁定至基础值。请在终端执行以下命令（需 sudo 权限）：

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

此外，若使用带有 ECC 功能的显卡（如 A10），建议关闭 ECC 以获得最大内存带宽：

sudo nvidia-smi -e 0