使用PyTorch在AMD GPU上进行INT8量化实现精简化的LLM推理

Leaner LLM Inference with INT8 Quantization on AMD GPUs using PyTorch — ROCm Blogs

随着大型语言模型（LLMs）规模达到数千亿参数，我们在这些庞大模型中表示数据的方式极大地影响了训练所需的资源（例如，用于推理的GPU数量）。在我们之前的博客（JAX混合精度训练；PyTorch AMP）中，我们已经展示了混合精度训练如何加速LLMs的训练过程。在这篇博客文章中，我们将进一步推进，向您展示如何通过量化到更低精度的数据格式来加速推理，节省时间和内存，而不会牺牲模型的整体性能。量化是一种技术，通过将模型参数的精度从32位浮点（FP32）或16位浮点（FP16）减少到8位整数（INT8）。标准模型通常使用32位浮点（FP32）精度。然而，对于推理任务来说，这种高精度并不是总是必要的。通过将模型权重和激活转换为较低精度格式如INT8（8位整数），我们可以实现更快的计算和更低的内存使用，有效地将模型大小减少到原来的四分之三（从32位）或一半（从16位），仅有轻微的准确性下降，这常常被速度提升所弥补。

在这篇博客文章中，我们将逐步向您展示如何使用ROCm、PyTorch和gpt-fast repository在AMD GPU上实现INT8量化，以及如何评估由此产生的推理性能。具体来说，我们将演示INT8量化如何显著提升Llama系列和Mistral LLM模型的推理速度。

如何进行量化

大多数模型量化技术可以分为以下两类：

• 训练后量化（PTQ）：在模型完全训练后应用。这种方法更简单，但可能会导致一定的性能损失。

• 量化感知训练（QAT）：在训练过程中引入量化，使量化后的权重能够更好地捕捉数据的信息。它通常能带来更好的结果，但需要更多的计算资源。

关于这两种策略的更多信息，请参见LLM 系列 - 量化概述。

对于我们下面的代码示例，我们将使用 PTQ。以下是 PTQ 的一般步骤，尽管实际步骤在不同应用中可能会有很大差异。

1. 模型准备

• *加载预训练模型*：从一个预训练模型开始，通常是 FP32、FP16 或 BF16 格式。

• *定义量化配置*：指定量化方案和配置，例如对称量化或非对称量化，以及按通道或按张量量化。

2. 校准

• *收集校准数据*：收集一个代表性的数据集，以捕捉模型在推理期间遇到的输入分布。

• *运行校准*：使用校准数据运行模型并收集统计数据，例如每层激活的最小值和最大值。此步骤确定权重和激活的量化参数（比例和零点）的值。这些比例和零点参数类似于标准化中的标准差和均值。

3. 量化和模型转换

• *量化权重和激活*：使用校准步骤中确定的量化参数值，将较高精度的权重和激活量化为 INT8。

• *转换模型格式*：使用像 PyTorch 这样的框架将模型转换为量化格式。

在我们的演示中，我们将只对权重进行量化，并跳过校准过程。权重的分布是已知且固定的，可以直接从权重值本身计算出量化比例和零点参数。

实现

在官方的 gpt-fast repository 中，作者测量了 meta-llama/Llama-2-7b-chat-hf 模型在 MI-250x GPU 上的推理速度，重点关注模型处理数据的速度。然而，在实际环境中评估推理效率时，还需要考虑吞吐量，它是衡量在给定时间内能处理多少数据（在这种情况下是 tokens）。吞吐量通常以 Tokens/Second 表示，提供了模型在实际应用中性能的更全面理解。

在我们的实现中，我们将使用与 MI-250x GPU 的一个 GCD（图形计算芯片）等效的 MI210 GPU（一个 MI-250x GPU 可以看作两个 MI210 GPU）。我们将测量 Llama-2-7B 推理吞吐量作为基准，然后将测试扩展到另外三个流行的模型：`meta-llama/Meta-Llama-3-8B`（Llama 家族模型的较新版本）、`mistralai/Mistral-7B-v0.1` 和 meta-llama/Llama-2-13b-chat-hf。对于每个模型，我们将测试三种不同优化水平的模式以确定其性能：

1. Eager 模式（无优化）

2. Torch.compile

3. Torch.compile + INT8 量化

torch.compile 是 PyTorch 的一个功能，通过将模型执行转换为更高效的编译形式来优化模型执行，从而提高运行时性能。

