适用于 Alauda AI 的 LLM Compressor

本文档介绍如何使用 LLM Compressor 与 Alauda AI 平台集成,以执行模型压缩工作流。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。本指南中的示例 notebooks 使用 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. 配置分片上传,每个分片 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

创建并运行压缩 notebooks

根据你的使用场景下载相应的示例 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 实现快速推理。 请按照 create inference service 中的说明完成此步骤。