在 Alauda AI 中使用 LLM Compressor

本文档介绍如何在 Alauda AI 平台中使用 LLM Compressor 集成来执行模型压缩工作流。LLM Compressor 的 Alauda AI 集成提供了两个示例工作流:

notebook

支持的模型压缩工作流

在 Alauda AI 平台上,你可以使用 Workbench 功能对存储在模型仓库中的模型运行 LLM Compressor。下面的工作流概述了压缩模型的典型步骤。

创建 Workbench

按照 Create Workbench 中的说明创建一个新的 Workbench 实例。创建 Workbench 时,请确保选择镜像 odh-workbench-jupyter-pytorch-llmcompressor-cuda-py312-ubi9,该镜像包含所需的 LLM Compressor 环境。目前模型压缩仅支持在 JupyterLab 中使用。

如果你的环境中没有该镜像,请参考链接文章中的说明,同步该镜像并修补 WorkspaceKind 资源,使该镜像可用于 Workbench 创建选项。

创建模型仓库并上传模型

有关创建模型仓库和上传模型文件的详细步骤,请参阅 Upload Models Using Notebook。本指南中的示例 notebook 使用 TinyLlama-1.1B-Chat-v1.0 模型。

data-free-compressor.ipynb
from llmcompressor.modifiers.quantization import QuantizationModifier

model_id = "./TinyLlama-1.1B-Chat-v1.0"
recipe = QuantizationModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"])
  1. 要压缩的模型。如果你想使用自己的模型,可以修改这一行
  2. 该 recipe 将量化除 lm_head 之外的所有 Linear 层,因为 lm_head 通常对量化比较敏感。W4A16 方案会将权重量化为 4-bit 整数,同时保留 16-bit 激活值。

(可选)准备并上传数据集

NOTE

如果你计划使用 data-free compressor notebook,可以跳过此步骤。

要使用 calibration compressor notebook,你必须准备并上传一个校准数据集。使用与 Upload Models Using Notebook 中描述的相同流程准备数据集。本示例校准 notebook 使用 ultrachat_200k 数据集。

calibration-compressor.ipynb
from datasets import load_dataset

dataset_id = "./ultrachat_200k"

num_calibration_samples = 512 if use_gpu else 4
max_sequence_length = 2048 if use_gpu else 16

ds = load_dataset(dataset_id, split="train_sft")
ds = ds.shuffle(seed=42).select(range(num_calibration_samples))

def preprocess(example): 
    text = tokenizer.apply_chat_template(
        example["messages"],
        tokenize=False,
    )
    return tokenizer(
        text,
        padding=False,
        max_length=max_sequence_length,
        truncation=True,
        add_special_tokens=False,
    )

ds = ds.map(preprocess, remove_columns=ds.column_names)
  1. 使用 Huggingface datasets API 创建校准数据集。如果你想使用自己的数据集,可以修改这一行
  2. 选择样本数量。512 个样本是一个不错的起点。增加样本数量可以提升准确性。
  3. 加载数据集。
  4. 打乱顺序并仅获取所需数量的样本。
  5. 预处理并分词为模型使用的格式。

(可选)将数据集上传到 S3 兼容对象存储

如果你想将数据集上传到 S3 兼容对象存储,可以在 JupyterLab 中运行以下代码。Alauda AI 支持 S3 兼容存储访问,在典型的产品部署中,存储实现通常是 Ceph 对象存储,因此你可以使用标准的 boto3 客户端。

import os
from boto3.s3.transfer import TransferConfig
import boto3

local_folder = "./ultrachat_200k"
bucket_name = "datasets"

config = TransferConfig(
    multipart_threshold=100*1024*1024,
    max_concurrency=10,
    multipart_chunksize=100*1024*1024,
    use_threads=True
)

