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": [
        "## Download the pre-quantized models"
      ]
    },
    {
      "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 model:"
      ]
    },
    {
      "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
}