Alauda AI 上的 LLM Compressor

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

notebook

支持的模型压缩工作流

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

创建 Workbench

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

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

创建模型仓库并上传模型

有关创建模型仓库和上传模型文件的详细步骤,请参考 使用 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,你必须准备并上传一个校准数据集。请按照 使用 Notebook 上传模型 中描述的相同步骤准备数据集。示例 calibration 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. 使用 Hugging Face 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. 配置分块上传,块大小为 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"] = "<access-key>"
os.environ["AWS_SECRET_ACCESS_KEY"] = "<secret-key>"

storage_options = {
  "key": "<access-key>",
  "secret": "<secret-key>",
  "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. 设置环境变量(作为备份,某些底层组件会使用它们)。将 <access-key><secret-key> 替换为对象存储账户的凭据。
  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

对于 calibration 数据集,也可以用类似方式下载:

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

创建并运行压缩 notebooks

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

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

压缩完成后,将压缩后的模型重新上传到模型仓库。有关将模型文件上传到模型仓库的详细步骤,请参见 使用 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 使用经过优化的部署内核实现快速推理。 请按照 创建推理服务 中的说明完成此步骤。