{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LLM Compressor Workbench -- Getting Started\n",
    "\n",
    "This notebook will demonstrate how common [LLM Compressor](https://github.com/vllm-project/llm-compressor) flows can be run on the Alauda AI.\n",
    "\n",
    "We will show how a user can compress a Large Language Model, without data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data-Free Model Compression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llmcompressor.modifiers.quantization import QuantizationModifier\n",
    "\n",
    "# model to compress\n",
    "model_id = \"./TinyLlama-1.1B-Chat-v1.0\"\n",
    "\n",
    "# This recipe will quantize all Linear layers except those in the `lm_head`,\n",
    "#  which is often sensitive to quantization. The W4A16 scheme compresses\n",
    "#  weights to 4-bit integers while retaining 16-bit activations.\n",
    "recipe = QuantizationModifier(targets=\"Linear\", scheme=\"W4A16\", ignore=[\"lm_head\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load up model using huggingface API\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id, device_map=\"auto\", torch_dtype=\"auto\"\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run compression using `oneshot`\n",
    "from llmcompressor import oneshot\n",
    "\n",
    "model = oneshot(model=model, recipe=recipe, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save model and tokenizer\n",
    "model_dir = \"./\" + model_id.split(\"/\")[-1] + \"-W4A16\"\n",
    "model.save_pretrained(model_dir)\n",
    "tokenizer.save_pretrained(model_dir);"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}