runing_falcon180b_on_single_a100_80g.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EWOtpjJMDmYE"
      },
      "source": [
        "# Runing Falcon-180B on a single A100 80GB"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This file including three section:\n",
        "- [(Optional) Train the quantization parameters of Falcon-180B by yourself.](#optional-train-the-quantization-parameters-of-falcon-180b-by-yourself)\n",
        "- [Download the pre-quantized models](#download-the-pre-quantized-models)\n",
        "- [Let's Infer!](#lets-infer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## (Optional) Train the quantization parameters of Falcon-180B by yourself.\n",
        "\n",
        "This section provids how to train the quantization parameters of Falcon-180B by yourself. You can skip this section because we have provided the pre-built quantized models in [Download the pre-quantized models](#download-the-pre-quantized-models).\n",
        "Additionally, This process wound take about 12 hours on a single A100-80GB GPU."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "1. Training the quantization parameters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!CUDA_VISIBLE_DEVICES=0 python main.py \\\n",
        "--model /PATH/TO/Falcon/falcon-180b \\\n",
        "--epochs 40 --output_dir ./log/falcon-180b-w3a16g512 \\\n",
        "--wbits 3 --abits 16 --group_size 512 --lwc --aug_loss \\\n",
        "--nsamples 32"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "2. Saving the really quantized model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!CUDA_VISIBLE_DEVICES=0 python main.py \\\n",
        "--model /PATH/TO/Falcon/falcon-180b \\\n",
        "--epochs 0  --output_dir ./log/temp \\\n",
        "--wbits 3 --abits 16 --group_size 512 --nsamples 32 \\\n",
        "--resume ./log/falcon-180b-w3a16g512/omni_parameters.pth \\\n",
        "--real_quant --save_dir /PATH/TO/SAVE/MODELS \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Download the prebuilt quantized model:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FwsWd1WbDmYE"
      },
      "source": [
        "We have provide the prebuilt quantized model on 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": "code",
      "execution_count": null,
      "metadata": {
        "id": "FSAe7Ew_DmYF"
      },
      "outputs": [],
      "source": [
        "!mkdir -p pre_quantized_models/\n",
        "# download falcon-180b with w3a16g512 quantization\n",
        "!git clone https://huggingface.co/ChenMnZ/falcon-180b-omniquant-w3a16g512 ./pre_quantized_models/falcon-180b-omniquant-w3a16g512"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "76Ru5__tDmYF"
      },
      "source": [
        "## Let's Infer!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Constraint in one GPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/cpfs01/user/chenmengzhao/anaconda3/envs/smoothquant/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[2023-09-11 05:21:59,584] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 80/80 [00:06<00:00, 12.17it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading pre-computed quantized weights...\n",
            "Loading pre-computed quantized weights Successfully\n"
          ]
        }
      ],
      "source": [
        "from accelerate import init_empty_weights, infer_auto_device_map, load_checkpoint_in_model\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig\n",
        "import torch\n",
        "import auto_gptq.nn_modules.qlinear.qlinear_cuda as qlinear_cuda\n",
        "from transformers.models.falcon.modeling_falcon import FalconLinear\n",
        "from tqdm import tqdm\n",
        "import gc   \n",
        "import time\n",
        "\n",
        "def get_named_linears(module):\n",
        "    return {name: m for name, m in module.named_modules() if isinstance(m, FalconLinear)}\n",
        "\n",
        "def set_op_by_name(layer, name, new_module):\n",
        "    levels = name.split('.')\n",
        "    if len(levels) > 1:\n",
        "        mod_ = layer\n",
        "        for l_idx in range(len(levels)-1):\n",
        "            if levels[l_idx].isdigit():\n",
        "                mod_ = mod_[int(levels[l_idx])]\n",
        "            else:\n",
        "                mod_ = getattr(mod_, levels[l_idx])\n",
        "        setattr(mod_, levels[-1], new_module)\n",
        "    else:\n",
        "        setattr(layer, name, new_module)\n",
        "\n",
        "model_path = './pre_quantized_models/falcon-180b-omniquant-w3a16g512'\n",
        "wbits = 3\n",
        "group_size = 512\n",
        "config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)\n",
        "enc = AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=True)\n",
        "with init_empty_weights():\n",
        "    model = AutoModelForCausalLM.from_config(config=config,torch_dtype=torch.float16, trust_remote_code=True)\n",
        "\n",
        "layers = model.transformer.h\n",
        "for i in tqdm(range(len(layers))):\n",
        "    layer = layers[i]\n",
        "    named_linears = get_named_linears(layer)\n",
        "    for name, module in named_linears.items():\n",
        "        q_linear = qlinear_cuda.QuantLinear(wbits, group_size, module.in_features,module.out_features,not module.bias is None,kernel_switch_threshold=128)\n",
        "        q_linear.to(next(layer.parameters()).device)\n",
        "        set_op_by_name(layer, name, q_linear)\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "model.tie_weights()\n",
        "device_map = infer_auto_device_map(model)\n",
        "print(\"Loading pre-computed quantized weights...\")\n",
        "load_checkpoint_in_model(model,checkpoint=model_path,device_map=device_map,offload_state_dict=True)\n",
        "print(\"Loading pre-computed quantized weights Successfully\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Start inference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
            "Setting `pad_token_id` to `eos_token_id`:11 for open-end generation.\n",
            "The current implementation of Falcon calls `torch.scaled_dot_product_attention` directly, this will be deprecated in the future in favor of the `BetterTransformer` API. Please install the latest optimum library with `pip install -U optimum` and call `model.to_bettertransformer()` to benefit from `torch.scaled_dot_product_attention` and future performance optimizations.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Give me a list of the top 10 dive sites you would recommend around the world. \n",
            "The list is:\n",
            "The Red Sea \n",
            "The Great Barrier Reef\n",
            "The Cayman Islands \n",
            "Belize \n",
            "The Bahamas \n",
            "The Galápagos Islands \n",
            "Palau \n",
            "Hawaii \n",
            "Thailand \n",
            "Fiji\n",
            "\n",
            "This is an article I found on a dive site, so I am sure the list is accurate.\n",
            "\n",
            "I am not sure about the \"best dive sites\" because that is a subjective question. However, the top ten list of dive sites that are most popular, most visited, and most famous would be:\n",
            "1. Red Sea, Egypt\n",
            "2. Great Barrier Reef, Australia\n",
            "3\n",
            "speed:0.4184875772231599token/s\n"
          ]
        }
      ],
      "source": [
        "model.eval()\n",
        "prompt = \"Give me a list of the top 10 dive sites you would recommend around the world. \\nThe list is:\"\n",
        "input_ids = enc(prompt, return_tensors='pt').input_ids.cuda()\n",
        "model = model.cuda()\n",
        "start_time = time.time()\n",
        "output = model.generate(inputs=input_ids, do_sample=True, top_k=10, max_new_tokens=128)\n",
        "end_time = time.time()\n",
        "speed = len(output[0])/(end_time-start_time)\n",
        "print(enc.decode(output[0]))\n",
        "print(f\"speed:{speed}token/s\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Although the quantized Falcon-180b can be loaded onto a single A100 80GB, its inference speed remains slow due to CUDA kernel incompatibility. Kernel improvements are in progress."
      ]
    }
  ],
  "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.8.16"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}