bloom_3b_4bit_inference.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lwieske/bloom-and-qlora/blob/main/bloom_3b_4bit_inference.ipynb)\n",
        "[![Open in SageMaker Studio Lab](https://studiolab.sagemaker.aws/studiolab.svg)](https://studiolab.sagemaker.aws/import/github/lwieske/bloom-and-qlora/blob/main/bloom_3b_4bit_inference.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Ewy67IeWGS4"
      },
      "source": [
        "# BLOOM 3B 4bit Inference with Nvidia T4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2023-12-03T16:33:09.226833Z",
          "iopub.status.busy": "2023-12-03T16:33:09.225995Z",
          "iopub.status.idle": "2023-12-03T16:33:26.708129Z",
          "shell.execute_reply": "2023-12-03T16:33:26.706945Z",
          "shell.execute_reply.started": "2023-12-03T16:33:09.226799Z"
        },
        "id": "_UyyuMGnCPjA",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "!pip -q install accelerate bitsandbytes transformers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2023-12-03T16:33:26.710511Z",
          "iopub.status.busy": "2023-12-03T16:33:26.710215Z",
          "iopub.status.idle": "2023-12-03T16:33:43.177908Z",
          "shell.execute_reply": "2023-12-03T16:33:43.176975Z",
          "shell.execute_reply.started": "2023-12-03T16:33:26.710484Z"
        },
        "id": "pBOE-h8sbrNK",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2023-12-03T16:33:43.179663Z",
          "iopub.status.busy": "2023-12-03T16:33:43.179300Z",
          "iopub.status.idle": "2023-12-03T16:33:43.185046Z",
          "shell.execute_reply": "2023-12-03T16:33:43.184221Z",
          "shell.execute_reply.started": "2023-12-03T16:33:43.179628Z"
        },
        "id": "vziwd2UuCYGl",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "#MODEL      = \"bigscience/bloom-560m\"\n",
        "#MODEL      = \"bigscience/bloom-1b1\"\n",
        "#MODEL      = \"bigscience/bloom-1b7\"\n",
        "MODEL      = \"bigscience/bloom-3b\"\n",
        "#MODEL      = \"bigscience/bloom-7b1\"\n",
        "\n",
        "TOKENS     = 20"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2023-12-03T16:33:43.187660Z",
          "iopub.status.busy": "2023-12-03T16:33:43.187386Z",
          "iopub.status.idle": "2023-12-03T16:33:43.200847Z",
          "shell.execute_reply": "2023-12-03T16:33:43.199911Z",
          "shell.execute_reply.started": "2023-12-03T16:33:43.187635Z"
        },
        "id": "3h_ydWGf6EAd",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "bnb_config = BitsAndBytesConfig(\n",
        "                                    load_in_4bit=True,\n",
        "                                    bnb_4bit_use_double_quant=True,\n",
        "                                    bnb_4bit_quant_type=\"nf4\",\n",
        "                                    bnb_4bit_compute_dtype=torch.bfloat16\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2023-12-03T16:33:43.202355Z",
          "iopub.status.busy": "2023-12-03T16:33:43.202018Z",
          "iopub.status.idle": "2023-12-03T16:35:43.973028Z",
          "shell.execute_reply": "2023-12-03T16:35:43.972073Z",
          "shell.execute_reply.started": "2023-12-03T16:33:43.202322Z"
        },
        "id": "W2EZhNQ66EAd",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "pretrained_model = AutoModelForCausalLM.from_pretrained(\n",
        "                                                            MODEL,\n",
        "                                                            quantization_config=bnb_config,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "execution": {
          "iopub.execute_input": "2023-12-03T16:35:43.974501Z",
          "iopub.status.busy": "2023-12-03T16:35:43.974215Z",
          "iopub.status.idle": "2023-12-03T16:35:45.073154Z",
          "shell.execute_reply": "2023-12-03T16:35:45.072163Z",
          "shell.execute_reply.started": "2023-12-03T16:35:43.974478Z"
        },
        "id": "WWgahTxOcXEM",
        "outputId": "42263134-1794-4439-b786-ceb61036aac3",
        "trusted": true
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Mon Dec  4 10:24:42 2023       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   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   41C    P0    25W /  70W |   3107MiB / 15360MiB |      1%      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",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2023-12-03T16:35:45.075209Z",
          "iopub.status.busy": "2023-12-03T16:35:45.074797Z",
          "iopub.status.idle": "2023-12-03T16:35:46.443819Z",
          "shell.execute_reply": "2023-12-03T16:35:46.442667Z",
          "shell.execute_reply.started": "2023-12-03T16:35:45.075174Z"
        },
        "id": "aU0awofs84q7",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(MODEL)\n",
        "\n",
        "tokenizer.pad_token = tokenizer.eos_token"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c_fR4WttssU9",
        "outputId": "8e1c54ea-7fa9-4f66-b209-4ef9f7360ecf"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Mon Dec  4 10:24:43 2023       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   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   41C    P0    25W /  70W |   3107MiB / 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",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qtc1gbK39Hp7"
      },
      "source": [
        "## Inference with Pretrained BLOOM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2023-12-03T16:35:46.445908Z",
          "iopub.status.busy": "2023-12-03T16:35:46.445465Z",
          "iopub.status.idle": "2023-12-03T16:35:46.453059Z",
          "shell.execute_reply": "2023-12-03T16:35:46.451983Z",
          "shell.execute_reply.started": "2023-12-03T16:35:46.445862Z"
        },
        "id": "jak6FzpvFTHk",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "def generate_output_tokens(model, input_tokens):\n",
        "    outputs = model.generate(\n",
        "        input_ids=input_tokens[\"input_ids\"],\n",
        "        attention_mask=input_tokens[\"attention_mask\"],\n",
        "        max_new_tokens=TOKENS,\n",
        "        repetition_penalty=1.5,\n",
        "        early_stopping=False,\n",
        "        eos_token_id=tokenizer.eos_token_id,\n",
        "    )\n",
        "    return outputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "execution": {
          "iopub.execute_input": "2023-12-03T16:35:46.454742Z",
          "iopub.status.busy": "2023-12-03T16:35:46.454391Z",
          "iopub.status.idle": "2023-12-03T16:35:56.957515Z",
          "shell.execute_reply": "2023-12-03T16:35:56.956523Z",
          "shell.execute_reply.started": "2023-12-03T16:35:46.454706Z"
        },
        "id": "3BAYg7czFYeK",
        "outputId": "1e3b2ffb-15aa-49f5-8ead-4cbaf536effa",
        "trusted": true
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['What is the meaning of life? What does it mean to be human?\\nThe answer lies in our relationship with God. The Bible says']\n"
          ]
        }
      ],
      "source": [
        "input_words = \"What is the meaning of life?\"\n",
        "\n",
        "input_tokens = tokenizer(input_words, return_tensors=\"pt\").to('cuda')\n",
        "\n",
        "pretrained_output_tokens = generate_output_tokens(pretrained_model, input_tokens)\n",
        "\n",
        "pretrained_output_words = tokenizer.batch_decode(pretrained_output_tokens, skip_special_tokens=True)\n",
        "\n",
        "print(pretrained_output_words)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ttIVW6j7ssU_",
        "outputId": "e6bccd90-e1e8-4fac-f463-4a83db8462cb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Mon Dec  4 10:24:53 2023       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   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   43C    P0    43W /  70W |   3639MiB / 15360MiB |     46%      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",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "!nvidia-smi"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "machine_shape": "hm",
      "provenance": []
    },
    "kaggle": {
      "accelerator": "nvidiaTeslaT4",
      "dataSources": [],
      "dockerImageVersionId": 30588,
      "isGpuEnabled": true,
      "isInternetEnabled": true,
      "language": "python",
      "sourceType": "notebook"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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": 0
}