BLOOM_BigScience_bitsandbytes_int8.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/R3gm/Colab-resources/blob/main/BLOOM_BigScience_bitsandbytes_int8.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gbr9iYfRy4GZ"
      },
      "source": [
        "# HuggingFace meets `bitsandbytes` for lighter models on GPU for inference"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zTr9qoBgbX_Z"
      },
      "source": [
        "Bloom: https://huggingface.co/docs/transformers/model_doc/bloom"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "| Code Credits | Link |\n",
        "| ----------- | ---- |\n",
        "| 🎉 Repository | [![GitHub Repository](https://img.shields.io/github/stars/TimDettmers/bitsandbytes?style=social)](https://github.com/TimDettmers/bitsandbytes) |\n",
        "| 🔥 Discover More Colab Notebooks | [![GitHub Repository](https://img.shields.io/badge/GitHub-Repository-black?style=flat-square&logo=github)](https://github.com/R3gm/InsightSolver-Colab/) |\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K9uJn2NczFBK"
      },
      "source": [
        "\n",
        "You can run your own 8-bit model on any HuggingFace 🤗 model with just few lines of code. Install the dependencies below first!\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0zvx56Q8Jvu4",
        "outputId": "1d18071f-b701-4a7b-b81b-93ac22871773"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.1/7.1 MB\u001b[0m \u001b[31m47.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m219.1/219.1 kB\u001b[0m \u001b[31m4.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m104.3/104.3 MB\u001b[0m \u001b[31m8.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m108.2/108.2 MB\u001b[0m \u001b[31m9.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m224.5/224.5 kB\u001b[0m \u001b[31m11.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.8/7.8 MB\u001b[0m \u001b[31m92.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "! pip install -q transformers accelerate bitsandbytes xformers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W4Cu-f2hjTza"
      },
      "source": [
        "## Hardware requirements 🔨"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TLl929-ajXH8"
      },
      "source": [
        "To run properly this feature you need to have GPU that supports 8-bit operation modules. Currently, Turing and Ampere GPUs (RTX20s, RTX30s, A40-A100, T4+) are supported, which means on colab we need to use a T4 GPU for this feature. You can check that using this code snippet and make sure you are using a supported GPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JDF6faQPjiuH",
        "outputId": "d072dea0-e8d8-4c9b-e1df-394a86cd9653"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Fri May 12 00:51:33 2023       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   47C    P8    10W /  70W |      0MiB / 15360MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Not connected to a GPU')\n",
        "else:\n",
        "  print(gpu_info)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jqSXxf9NlQJt"
      },
      "source": [
        "Here we are using a `Tesla T4` GPU that should support 8-bit tensor cores! We are good to go 🚀"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_rWAkaq05NdH"
      },
      "source": [
        "## Utility variables & functions 🧰"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JM1NgY4B5Mx2"
      },
      "outputs": [],
      "source": [
        "name = \"bigscience/bloom-3b\"\n",
        "#name = \"facebook/opt-1.3b\"\n",
        "text = \"Hello my name is\"\n",
        "max_new_tokens = 20\n",
        "\n",
        "def generate_from_model(model, tokenizer):\n",
        "  encoded_input = tokenizer(text, return_tensors='pt')\n",
        "  output_sequences = model.generate(input_ids=encoded_input['input_ids'].cuda())\n",
        "  return tokenizer.decode(output_sequences[0], skip_special_tokens=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uxyjbQG54ZPm"
      },
      "source": [
        "## Use 8bit models and `pipeline` 🤗"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z7YKvh8HCF32"
      },
      "source": [
        "You can use 8bit quantized models together with `pipeline` as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9VOMsHmY37Sx",
        "outputId": "b8857f1a-e951-4e12-859a-52a7aef255a5"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Overriding torch_dtype=None with `torch_dtype=torch.float16` due to requirements of `bitsandbytes` to enable model loading in mixed int8. Either pass torch_dtype=torch.float16 or don't pass this argument at all to remove this warning.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "===================================BUG REPORT===================================\n",
            "Welcome to bitsandbytes. For bug reports, please run\n",
            "\n",
            "python -m bitsandbytes\n",
            "\n",
            " and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
            "================================================================================\n",
            "bin /usr/local/lib/python3.10/dist-packages/bitsandbytes/libbitsandbytes_cuda118.so\n",
            "CUDA_SETUP: WARNING! libcudart.so not found in any environmental path. Searching in backup paths...\n",
            "CUDA SETUP: CUDA runtime path found: /usr/local/cuda/lib64/libcudart.so\n",
            "CUDA SETUP: Highest compute capability among GPUs detected: 7.5\n",
            "CUDA SETUP: Detected CUDA version 118\n",
            "CUDA SETUP: Loading binary /usr/local/lib/python3.10/dist-packages/bitsandbytes/libbitsandbytes_cuda118.so...\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: /usr/lib64-nvidia did not contain ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] as expected! Searching further paths...\n",
            "  warn(msg)\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/sys/fs/cgroup/memory.events /var/colab/cgroup/jupyter-children/memory.events')}\n",
            "  warn(msg)\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('//172.28.0.1'), PosixPath('http'), PosixPath('8013')}\n",
            "  warn(msg)\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('//colab.research.google.com/tun/m/cc48301118ce562b961b3c22d803539adc1e0c19/gpu-t4-s-3iyksaxwto648 --tunnel_background_save_delay=10s --tunnel_periodic_background_save_frequency=30m0s --enable_output_coalescing=true --output_coalescing_required=true'), PosixPath('--logtostderr --listen_host=172.28.0.12 --target_host=172.28.0.12 --tunnel_background_save_url=https')}\n",
            "  warn(msg)\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/env/python')}\n",
            "  warn(msg)\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('//ipykernel.pylab.backend_inline'), PosixPath('module')}\n",
            "  warn(msg)\n",
            "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py:145: UserWarning: Found duplicate ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] files: {PosixPath('/usr/local/cuda/lib64/libcudart.so'), PosixPath('/usr/local/cuda/lib64/libcudart.so.11.0')}.. We'll flip a coin and try one of these, in order to fail forward.\n",
            "Either way, this might cause trouble in the future:\n",
            "If you get `CUDA error: invalid device function` errors, the above might be the cause and the solution is to make sure only one ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] in the paths that we search based on your env.\n",
            "  warn(msg)\n"
          ]
        }
      ],
      "source": [
        "from transformers import pipeline\n",
        "\n",
        "pipe = pipeline(model=name, model_kwargs= {\"device_map\": \"auto\", \"load_in_8bit\": True}, max_new_tokens=max_new_tokens)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7oPPkt-IChHn"
      },
      "source": [
        "Let's check the output!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6xYYcAxO4v4m",
        "outputId": "cd04c1d2-2cce-480f-9a4e-8faed6eeeb99"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[{'generated_text': 'Hello my name is John and I am a student at the University of the West of England. I am currently studying for'}]"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "pipe(text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BejbQStU5Eik"
      },
      "source": [
        "## Use 8bit models and `.generate` 📖"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YJlldexxwnhM",
        "outputId": "10ae1616-56dd-4fe1-9dec-18606b5edee6"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Overriding torch_dtype=None with `torch_dtype=torch.float16` due to requirements of `bitsandbytes` to enable model loading in mixed int8. Either pass torch_dtype=torch.float16 or don't pass this argument at all to remove this warning.\n"
          ]
        }
      ],
      "source": [
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "import torch\n",
        "\n",
        "model_8bit = AutoModelForCausalLM.from_pretrained(name, device_map=\"auto\", load_in_8bit=True)\n",
        "tokenizer = AutoTokenizer.from_pretrained(name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 93
        },
        "id": "Uol1MHIRyFDn",
        "outputId": "49ccfd40-f649-49d7-bfe2-9ab0b8792d84"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py:1346: UserWarning: Using `max_length`'s default (20) to control the generation length. This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we recommend using `max_new_tokens` to control the maximum length of the generation.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'Hello my name is John and I am a student at the University of the West of England. I'"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "generate_from_model(model_8bit, tokenizer)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oIltcTCscANs"
      },
      "outputs": [],
      "source": [
        "# memory size\n",
        "mem_int8 = model_8bit.get_memory_footprint()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AatjsZQbdw_Q",
        "outputId": "cb2c0914-13aa-4a9d-be17-9960c405a09c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Memory footprint int8 model: 3645818880 \n"
          ]
        }
      ],
      "source": [
        "print(\"Memory footprint int8 model: {} \".format(mem_int8))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vu2T9mWr5Vpx"
      },
      "source": [
        "Let's compare the qualitative results between our quantized model and the original model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 93
        },
        "id": "JECvvs3SzzeM",
        "outputId": "d761a38f-967f-4a90-9713-1bf2bd60cb2e"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py:1346: UserWarning: Using `max_length`'s default (20) to control the generation length. This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we recommend using `max_new_tokens` to control the maximum length of the generation.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'Hello my name is John and I am a student at the University of the West Indies. I am'"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "name = \"bigscience/bloom-3b\"\n",
        "#name = \"facebook/opt-1.3b\"\n",
        "text = \"Hello my name is\"\n",
        "max_new_tokens = 20\n",
        "\n",
        "def generate_from_model(model, tokenizer):\n",
        "  encoded_input = tokenizer(text, return_tensors='pt')\n",
        "  output_sequences = model.generate(input_ids=encoded_input['input_ids'].cuda())\n",
        "  return tokenizer.decode(output_sequences[0], skip_special_tokens=True)\n",
        "\n",
        "model_native = AutoModelForCausalLM.from_pretrained(name, device_map=\"auto\", torch_dtype=\"auto\")\n",
        "tokenizer = AutoTokenizer.from_pretrained(name)\n",
        "generate_from_model(model_native, tokenizer)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iYJom2z-cB_e"
      },
      "outputs": [],
      "source": [
        "# memory size\n",
        "mem_fp16 = model_native.get_memory_footprint()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ODzs-N2SeUe0",
        "outputId": "0f1c3012-792a-4c70-bb33-70ddb56fc9be"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Memory footprint fp16 model: 12010229760\n"
          ]
        }
      ],
      "source": [
        "print(\"Memory footprint fp16 model: {}\".format( mem_fp16))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aoIOaIxoCnBe"
      },
      "source": [
        "## Memory footprint comparison 🪶"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RtoXgrDUz8gc",
        "outputId": "f62bb9d5-dcb9-4573-b50a-4ac7e33ebbd0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Memory footprint int8 model: 2236858368 | Memory footprint fp16 model: 6889635840 | Relative difference: 3.0800500999802236\n"
          ]
        }
      ],
      "source": [
        "#print(\"Memory footprint int8 model: {} | Memory footprint fp16 model: {} | Relative difference: {}\".format(mem_int8, mem_fp16, mem_fp16/mem_int8))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dFsxYDzVEAeJ"
      },
      "source": [
        "We saved 1.65x memory for a 3-billion parameters models! Note that internally we replace all the linear layers by the ones implemented in `bitsandbytes`. By scaling up the model the number of linear layers will increase therefore the impact of saving memory on those layers will be huge for very large models. For example quantizing BLOOM-176 (176 Billion parameter model) gives a gain of 1.96x memory footprint which can save a lot of compute power in practice."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "include_colab_link": true,
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}