for root, dirs, files in os.walk(local_folder):
    for filename in files:
        local_path = os.path.join(root, filename)
        relative_path = os.path.relpath(local_path, local_folder)
        s3_key = f"ultrachat_200k/{relative_path.replace(os.sep, '/')}"
        s3.upload_file(local_path, bucket_name, s3_key, Config=config)
        print(f"Uploaded {local_path} -> {s3_key}")
  1. 如果你想使用自己的数据集,可以修改这一行
  2. 配置 multipart 上传,使用 100 MB 分块,并最多支持 10 个并发线程。

(可选)从 S3 兼容对象存储中使用数据集

如果你想使用存储在 S3 兼容对象存储中的数据集,首先安装 s3fs 工具,然后按如下所示修改示例中的数据集加载部分。在 Alauda AI 环境中,这种 S3 兼容存储通常由 Ceph 对象存储提供支持。

pip install s3fs -i https://pypi.tuna.tsinghua.edu.cn/simple
calibration-compressor.ipynb
import os
from datasets import load_dataset

os.environ["AWS_ACCESS_KEY_ID"] = "@7Apples@"
os.environ["AWS_SECRET_ACCESS_KEY"] = "07Apples@"

storage_options = {
  "key": "07Apples@",
  "secret": "O7Apples@",
  "client_kwargs": {
    "endpoint_url": "https://ceph-obj.example.com"
  }
}

ds = load_dataset(
      'parquet',
      data_files='s3://datasets/ultrachat_200k/data/train_sft-*.parquet', 
      storage_options=storage_options, 
      split="train"
)
  1. 设置环境变量(作为备用,一些底层组件会使用它们)。
  2. 定义存储配置;你必须为 S3 兼容对象存储服务显式指定 endpoint_url,例如 Ceph 对象存储端点。
  3. 如果数据集已拆分,则这相当于示例中的 split="train_sft"

在 JupyterLab 中克隆模型和数据集

在 JupyterLab 终端中,使用 git clone 将模型仓库(如果适用,也包括数据集)下载到工作区。data-free compressor notebook 不需要数据集。

INFO

如果网络访问允许,我们建议使用 Hugging Face(hf)命令行工具直接下载模型和数据集。这种方式通常比手动克隆仓库更快,也更简单。

如果你处于受限网络环境中,可以使用 Hugging Face 镜像来加速访问:

export HF_ENDPOINT=https://hf-mirror.com

你可以使用 hf 命令行工具直接下载模型和数据集。例如,要下载 TinyLlama 模型:

hf download TinyLlama/TinyLlama-1.1B-Chat-v1.0 --local-dir TinyLlama-1.1B-Chat-v1.0

对于校准数据集,也可以类似下载:

hf download --repo-type dataset HuggingFaceH4/ultrachat_200k --local-dir ultrachat_200k

创建并运行压缩 notebook

根据你的使用场景下载相应的示例 notebook:如果你使用了数据集,则下载 calibration compressor notebook;否则下载 data-free compressor notebook。然后点击 JupyterLab 页面上的向上箭头按钮,上传已下载的 notebook 文件。

将压缩后的模型上传回仓库

压缩完成后,将压缩后的模型重新上传到模型仓库。有关将模型文件上传到模型仓库的详细步骤,请参阅 Upload Models Using Notebook

model_dir = "./" + model_id.split("/")[-1] + "-W4A16"
model.save_pretrained(model_dir)
tokenizer.save_pretrained(model_dir);
  1. 保存模型和 tokenizer。如果你想更改输出名称,可以修改这一行

部署并使用压缩后的模型进行推理

你使用 LLM Compressor 创建的量化和稀疏模型会使用 compressed-tensors 库保存(该库是 Safetensors 的扩展)。 压缩格式与模型的量化或稀疏类型相匹配。这些格式在 vLLM 中原生支持,从而可以通过 Alauda AI Inference Server 使用优化后的部署 kernel 实现快速推理。 按照 创建推理服务 中的说明完成此步骤。