#LLM Compressor with Alauda AI

This document describes how to use the LLM Compressor integration with the Alauda AI platform to perform model compression workflows. The Alauda AI integration of LLM Compressor provides two example workflows:

• A workbench image and the data-free-compressor.ipynb that demonstrate how to compress a model.
• A workbench image and the calibration-compressor.ipynb that demonstrate how to compress a model using a calibration dataset.

#Supported Model Compression Workflows

On the Alauda AI platform, you can use the Workbench feature to run LLM Compressor on models stored in S3, a PVC, or an OCI registry. The following workflow outlines the typical steps for compressing a model.

#Create a Workbench

Follow the instructions in Create Workbench to create a new Workbench instance. When creating the Workbench, make sure to select the image odh-workbench-jupyter-pytorch-llmcompressor-cuda-py312-ubi9, which includes the required LLM Compressor environment. Model compression is currently supported only within JupyterLab.

If this image is not available in your environment, refer to the linked article for instructions on how to synchronize the image and patch the WorkspaceKind resource so that the image becomes available in the Workbench creation options.

#Prepare the Model Artifact

Download the model into the Workbench or stage it on a PVC/S3 location accessible from the Workbench. After compression, upload the output to a supported storage backend. See Upload Models Using a Workbench. Register the resulting Model URI as a new Model Registry version using Model Registry. The example notebooks in this guide use the TinyLlama-1.1B-Chat-v1.0 model.

from llmcompressor.modifiers.quantization import QuantizationModifier

model_id = "./TinyLlama-1.1B-Chat-v1.0"
recipe = QuantizationModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"])

#(Optional) Prepare and Upload a Dataset

To use the calibration compressor notebook, you must prepare and upload a calibration dataset. Prepare your dataset using the same process described in Upload Models Using Notebook. The example calibration notebook uses the ultrachat_200k dataset.

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)

#(Optional) Upload Dataset into S3-Compatible Object Storage

If you want to upload datasets into S3-compatible object storage, you can run the following code in JupyterLab. Alauda AI supports S3-compatible storage access, and in typical product deployments the storage implementation is Ceph object storage, so you can use the standard boto3 client.

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}")

#(Optional) Use Dataset from S3-Compatible Object Storage

If you want to use datasets stored in S3-compatible object storage, first install the s3fs tool and then modify the dataset loading section in the example as shown below. In Alauda AI environments, this S3-compatible storage is typically backed by Ceph object storage.

pip install s3fs -i https://pypi.tuna.tsinghua.edu.cn/simple

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"
)

#Download Models and Datasets in JupyterLab

In the JupyterLab terminal, use the hf command-line tool or your approved object-storage client to download the model and dataset into the workspace. The data-free compressor notebook does not require a dataset.

If you are in a restricted network environment, you can use the Hugging Face mirror for accelerated access:

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

You can use the hf command-line tool to download models and datasets directly. For example, to download the TinyLlama model:

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

For calibration datasets, download them similarly:

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

#Create and Run Compression Notebooks

Download the appropriate example notebook for your use case: the calibration compressor notebook if you are using a dataset, or the data-free compressor notebook otherwise. Click the upward arrow button on the JupyterLab page to upload the downloaded notebook file.

#Register the Compressed Model

Once compression is complete, upload the compressed model to S3, a PVC, or an OCI registry using Upload Models Using a Workbench. Then register its Model URI as a new Model Registry version using Model Registry.

model_dir = "./" + model_id.split("/")[-1] + "-W4A16"
model.save_pretrained(model_dir)
tokenizer.save_pretrained(model_dir);

#Deploy and Use the Compressed Model for Inference

Quantized and sparse models that you create with LLM Compressor are saved using the compressed-tensors library (an extension of Safetensors). The compression format matches the model's quantization or sparsity type. These formats are natively supported in vLLM, enabling fast inference through optimized deployment kernels by using Alauda AI Inference Server. Follow the instructions in create inference service to complete this step.