我们将使用带有每晚构建 PyTorch 的 ROCm Docker 容器进行演示。PyTorch 持续改进，每晚版本通常包含最新的优化。Docker 容器将在运行 Ubuntu 的 AMD GPUs 服务器上运行。

请参阅 System requirements (Linux) 了解 AMD 支持的硬件和操作系统的完整列表。

• 使用以下命令在 Linux 终端中拉取和运行 Docker 容器:

  docker run -it --ipc=host --network=host --device=/dev/kfd --device=/dev/dri \
            --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
            --name=pt_nightly rocm/pytorch-nightly:latest /bin/bash
  

  您可以通过在 Python 控制台中运行以下代码来查看 PyTorch 检测到的 GPUs 数量。PyTorch 至少需要检测到一块 GPU。

  import torch
  torch.cuda.device_count()
  

• 安装所需的 Python 包:

  python3 -m pip install --upgrade pip
  pip install sentencepiece huggingface_hub tiktoken blobfile
  

• 使用以下命令下载 gpt-fast 仓库:

  git clone https://github.com/pytorch-labs/gpt-fast.git
  

• 下载 generate.py 和 run_commands.sh 文件，获取地址为 此博客srcGitHub 文件夹。将 generate.py 和 run_commands.sh 放在下载的 gpt-fast 文件夹中，替换原始的 generate.py 文件。`gpt-fast` 仓库中的原始 generate.py 文件仅仅在基准模型上计算每秒处理的平均 tokens 数，而不排除前几轮的热身，这可能导致结果偏差。我们修改并增强了 generate.py 文件：

  • 运行总共 30 次迭代，前 20 次迭代作为热身，并计算最后 10 轮的平均值和标准差。

  • 计算内存带宽的平均值和标准差。

运行基准模型并收集指标

下面是一个名为`run_commands.sh`的文件，它包含了下载、量化和运行基准模型以收集推理指标 Tokens/Second 和 Memory Bandwidth (GB/s) 的命令。

您需要提供您的 Hugging Face 凭据才能运行 run_commands.sh。通过运行 huggingface-cli login 并按照提示操作来提供您的凭据。有关如何获得 Hugging Face 访问令牌的信息，请参阅 Hugging Face 用户文档中的 用户访问令牌 。

#!/bin/bash

# 日志文件
LOGFILE="output_4_models.log"

# 如果日志文件存在则清除
> $LOGFILE

# 模型仓库数组
MODEL_REPOS=("meta-llama/Llama-2-7b-chat-hf" "meta-llama/Meta-Llama-3-8B" "mistralai/Mistral-7B-v0.1" "meta-llama/Llama-2-13b-chat-hf")

# 运行命令并记录输出
{
    echo "Processing models"
    
    # 循环遍历模型仓库
    for MODEL_REPO in "${MODEL_REPOS[@]}"; do
        # 准备/下载模型
        ./scripts/prepare.sh $MODEL_REPO

        echo -e "\n**************正在运行基线模型 $MODEL_REPO..."
        python generate.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --prompt "Hello, my name is"
        echo -e "\n**************正在运行 torch.compile 模型 $MODEL_REPO..."
        python generate.py --compile --checkpoint_path checkpoints/$MODEL_REPO/model.pth --prompt "Hello, my name is"

        echo "Setting DEVICE to cuda..."
        export DEVICE=cuda

        echo -e "\n**************正在量化并运行命令  $MODEL_REPO..."
        python quantize.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --mode int8
        echo -e "\n**************运行 int8 模型 $MODEL_REPO..."
        python generate.py --compile --checkpoint_path checkpoints/$MODEL_REPO/model_int8.pth --device $DEVICE
    done
} &> $LOGFILE

在 gpt-fast 文件夹中运行 ./run_commands.sh 来收集指标。

基准测试结果

在下表中，我们展示了每种推理模式下各个模型的指标。在最后一列中可以看到，与torch compile模式相比，INT8量化通过量提高了大约25-45％，与eager模式相比甚至更多。这证实了模型量化可以用来提升大规模语言模型（LLMs）的推理性能。

总结

在这篇博客文章中，我们逐步向您展示了如何使用AMD GPU实现INT8量化，并如何基准测试结果推理。我们展示了INT8量化对Llama系列和Mistral大型语言模型训练的加速效果。

致谢

我们要感谢PyTorch Labs开发的 gpt-fast 仓库，它为我们的工作提供了指南。