runing_quantized_models_with_mlc_llm.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EWOtpjJMDmYE"
      },
      "source": [
        "# Runing quantized models with MLC-LLM"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Acknowledgement: This file is modified from the official MLC-LLM instruction [https://colab.research.google.com/github/mlc-ai/notebooks/blob/main/mlc-llm/tutorial_chat_module_getting_started.ipynb](https://colab.research.google.com/github/mlc-ai/notebooks/blob/main/mlc-llm/tutorial_chat_module_getting_started.ipynb)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This file including three section:\n",
        "- [Environment setup and file download](#environment-setup-and-file-download)\n",
        "- [Let's Chat!](#lets-chat)\n",
        "- [Compile your own quantized models with MLC-LLM](#compile-your-own-quantized-models-with-mlc-llm)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Environment setup and file download"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Firstly, let's download the MLC-AI and MLC-Chat nightly build packages. Go to https://mlc.ai/package/ and replace the command below with the one that is appropriate for your hardware and OS."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PgW-5OAADmYE"
      },
      "outputs": [],
      "source": [
        "!pip install --pre --force-reinstall mlc-ai-nightly-cu118 mlc-chat-nightly-cu118 -f https://mlc.ai/wheels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FwsWd1WbDmYE"
      },
      "source": [
        "Next, let's download the prebuilt model libraries from HuggingFace. In order to download the large weights, we'll have to use `git lfs`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V0GjINnMDmYF"
      },
      "outputs": [],
      "source": [
        "!conda install git git-lfs\n",
        "!git lfs install"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yYwjsCOK7Jij"
      },
      "source": [
        "Download the prebuilt quantized models weights:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "FSAe7Ew_DmYF"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "mkdir: cannot create directory ‘dist/’: File exists\n",
            "fatal: destination path './dist' already exists and is not an empty directory.\n"
          ]
        }
      ],
      "source": [
        "!mkdir -p dist/\n",
        "# download llama-2-7b-chat with w3a16g128 quantization\n",
        "!git clone https://huggingface.co/ChenMnZ/Llama-2-7b-chat-omniquant-w3a16g128asym ./dist/Llama-2-7b-chat-omniquant-w3a16g128asym\n",
        "# optional, you can also download llama-2-13b-chat with w3a16g128 quantization\n",
        "#!git clone https://huggingface.co/ChenMnZ/Llama-2-13b-chat-omniquant-w3a16g128asym ./dist/Llama-2-13b-chat-omniquant-w3a16g128asym"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "76Ru5__tDmYF"
      },
      "source": [
        "## Let's Chat!\n",
        "\n",
        "Before we can chat with the model, we must first import a library and instantiate a `ChatModule` instance. The `ChatModule` must be initialized with the appropriate model name."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "AJAt6oW7DmYF"
      },
      "outputs": [],
      "source": [
        "from mlc_chat import ChatModule\n",
        "from mlc_chat.callback import StreamToStdout\n",
        "\n",
        "cm = ChatModule(model=\"dist/Llama-2-7b-chat-omniquant-w3a16g128asym/params\", model_lib_path=\"dist/Llama-2-7b-chat-omniquant-w3a16g128asym/Llama-2-7b-chat-omniquant-w3a16g128asym-cuda.so\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zEaVXnnJDmYF"
      },
      "source": [
        "That is all what needed to set up the `ChatModule`. You can now chat with the model by entering any prompt you'd like. Try it out below!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "TNmg9N_NDmYF"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Prompt:What is the meaning of life?\n",
            "\n",
            "Ah, a question that has puzzled philosophers and theologians for centuries! The meaning of life is a complex and multi-faceted topic, and there is no one definitive answer. However, here are some possible perspectives on the meaning of life:\n",
            "\n",
            "1. Religious or Spiritual Perspective: Many people believe that the meaning of life is to fulfill a divine or spiritual purpose. According to this view, the meaning of life is to follow the will of a higher power or to achieve spiritual enlightenment.\n",
            "2. Personal Fulfillment Perspective: Others believe that the meaning of life is to achieve personal fulfillment and happiness. According to this view, the meaning of life is to pursue one's passions and interests, and to live a life that is fulfilling and satisfying.\n",
            "3. Social Perspective: Some people believe that the meaning of life is to contribute to the greater good of society. According to this view, the meaning of life is to make a positive impact on the world and to leave a lasting legacy.\n",
            "4. Existentialist Perspective: Existentialists believe that the meaning of life is not predetermined or inherent, but rather something that is created by each individual. According to this view, the meaning of life is something that each person must create for themselves through their experiences and choices.\n",
            "5. Happiness Perspective: Some people believe that the meaning of life is to achieve happiness and well-being. According to this view, the meaning of life is to pursue happiness and well-being, either through personal fulfillment or through the love and support of others.\n",
            "6. Purpose-Driven Perspective: Others believe that the meaning of life is to fulfill a purpose or mission. According to this view, the meaning of life is to identify one's purpose in life and to work towards fulfilling it.\n",
            "7. Cultural Perspective: The meaning of life can also be influenced by cultural beliefs and values. According to this view, the meaning of life is shaped by the cultural and social norms of one's community and upbringing.\n",
            "8. Personal Identity Perspective: Some people believe that the meaning of life is closely tied to their personal identity and sense of self. According to this view, the meaning of life is to discover and cultivate one's personal identity and sense of purpose\n"
          ]
        }
      ],
      "source": [
        "prompt = \"What is the meaning of life?\"\n",
        "print(f\"Prompt:{prompt}\\n\")\n",
        "output = cm.generate(\n",
        "    prompt=prompt,\n",
        "    progress_callback=StreamToStdout(callback_interval=3),\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Prompt:How many points did you list out?\n",
            "\n",
            "I listed out 8 possible perspectives on the meaning of life in my previous response:\n",
            "1. Religious or Spiritual Perspective\n",
            "2. Personal Fulfillment Perspective\n",
            "3. Social Perspective\n",
            "4. Existentialist Perspective\n",
            "5. Happiness Perspective\n",
            "6. Purpose-Driven Perspective\n",
            "7. Cultural Perspective\n",
            "8. Personal Identity Perspective\n",
            "\n",
            "I hope this helps! Let me know if you have any other questions.\n"
          ]
        }
      ],
      "source": [
        "prompt = \"How many points did you list out?\"\n",
        "print(f\"Prompt:{prompt}\\n\")\n",
        "output = cm.generate(\n",
        "    prompt=prompt,\n",
        "    progress_callback=StreamToStdout(callback_interval=3),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yso2UiUTyZny"
      },
      "source": [
        "You can also repeat running the code block below for multiple rounds to interact with the model in a chat style."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "qJyz3a7vyZny"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Of course, here are more details on the third perspective on the meaning of life:\n",
            "The third perspective on the meaning of life is that it is to contribute to the greater good of society. According to this view, the meaning of life is to make a positive impact on the world and to leave a lasting legacy. This perspective is often associated with altruism and a sense of social responsibility.\n",
            "Some of the key aspects of this perspective include:\n",
            "1. Social Justice: The belief that everyone should have equal access to opportunities and resources, and that it is important to work towards creating a more just and equitable society.\n",
            "2. Community Involvement: The belief that individuals have a responsibility to contribute to the well-being of their community, whether through volunteering, activism, or other forms of service.\n",
            "3. Personal Responsibility: The belief that individuals have a responsibility to take care of themselves and their loved ones, and to work towards creating a better world for themselves and future generations.\n",
            "4. Legacy: The belief that individuals should strive to leave a lasting legacy that will continue to have a positive impact on the world after they are gone.\n",
            "5. Personal Fulfillment: The belief that contributing to the greater good can also bring personal fulfillment and satisfaction, as individuals find meaning and purpose in their work.\n",
            "6. Ethical Considerations: The belief that it is important to act ethically and with integrity, and to consider the potential consequences of one's actions on society and the environment.\n",
            "I hope this helps! Let me know if you have any other questions.\n"
          ]
        }
      ],
      "source": [
        "prompt = input(\"Prompt: \")\n",
        "output = cm.generate(prompt=prompt, progress_callback=StreamToStdout(callback_interval=3))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I4bOyUk7DmYF"
      },
      "source": [
        "To check the generation speed of the chat bot, you can print the statistics."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "PPbPj6vpDmYF"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Statistics: prefill: 240.0 tok/s, decode: 83.6 tok/s\n",
            "\n"
          ]
        }
      ],
      "source": [
        "print(f\"\\nStatistics: {cm.stats()}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XAb-XZPnDmYF"
      },
      "source": [
        "By default, the `ChatModule` will keep a history of your chat. You can reset the chat history by running the following."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iKpKgVxNDmYF"
      },
      "outputs": [],
      "source": [
        "cm.reset_chat()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FifWDyOeyZny"
      },
      "source": [
        "### Benchmark Performance\n",
        "\n",
        "To benchmark the performance, we can use the `benchmark_generate` method of ChatModule. It takes an input prompt and the number of tokens to generate, ignores the system prompt and model stop criterion, generates tokens in a language model way and stops until finishing generating the desired number of tokens. After calling `benchmark_generate`, we can use `stats` to check the performance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UHgsAAZByZnz"
      },
      "outputs": [],
      "source": [
        "print(cm.benchmark_generate(prompt=\"What is benchmark?\", generate_length=512))\n",
        "cm.stats()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compile your own quantized models with MLC-LLM"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Firstly, saveing the fakely quantized models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!CUDA_VISIBLE_DEVICES=0 python main.py \\\n",
        "--epochs 0 --output_dir ./log/temp \\\n",
        "--wbits 3 --abits 16 --group_size 128 --lwc \\\n",
        "--model /PATH/TO/LLaMA/Llama-2-7b-chat \\\n",
        "--save_dir /PATH/TO/SAVE/Llama-2-7b-chat-omniquant-w3a16g128 \\\n",
        "--resume /PATH/TO/CHECKPOINTS/Llama-2-7b-chat-w3a16g128.pth"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "To compile the pseudo-quantized models, refer to [https://mlc.ai/mlc-llm/docs/compilation/compile_models.html](https://mlc.ai/mlc-llm/docs/compilation/compile_models.html). Be aware that the MLC-LLM project has its customized quantization schemes in [mlc_llm/quantization/__init__.py](https://github.com/mlc-ai/mlc-llm/blob/main/mlc_llm/quantization/__init__.py). Please incorporate the W3A16g128 quantization scheme into the aforementioned [mlc_llm/quantization/__init__.py](https://github.com/mlc-ai/mlc-llm/blob/main/mlc_llm/quantization/__init__.py)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "    \"w3a16g128asym\": QuantizationScheme(\n",
        "        name=\"w3a16g128asym\",\n",
        "        linear_weight=GroupQuantizationSpec(\n",
        "            dtype=\"float16\",\n",
        "            mode=\"int3\",\n",
        "            sym=False,\n",
        "            storage_nbit=16,\n",
        "            group_size=128,\n",
        "            transpose=False,\n",
        "        ),\n",
        "        embedding_table=None,\n",
        "        final_fc_weight=None,\n",
        "    ),"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "You can probably figure out the general pattern for any other quantization scheme you want to do. Once you have this, then you can compile the MLC-LLM model, which should be quick:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!python build.py --target cuda --quantization w3a16g128asym --model /PATH/TO/SAVE/Llama-2-7b-chat-omniquant-w3a16g128 --use-cache=0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now you will have a new dist/ folder with the compiled model.You can refer the official [MLC-LLM docs](https://mlc.ai/mlc-llm/docs/) for more diverse usage of quantized models."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.11.0"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}