python/models/llama/LLaMA-2 E2E Notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) Microsoft Corporation. All rights reserved.  \n",
    "Licensed under the MIT License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Accelerating LLaMA-2 Inference with ONNX Runtime\n",
    "\n",
    "In this tutorial, you will export, optimize, and run the LLaMA-2 model using ONNX Runtime.\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "0. Use a machine with at least 64GB of memory. Exporting LLaMA-2 requires a significant amount of memory because of the model's size."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Install [Anaconda](https://www.anaconda.com/distribution/). Once installed, create a `conda` environment named `llama2` by running the following in your terminal (outside of this notebook).\n",
    "\n",
    "```console\n",
    "$ conda create -n llama2 python=3.9\n",
    "$ conda activate llama2\n",
    "```\n",
    "\n",
    "If you don't have Jupyter installed to run this notebook, here is how you can install it and connect it to your new `conda` environment (run in your terminal outside of this notebook).\n",
    "```console\n",
    "$ pip install jupyterlab\n",
    "$ conda install ipykernel\n",
    "$ conda install -c conda-forge ipywidgets\n",
    "$ ipython kernel install --user --name llama2\n",
    "$ jupyter-lab\n",
    "```\n",
    "\n",
    "Once you have this notebook open in Jupyter, you can select the `llama2` environment that you created as the kernel for this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Select the `torch` package for your environment. For this notebook, you need to install `torch` with CUDA enabled for your installed CUDA version.\n",
    "\n",
    "First, you need to identify your installed CUDA version. Run the following command in your terminal (outside of this notebook).\n",
    "\n",
    "```console\n",
    "$ nvidia-smi\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A table should print that shows the status of your GPUs. Your CUDA version is located in the top right of the table. For this notebook, CUDA 11.8 is used as the version.\n",
    "\n",
    "Once you have identified your CUDA version, you need to visit the [PyTorch website](https://pytorch.org/get-started/locally/) to get the download instructions for your specific `torch` package. For CUDA version 11.8 and lower, you can select `CUDA 11.8` on the PyTorch website. For CUDA version 12.0 and higher, you can select `CUDA 12.1` on the PyTorch website.\n",
    "\n",
    "Important: You must select the `Preview (Nightly)` PyTorch build option. Otherwise, the export may fail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Skipping torch as it is not installed.\u001b[0m\u001b[33m\n",
      "\u001b[0mLooking in indexes: https://download.pytorch.org/whl/nightly/cu118\n",
      "Collecting torch\n",
      "  Using cached https://download.pytorch.org/whl/nightly/cu118/torch-2.2.0.dev20231111%2Bcu118-cp39-cp39-linux_x86_64.whl (2532.6 MB)\n",
      "Collecting filelock (from torch)\n",
      "  Using cached https://download.pytorch.org/whl/nightly/filelock-3.9.0-py3-none-any.whl (9.7 kB)\n",
      "Requirement already satisfied: typing-extensions in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from torch) (4.7.1)\n",
      "Collecting sympy (from torch)\n",
      "  Using cached https://download.pytorch.org/whl/nightly/sympy-1.11.1-py3-none-any.whl (6.5 MB)\n",
      "Collecting networkx (from torch)\n",
      "  Using cached https://download.pytorch.org/whl/nightly/networkx-3.0rc1-py3-none-any.whl (2.0 MB)\n",
      "Requirement already satisfied: jinja2 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from torch) (3.1.2)\n",
      "Collecting fsspec (from torch)\n",
      "  Using cached https://download.pytorch.org/whl/nightly/fsspec-2023.4.0-py3-none-any.whl (153 kB)\n",
      "Collecting pytorch-triton==2.1.0+6e4932cda8 (from torch)\n",
      "  Using cached https://download.pytorch.org/whl/nightly/pytorch_triton-2.1.0%2B6e4932cda8-cp39-cp39-linux_x86_64.whl (125.4 MB)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from jinja2->torch) (2.1.3)\n",
      "Collecting mpmath>=0.19 (from sympy->torch)\n",
      "  Using cached https://download.pytorch.org/whl/nightly/mpmath-1.2.1-py3-none-any.whl (532 kB)\n",
      "Installing collected packages: mpmath, sympy, networkx, fsspec, filelock, pytorch-triton, torch\n",
      "Successfully installed filelock-3.9.0 fsspec-2023.4.0 mpmath-1.2.1 networkx-3.0rc1 pytorch-triton-2.1.0+6e4932cda8 sympy-1.11.1 torch-2.2.0.dev20231111+cu118\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "\n",
    "# Uninstall existing torch\n",
    "!{sys.executable} -m pip uninstall -y torch\n",
    "\n",
    "# Example installation command for CUDA 11.8\n",
    "!{sys.executable} -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu118"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. Install the packages from the `requirements-*.txt` file that [fits your scenario](https://github.com/microsoft/onnxruntime/tree/main/onnxruntime/python/tools/transformers/models/llama). Since this notebook is showing inference with CUDA, you can use `requirements-cuda.txt`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2023-11-12 08:59:28--  https://raw.githubusercontent.com/microsoft/onnxruntime/main/onnxruntime/python/tools/transformers/models/llama/requirements-cuda.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.109.133, 185.199.110.133, 185.199.108.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.109.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 234 [text/plain]\n",
      "Saving to: ‘requirements-cuda.txt’\n",
      "\n",
      "requirements-cuda.t 100%[===================>]     234  --.-KB/s    in 0s      \n",
      "\n",
      "2023-11-12 08:59:28 (24.4 MB/s) - ‘requirements-cuda.txt’ saved [234/234]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Replace requirements-*.txt filename with the one for your scenario\n",
    "!wget https://raw.githubusercontent.com/microsoft/onnxruntime/main/onnxruntime/python/tools/transformers/models/llama/requirements-cuda.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you have downloaded the desired `requirements-*.txt` file, you need to download the common `requirements.txt` file that contains all shared packages across the `requirements-*.txt` files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2023-11-12 08:59:33--  https://raw.githubusercontent.com/microsoft/onnxruntime/main/onnxruntime/python/tools/transformers/models/llama/requirements.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.109.133, 185.199.110.133, 185.199.108.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.109.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 138 [text/plain]\n",
      "Saving to: ‘requirements.txt’\n",
      "\n",
      "requirements.txt    100%[===================>]     138  --.-KB/s    in 0s      \n",
      "\n",
      "2023-11-12 08:59:34 (11.3 MB/s) - ‘requirements.txt’ saved [138/138]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget https://raw.githubusercontent.com/microsoft/onnxruntime/main/onnxruntime/python/tools/transformers/models/llama/requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once downloaded, you can install the packages from the `requirements-*.txt` file for your scenario. It will use `requirements.txt` to download the shared packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting git+https://github.com/huggingface/optimum.git (from -r requirements.txt (line 1))\n",
      "  Cloning https://github.com/huggingface/optimum.git to /tmp/pip-req-build-6v3xkch4\n",
      "  Running command git clone --filter=blob:none --quiet https://github.com/huggingface/optimum.git /tmp/pip-req-build-6v3xkch4\n",
      "  Resolved https://github.com/huggingface/optimum.git to commit 832f3b292b501c7ab920ecc7913de3c6b7894d60\n",
      "  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25hCollecting transformers>=4.33.2 (from -r requirements.txt (line 2))\n",
      "  Using cached transformers-4.35.0-py3-none-any.whl.metadata (123 kB)\n",
      "Requirement already satisfied: torch>=2.2.0.dev20230920 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from -r requirements.txt (line 3)) (2.2.0.dev20231111+cu118)\n",
      "Collecting onnx>=1.14.0 (from -r requirements.txt (line 4))\n",
      "  Using cached onnx-1.15.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (15 kB)\n",
      "Collecting datasets>=2.8.0 (from -r requirements.txt (line 5))\n",
      "  Using cached datasets-2.14.6-py3-none-any.whl.metadata (19 kB)\n",
      "Collecting protobuf==3.20.2 (from -r requirements.txt (line 6))\n",
      "  Using cached protobuf-3.20.2-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.whl (1.0 MB)\n",
      "Collecting onnxruntime-gpu>=1.16.2 (from -r requirements-cuda.txt (line 4))\n",
      "  Using cached onnxruntime_gpu-1.16.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (4.2 kB)\n",
      "Collecting coloredlogs (from optimum==1.15.0.dev0->-r requirements.txt (line 1))\n",
      "  Using cached coloredlogs-15.0.1-py2.py3-none-any.whl (46 kB)\n",
      "Requirement already satisfied: sympy in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from optimum==1.15.0.dev0->-r requirements.txt (line 1)) (1.11.1)\n",
      "Requirement already satisfied: packaging in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from optimum==1.15.0.dev0->-r requirements.txt (line 1)) (23.1)\n",
      "Collecting numpy (from optimum==1.15.0.dev0->-r requirements.txt (line 1))\n",
      "  Using cached numpy-1.26.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (61 kB)\n",
      "Collecting huggingface-hub>=0.8.0 (from optimum==1.15.0.dev0->-r requirements.txt (line 1))\n",
      "  Using cached huggingface_hub-0.19.0-py3-none-any.whl.metadata (13 kB)\n",
      "Requirement already satisfied: filelock in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from transformers>=4.33.2->-r requirements.txt (line 2)) (3.9.0)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from transformers>=4.33.2->-r requirements.txt (line 2)) (6.0.1)\n",
      "Collecting regex!=2019.12.17 (from transformers>=4.33.2->-r requirements.txt (line 2))\n",
      "  Using cached regex-2023.10.3-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (40 kB)\n",
      "Requirement already satisfied: requests in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from transformers>=4.33.2->-r requirements.txt (line 2)) (2.31.0)\n",
      "Collecting tokenizers<0.15,>=0.14 (from transformers>=4.33.2->-r requirements.txt (line 2))\n",
      "  Using cached tokenizers-0.14.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (6.7 kB)\n",
      "Collecting safetensors>=0.3.1 (from transformers>=4.33.2->-r requirements.txt (line 2))\n",
      "  Using cached safetensors-0.4.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (3.8 kB)\n",
      "Collecting tqdm>=4.27 (from transformers>=4.33.2->-r requirements.txt (line 2))\n",
      "  Using cached tqdm-4.66.1-py3-none-any.whl.metadata (57 kB)\n",
      "Requirement already satisfied: typing-extensions in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from torch>=2.2.0.dev20230920->-r requirements.txt (line 3)) (4.7.1)\n",
      "Requirement already satisfied: networkx in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from torch>=2.2.0.dev20230920->-r requirements.txt (line 3)) (3.0rc1)\n",
      "Requirement already satisfied: jinja2 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from torch>=2.2.0.dev20230920->-r requirements.txt (line 3)) (3.1.2)\n",
      "Requirement already satisfied: fsspec in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from torch>=2.2.0.dev20230920->-r requirements.txt (line 3)) (2023.4.0)\n",
      "Requirement already satisfied: pytorch-triton==2.1.0+6e4932cda8 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from torch>=2.2.0.dev20230920->-r requirements.txt (line 3)) (2.1.0+6e4932cda8)\n",
      "Collecting pyarrow>=8.0.0 (from datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached pyarrow-14.0.1-cp39-cp39-manylinux_2_28_x86_64.whl.metadata (3.0 kB)\n",
      "Collecting dill<0.3.8,>=0.3.0 (from datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached dill-0.3.7-py3-none-any.whl.metadata (9.9 kB)\n",
      "Collecting pandas (from datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached pandas-2.1.3-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (18 kB)\n",
      "Collecting xxhash (from datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached xxhash-3.4.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (12 kB)\n",
      "Collecting multiprocess (from datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached multiprocess-0.70.15-py39-none-any.whl.metadata (7.2 kB)\n",
      "Collecting aiohttp (from datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached aiohttp-3.8.6-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (7.7 kB)\n",
      "Collecting flatbuffers (from onnxruntime-gpu>=1.16.2->-r requirements-cuda.txt (line 4))\n",
      "  Using cached flatbuffers-23.5.26-py2.py3-none-any.whl.metadata (850 bytes)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from aiohttp->datasets>=2.8.0->-r requirements.txt (line 5)) (23.1.0)\n",
      "Requirement already satisfied: charset-normalizer<4.0,>=2.0 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from aiohttp->datasets>=2.8.0->-r requirements.txt (line 5)) (3.3.2)\n",
      "Collecting multidict<7.0,>=4.5 (from aiohttp->datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached multidict-6.0.4-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (114 kB)\n",
      "Collecting async-timeout<5.0,>=4.0.0a3 (from aiohttp->datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached async_timeout-4.0.3-py3-none-any.whl.metadata (4.2 kB)\n",
      "Collecting yarl<2.0,>=1.0 (from aiohttp->datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached yarl-1.9.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (269 kB)\n",
      "Collecting frozenlist>=1.1.1 (from aiohttp->datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached frozenlist-1.4.0-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (5.2 kB)\n",
      "Collecting aiosignal>=1.1.2 (from aiohttp->datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached aiosignal-1.3.1-py3-none-any.whl (7.6 kB)\n",
      "INFO: pip is looking at multiple versions of huggingface-hub to determine which version is compatible with other requirements. This could take a while.\n",
      "Collecting huggingface-hub>=0.8.0 (from optimum==1.15.0.dev0->-r requirements.txt (line 1))\n",
      "  Using cached huggingface_hub-0.18.0-py3-none-any.whl.metadata (13 kB)\n",
      "  Using cached huggingface_hub-0.17.3-py3-none-any.whl.metadata (13 kB)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from requests->transformers>=4.33.2->-r requirements.txt (line 2)) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from requests->transformers>=4.33.2->-r requirements.txt (line 2)) (2.0.7)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from requests->transformers>=4.33.2->-r requirements.txt (line 2)) (2023.7.22)\n",
      "Collecting sentencepiece!=0.1.92,>=0.1.91 (from transformers[sentencepiece]>=4.26.0->optimum==1.15.0.dev0->-r requirements.txt (line 1))\n",
      "  Using cached sentencepiece-0.1.99-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
      "Collecting humanfriendly>=9.1 (from coloredlogs->optimum==1.15.0.dev0->-r requirements.txt (line 1))\n",
      "  Using cached humanfriendly-10.0-py2.py3-none-any.whl (86 kB)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from jinja2->torch>=2.2.0.dev20230920->-r requirements.txt (line 3)) (2.1.3)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from pandas->datasets>=2.8.0->-r requirements.txt (line 5)) (2.8.2)\n",
      "Collecting pytz>=2020.1 (from pandas->datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached pytz-2023.3.post1-py2.py3-none-any.whl.metadata (22 kB)\n",
      "Collecting tzdata>=2022.1 (from pandas->datasets>=2.8.0->-r requirements.txt (line 5))\n",
      "  Using cached tzdata-2023.3-py2.py3-none-any.whl (341 kB)\n",
      "Requirement already satisfied: mpmath>=0.19 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from sympy->optimum==1.15.0.dev0->-r requirements.txt (line 1)) (1.2.1)\n",
      "Requirement already satisfied: six>=1.5 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from python-dateutil>=2.8.2->pandas->datasets>=2.8.0->-r requirements.txt (line 5)) (1.16.0)\n",
      "Using cached transformers-4.35.0-py3-none-any.whl (7.9 MB)\n",
      "Using cached onnx-1.15.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (15.7 MB)\n",
      "Using cached datasets-2.14.6-py3-none-any.whl (493 kB)\n",
      "Using cached onnxruntime_gpu-1.16.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (157.1 MB)\n",
      "Using cached dill-0.3.7-py3-none-any.whl (115 kB)\n",
      "Using cached aiohttp-3.8.6-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.1 MB)\n",
      "Using cached huggingface_hub-0.17.3-py3-none-any.whl (295 kB)\n",
      "Using cached numpy-1.26.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (18.2 MB)\n",
      "Using cached pyarrow-14.0.1-cp39-cp39-manylinux_2_28_x86_64.whl (38.0 MB)\n",
      "Using cached regex-2023.10.3-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (773 kB)\n",
      "Using cached safetensors-0.4.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
      "Using cached tokenizers-0.14.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.8 MB)\n",
      "Using cached tqdm-4.66.1-py3-none-any.whl (78 kB)\n",
      "Using cached flatbuffers-23.5.26-py2.py3-none-any.whl (26 kB)\n",
      "Using cached multiprocess-0.70.15-py39-none-any.whl (133 kB)\n",
      "Using cached pandas-2.1.3-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (12.3 MB)\n",
      "Using cached xxhash-3.4.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (193 kB)\n",
      "Using cached async_timeout-4.0.3-py3-none-any.whl (5.7 kB)\n",
      "Using cached frozenlist-1.4.0-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (228 kB)\n",
      "Using cached pytz-2023.3.post1-py2.py3-none-any.whl (502 kB)\n",
      "Building wheels for collected packages: optimum\n",
      "  Building wheel for optimum (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for optimum: filename=optimum-1.15.0.dev0-py3-none-any.whl size=399756 sha256=b06d30f0227068397a88bdd46db42138c01383673c55dbf542d2671be20e967b\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-i3f8o2c0/wheels/a2/2f/00/2a677c0c9bccd235c10e1cda441d7bb0d48ab534cff60239e1\n",
      "Successfully built optimum\n",
      "Installing collected packages: sentencepiece, pytz, flatbuffers, xxhash, tzdata, tqdm, safetensors, regex, protobuf, numpy, multidict, humanfriendly, frozenlist, dill, async-timeout, yarl, pyarrow, pandas, onnx, multiprocess, huggingface-hub, coloredlogs, aiosignal, tokenizers, onnxruntime-gpu, aiohttp, transformers, datasets, optimum\n",
      "Successfully installed aiohttp-3.8.6 aiosignal-1.3.1 async-timeout-4.0.3 coloredlogs-15.0.1 datasets-2.14.6 dill-0.3.7 flatbuffers-23.5.26 frozenlist-1.4.0 huggingface-hub-0.17.3 humanfriendly-10.0 multidict-6.0.4 multiprocess-0.70.15 numpy-1.26.1 onnx-1.15.0 onnxruntime-gpu-1.16.2 optimum-1.15.0.dev0 pandas-2.1.3 protobuf-3.20.2 pyarrow-14.0.1 pytz-2023.3.post1 regex-2023.10.3 safetensors-0.4.0 sentencepiece-0.1.99 tokenizers-0.14.1 tqdm-4.66.1 transformers-4.35.0 tzdata-2023.3 xxhash-3.4.1 yarl-1.9.2\n"
     ]
    }
   ],
   "source": [
    "!{sys.executable} -m pip install -r requirements-cuda.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. Create an account on Hugging Face to access the LLaMA-2 models. Once you have created your account, you can apply for access to one of the models [here](https://huggingface.co/meta-llama/Llama-2-7b-hf). Once you apply for access to one model and accept Meta's license, you will have access to [all LLaMA-2 models](https://huggingface.co/meta-llama/) in Hugging Face."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. Once you have access, you will need to install Hugging Face's CLI interface in order to download the model and authenticate your account with Hugging Face."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: huggingface_hub in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (0.17.3)\n",
      "Requirement already satisfied: filelock in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from huggingface_hub) (3.9.0)\n",
      "Requirement already satisfied: fsspec in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from huggingface_hub) (2023.4.0)\n",
      "Requirement already satisfied: requests in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from huggingface_hub) (2.31.0)\n",
      "Requirement already satisfied: tqdm>=4.42.1 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from huggingface_hub) (4.66.1)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from huggingface_hub) (6.0.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from huggingface_hub) (4.7.1)\n",
      "Requirement already satisfied: packaging>=20.9 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from huggingface_hub) (23.1)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from requests->huggingface_hub) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from requests->huggingface_hub) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from requests->huggingface_hub) (2.0.7)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages (from requests->huggingface_hub) (2023.7.22)\n"
     ]
    }
   ],
   "source": [
    "# Install CLI interface\n",
    "!{sys.executable} -m pip install huggingface_hub"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before you authenticate, make sure you set up a [user access token](https://huggingface.co/docs/hub/security-tokens) in your Hugging Face account. You can then run the following command and enter your token to authenticate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a76d6237a55b48819aee6be79a7d0a00",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the above does not work, you can enter the following command in your terminal (outside of this notebook).\n",
    "\n",
    "```console\n",
    "$ huggingface-cli login\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6. Verify that you have the following versions.\n",
    "- Transformers: v4.33.2 or higher\n",
    "  - Without this version, ONNX Runtime optimizations may not apply and you will miss performance benefits.\n",
    "- Protobuf: v3.20.2\n",
    "  - Without this version, you will not be able to optimize the exported ONNX model and will likely see a `Segmentation fault`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To see what versions you have installed, you can run the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: transformers\n",
      "Version: 4.35.0\n",
      "Summary: State-of-the-art Machine Learning for JAX, PyTorch and TensorFlow\n",
      "Home-page: https://github.com/huggingface/transformers\n",
      "Author: The Hugging Face team (past and future) with the help of all our contributors (https://github.com/huggingface/transformers/graphs/contributors)\n",
      "Author-email: transformers@huggingface.co\n",
      "License: Apache 2.0 License\n",
      "Location: /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages\n",
      "Requires: filelock, huggingface-hub, numpy, packaging, pyyaml, regex, requests, safetensors, tokenizers, tqdm\n",
      "Required-by: optimum\n"
     ]
    }
   ],
   "source": [
    "!{sys.executable} -m pip show transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: protobuf\n",
      "Version: 3.20.2\n",
      "Summary: Protocol Buffers\n",
      "Home-page: https://developers.google.com/protocol-buffers/\n",
      "Author: \n",
      "Author-email: \n",
      "License: BSD-3-Clause\n",
      "Location: /data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages\n",
      "Requires: \n",
      "Required-by: onnx, onnxruntime-gpu\n"
     ]
    }
   ],
   "source": [
    "!{sys.executable} -m pip show protobuf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If there is a version mismatch, please uninstall the existing version and install the following versions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If `transformers` version is wrong\n",
    "!{sys.executable} -m pip uninstall -y transformers\n",
    "!{sys.executable} -m pip install transformers==4.33.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If `protobuf` version is wrong\n",
    "!{sys.executable} -m pip uninstall -y protobuf\n",
    "!{sys.executable} -m pip install protobuf==3.20.2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Export + Optimize + Quantize LLaMA-2\n",
    "\n",
    "Now that all prerequisites have been completed, you are ready to export LLaMA-2 to an optimized (and quantized, if requested) ONNX model. Before you begin, let's define a cache directory to store downloaded files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "cache_dir = \"./cache_dir\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run all of these steps in one command, there is a `convert_to_onnx` script in ONNX Runtime for LLaMA-2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "usage: convert_to_onnx.py [-h] -m MODEL_NAME [-i INPUT] [-o OUTPUT]\n",
      "                          [-p {fp32,fp16,int8,int4}] [-e {cpu,cuda,rocm}] [-r]\n",
      "                          [--use_gqa] [--no_merged]\n",
      "                          [-q {blockwise,smooth_quant,quantize_dynamic}]\n",
      "                          [--block_size BLOCK_SIZE]\n",
      "                          [--smooth_quant_alpha SMOOTH_QUANT_ALPHA]\n",
      "                          [--smooth_quant_dataset SMOOTH_QUANT_DATASET]\n",
      "                          [--pad_max PAD_MAX]\n",
      "                          [--calibration_sampling_size CALIBRATION_SAMPLING_SIZE]\n",
      "                          [--nc_workspace NC_WORKSPACE]\n",
      "                          [--quantize_embedding_layer]\n",
      "                          [--quantize_per_channel] [--quantize_reduce_range]\n",
      "                          [-v] [-d] [--cache_dir CACHE_DIR]\n",
      "\n",
      "optional arguments:\n",
      "  -h, --help            show this help message and exit\n",
      "  -m MODEL_NAME, --model_name MODEL_NAME\n",
      "                        Model name in Hugging Face\n",
      "  -i INPUT, --input INPUT\n",
      "                        Directory path to PyTorch model and associated files\n",
      "                        if saved on disk\n",
      "  -o OUTPUT, --output OUTPUT\n",
      "                        Directory path to save exported model files in\n",
      "  -p {fp32,fp16,int8,int4}, --precision {fp32,fp16,int8,int4}\n",
      "                        Precision to export model in\n",
      "  -e {cpu,cuda,rocm}, --execution_provider {cpu,cuda,rocm}\n",
      "                        Execution provider to verify parity with\n",
      "  -r, --reexport        Re-export models and overwrite existing models in\n",
      "                        output folder\n",
      "  --use_gqa             Use GroupQueryAttention instead of MultiHeadAttention\n",
      "  --no_merged           Export models into 2 ONNX files instead of 1.\n",
      "                        Deprecated in favor of exporting into 1 ONNX file.\n",
      "  -q {blockwise,smooth_quant,quantize_dynamic}, --quantization_method {blockwise,smooth_quant,quantize_dynamic}\n",
      "                        Run a specific quantization algorithm (blockwise for\n",
      "                        int4, smooth_quant for int8, quantize_dynamic for\n",
      "                        int8). Blockwise is recommended. Need to install extra\n",
      "                        packages in `requirements-quant.txt` for SmoothQuant.\n",
      "  -v, --verbose         Print verbose logs\n",
      "  -d, --use_dynamo_export\n",
      "                        Use the new Dynamo exporter instead of the old\n",
      "                        TorchScript exporter\n",
      "  --cache_dir CACHE_DIR\n",
      "                        model cache dir to override default HF cache dir to\n",
      "                        avoid overflood the /home dir\n",
      "\n",
      "4-bit quantization:\n",
      "  --block_size BLOCK_SIZE\n",
      "                        Block size to quantize with. See https://github.com/mi\n",
      "                        crosoft/onnxruntime/blob/main/onnxruntime/python/tools\n",
      "                        /quantization/matmul_4bits_quantizer.py for details.\n",
      "\n",
      "smooth_quant (8-bit quantization):\n",
      "  --smooth_quant_alpha SMOOTH_QUANT_ALPHA\n",
      "                        Strength to control migration difficulty from\n",
      "                        activation to weights. Default is 0.8 to match value\n",
      "                        used in original paper for LLaMA. Paper recommends\n",
      "                        using values in [0.4, 0.6] range. Link to paper:\n",
      "                        https://arxiv.org/pdf/2211.10438.pdf\n",
      "  --smooth_quant_dataset SMOOTH_QUANT_DATASET\n",
      "                        Path to dataset for calibration during quantization\n",
      "  --pad_max PAD_MAX     Max padding size\n",
      "  --calibration_sampling_size CALIBRATION_SAMPLING_SIZE\n",
      "                        Calibration sampling size for quantization config\n",
      "  --nc_workspace NC_WORKSPACE\n",
      "                        Workspace to save intermediate files generated by\n",
      "                        Intel's Neural Compressor package.\n",
      "\n",
      "quantize_dynamic (8-bit quantization):\n",
      "  --quantize_embedding_layer\n",
      "                        Quantize MatMul, GEMM, and Gather.\n",
      "  --quantize_per_channel\n",
      "                        Quantize weights per each channel.\n",
      "  --quantize_reduce_range\n",
      "                        Quantize weights with 7 bits.\n"
     ]
    }
   ],
   "source": [
    "# List all flag options with convert_to_onnx\n",
    "import sys\n",
    "!{sys.executable} -m onnxruntime.transformers.models.llama.convert_to_onnx --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is what each of the main flags do.\n",
    "- `-m/--model_name`: corresponds to the model name in Hugging Face\n",
    "- `--output`: folder to store the exported ONNX model in\n",
    "- `--precision`: precision you want the final exported ONNX model to be in\n",
    "- `--execution_provider`: the execution provider to run the model with\n",
    "- `--quantization_method`: the method by which to quantize the model. For INT4, the quantization method is called `blockwise`.\n",
    "- `--use_gqa`: replace MultiHeadAttention with GroupQueryAttention. This replacement can only happen for FP16 CUDA and INT4 CUDA. This flag must also be used if `num_attention_heads != num_key_value_heads` in your model. You can determine this by running the below code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages/transformers/configuration_utils.py:486: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "90910f3c20384aeab1c37e31fd3a0189",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)lve/main/config.json:   0%|          | 0.00/609 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import LlamaConfig\n",
    "\n",
    "model_name = \"meta-llama/Llama-2-7b-hf\"  # Replace with your model name\n",
    "config = LlamaConfig.from_pretrained(model_name, use_auth_token=True, cache_dir=cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Does meta-llama/Llama-2-7b-hf require GroupQueryAttention? False'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f\"Does {model_name} require GroupQueryAttention? {config.num_attention_heads != config.num_key_value_heads}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are unsure if your machine satisfies the memory requirements, you should run `convert_to_onnx` outside of this notebook so that the notebook does not crash. Later in this notebook, you can load the model back.\n",
    "\n",
    "There are a [wide range of scenarios](https://github.com/microsoft/onnxruntime/tree/main/onnxruntime/python/tools/transformers/models/llama) for which you can export LLaMA-2. Here are some of the common options and the export commands for them.\n",
    "\n",
    "FP16 CUDA (with GroupQueryAttention)\n",
    "\n",
    "```console\n",
    "// Model will be stored at ./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp16.onnx\n",
    "$ python -m onnxruntime.transformers.models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-fp16-gqa --precision fp16 --execution_provider cuda --use_gqa\n",
    "```\n",
    "\n",
    "INT4 CPU (with FP32 inputs/outputs)\n",
    "```console\n",
    "// Model will be stored at ./llama2-7b-int4-cpu/rank_0_Llama-2-7b-hf_decoder_merged_model_int4.onnx\n",
    "$ python -m onnxruntime.transformers.models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-int4-cpu --precision int4 --execution_provider cpu --quantization_method blockwise\n",
    "```\n",
    "\n",
    "INT4 CUDA (with FP16 inputs/outputs)\n",
    "```console\n",
    "// Model will be stored at ./llama2-7b-int4-gpu/rank_0_Llama-2-7b-hf_decoder_merged_model_int4.onnx\n",
    "$ python -m onnxruntime.transformers.models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-int4-gpu --precision int4 --execution_provider cuda --quantization_method blockwise --use_gqa\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch Version:2.2.0.dev20231111+cu118\n",
      "Transformers Version:4.35.0\n",
      "OnnxRuntime Version:1.16.2\n",
      "Arguments: Namespace(model_name='meta-llama/Llama-2-7b-hf', input='.', output='./llama2-7b-fp16-gqa', precision=<Precision.FLOAT16: 'fp16'>, execution_provider='cuda', reexport=False, use_gqa=True, no_merged=False, quantization_method='', block_size=32, smooth_quant_alpha=0.8, smooth_quant_dataset='NeelNanda/pile-10k', pad_max=196, calibration_sampling_size=8, nc_workspace='./nc_workspace', quantize_embedding_layer=False, quantize_per_channel=False, quantize_reduce_range=False, verbose=False, use_dynamo_export=False, cache_dir='./cache_dir')\n",
      "world_size: 1\n",
      "/data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages/transformers/models/auto/configuration_auto.py:1033: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.\n",
      "  warnings.warn(\n",
      "/data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages/transformers/models/auto/auto_factory.py:472: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.\n",
      "  warnings.warn(\n",
      "Downloading (…)fetensors.index.json: 100%|█| 26.8k/26.8k [00:00<00:00, 14.3MB/s]\n",
      "Downloading shards:   0%|                                 | 0/2 [00:00<?, ?it/s]\n",
      "Downloading (…)of-00002.safetensors:   0%|          | 0.00/9.98G [00:00<?, ?B/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   0%|  | 31.5M/9.98G [00:00<00:37, 266MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   1%|  | 73.4M/9.98G [00:00<00:30, 329MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   1%|   | 115M/9.98G [00:00<00:28, 347MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   2%|   | 157M/9.98G [00:00<00:27, 362MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   2%|   | 199M/9.98G [00:00<00:26, 373MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   2%|   | 241M/9.98G [00:00<00:25, 377MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   3%|   | 283M/9.98G [00:00<00:25, 381MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   3%|   | 325M/9.98G [00:00<00:25, 380MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   4%|   | 367M/9.98G [00:00<00:25, 382MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   4%|   | 409M/9.98G [00:01<00:24, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   5%|▏  | 451M/9.98G [00:01<00:24, 386MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   5%|▏  | 493M/9.98G [00:01<00:24, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   5%|▏  | 535M/9.98G [00:01<00:24, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   6%|▏  | 577M/9.98G [00:01<00:24, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   6%|▏  | 619M/9.98G [00:01<00:24, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   7%|▏  | 661M/9.98G [00:01<00:24, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   7%|▏  | 703M/9.98G [00:01<00:23, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   7%|▏  | 744M/9.98G [00:01<00:23, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   8%|▏  | 786M/9.98G [00:02<00:23, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   8%|▏  | 828M/9.98G [00:02<00:23, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   9%|▎  | 870M/9.98G [00:02<00:23, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   9%|▎  | 912M/9.98G [00:02<00:23, 384MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  10%|▎  | 954M/9.98G [00:02<00:23, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  10%|▎  | 996M/9.98G [00:02<00:23, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  10%|▏ | 1.04G/9.98G [00:02<00:23, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  11%|▏ | 1.08G/9.98G [00:02<00:22, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  11%|▏ | 1.12G/9.98G [00:02<00:22, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  12%|▏ | 1.16G/9.98G [00:03<00:22, 392MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  12%|▏ | 1.21G/9.98G [00:03<00:22, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  13%|▎ | 1.25G/9.98G [00:03<00:22, 381MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  13%|▎ | 1.29G/9.98G [00:03<00:22, 384MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  13%|▎ | 1.33G/9.98G [00:03<00:22, 386MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  14%|▎ | 1.37G/9.98G [00:03<00:22, 386MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  14%|▎ | 1.42G/9.98G [00:03<00:22, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  15%|▎ | 1.46G/9.98G [00:03<00:21, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  15%|▎ | 1.50G/9.98G [00:03<00:21, 389MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  15%|▎ | 1.54G/9.98G [00:04<00:21, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  16%|▎ | 1.58G/9.98G [00:04<00:21, 384MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  16%|▎ | 1.63G/9.98G [00:04<00:21, 383MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  17%|▎ | 1.67G/9.98G [00:04<00:21, 389MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  17%|▎ | 1.71G/9.98G [00:04<00:21, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  18%|▎ | 1.75G/9.98G [00:04<00:21, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  18%|▎ | 1.79G/9.98G [00:04<00:20, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  18%|▎ | 1.84G/9.98G [00:04<00:20, 389MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  19%|▍ | 1.88G/9.98G [00:04<00:20, 386MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  19%|▍ | 1.92G/9.98G [00:04<00:20, 386MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  20%|▍ | 1.96G/9.98G [00:05<00:20, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  20%|▍ | 2.00G/9.98G [00:05<00:20, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  20%|▍ | 2.04G/9.98G [00:05<00:20, 386MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  21%|▍ | 2.09G/9.98G [00:05<00:20, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  21%|▍ | 2.13G/9.98G [00:05<00:20, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  22%|▍ | 2.17G/9.98G [00:05<00:19, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  22%|▍ | 2.21G/9.98G [00:05<00:20, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  23%|▍ | 2.25G/9.98G [00:05<00:19, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  23%|▍ | 2.30G/9.98G [00:05<00:19, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  23%|▍ | 2.34G/9.98G [00:06<00:19, 393MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  24%|▍ | 2.38G/9.98G [00:06<00:19, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  24%|▍ | 2.42G/9.98G [00:06<00:19, 392MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  25%|▍ | 2.46G/9.98G [00:06<00:19, 392MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  25%|▌ | 2.51G/9.98G [00:06<00:19, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  26%|▌ | 2.55G/9.98G [00:06<00:18, 393MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  26%|▌ | 2.59G/9.98G [00:06<00:18, 392MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  26%|▌ | 2.63G/9.98G [00:06<00:18, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  27%|▌ | 2.67G/9.98G [00:06<00:18, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  27%|▌ | 2.72G/9.98G [00:07<00:18, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  28%|▌ | 2.76G/9.98G [00:07<00:18, 389MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  28%|▌ | 2.80G/9.98G [00:07<00:18, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  28%|▌ | 2.84G/9.98G [00:07<00:18, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  29%|▌ | 2.88G/9.98G [00:07<00:18, 389MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  29%|▌ | 2.93G/9.98G [00:07<00:18, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  30%|▌ | 2.97G/9.98G [00:07<00:18, 389MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  30%|▌ | 3.01G/9.98G [00:07<00:17, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  31%|▌ | 3.05G/9.98G [00:07<00:17, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  31%|▌ | 3.09G/9.98G [00:08<00:17, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  31%|▋ | 3.14G/9.98G [00:08<00:17, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  32%|▋ | 3.18G/9.98G [00:08<00:17, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  32%|▋ | 3.22G/9.98G [00:08<00:17, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  33%|▋ | 3.26G/9.98G [00:08<00:17, 381MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  33%|▋ | 3.30G/9.98G [00:08<00:17, 384MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  34%|▋ | 3.34G/9.98G [00:08<00:17, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  34%|▋ | 3.39G/9.98G [00:08<00:16, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  34%|▋ | 3.43G/9.98G [00:08<00:16, 392MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  35%|▋ | 3.47G/9.98G [00:08<00:16, 393MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  35%|▋ | 3.51G/9.98G [00:09<00:16, 392MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  36%|▋ | 3.55G/9.98G [00:09<00:16, 394MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  36%|▋ | 3.60G/9.98G [00:09<00:16, 394MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  36%|▋ | 3.64G/9.98G [00:09<00:16, 393MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  37%|▋ | 3.68G/9.98G [00:09<00:16, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  37%|▋ | 3.72G/9.98G [00:09<00:16, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  38%|▊ | 3.76G/9.98G [00:09<00:16, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  38%|▊ | 3.81G/9.98G [00:09<00:15, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  39%|▊ | 3.85G/9.98G [00:09<00:15, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  39%|▊ | 3.89G/9.98G [00:10<00:15, 383MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  39%|▊ | 3.93G/9.98G [00:10<00:15, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  40%|▊ | 3.97G/9.98G [00:10<00:15, 389MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  40%|▊ | 4.02G/9.98G [00:10<00:15, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  41%|▊ | 4.06G/9.98G [00:10<00:15, 379MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  41%|▊ | 4.10G/9.98G [00:10<00:17, 339MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  42%|▊ | 4.14G/9.98G [00:10<00:16, 352MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  42%|▊ | 4.18G/9.98G [00:10<00:15, 366MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  42%|▊ | 4.23G/9.98G [00:10<00:15, 372MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  43%|▊ | 4.27G/9.98G [00:11<00:15, 373MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  43%|▊ | 4.31G/9.98G [00:11<00:15, 377MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  44%|▊ | 4.35G/9.98G [00:11<00:14, 380MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  44%|▉ | 4.39G/9.98G [00:11<00:14, 373MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  44%|▉ | 4.44G/9.98G [00:11<00:14, 378MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  45%|▉ | 4.48G/9.98G [00:11<00:14, 379MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  45%|▉ | 4.52G/9.98G [00:11<00:14, 380MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  46%|▉ | 4.56G/9.98G [00:11<00:14, 383MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  46%|▉ | 4.60G/9.98G [00:11<00:13, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  47%|▉ | 4.65G/9.98G [00:12<00:13, 384MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  47%|▉ | 4.69G/9.98G [00:12<00:13, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  47%|▉ | 4.73G/9.98G [00:12<00:13, 386MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  48%|▉ | 4.77G/9.98G [00:12<00:13, 384MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  48%|▉ | 4.81G/9.98G [00:12<00:13, 386MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  49%|▉ | 4.85G/9.98G [00:12<00:13, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  49%|▉ | 4.90G/9.98G [00:12<00:13, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  50%|▉ | 4.94G/9.98G [00:12<00:12, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  50%|▉ | 4.98G/9.98G [00:12<00:12, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  50%|█ | 5.02G/9.98G [00:13<00:12, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  51%|█ | 5.06G/9.98G [00:13<00:12, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  51%|█ | 5.11G/9.98G [00:13<00:12, 389MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  52%|█ | 5.15G/9.98G [00:13<00:12, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  52%|█ | 5.19G/9.98G [00:13<00:12, 389MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  52%|█ | 5.23G/9.98G [00:13<00:12, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  53%|█ | 5.27G/9.98G [00:13<00:12, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  53%|█ | 5.32G/9.98G [00:13<00:12, 388MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  54%|█ | 5.36G/9.98G [00:13<00:11, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  54%|█ | 5.40G/9.98G [00:14<00:11, 386MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  55%|█ | 5.44G/9.98G [00:14<00:11, 386MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  55%|█ | 5.48G/9.98G [00:14<00:11, 384MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  55%|█ | 5.53G/9.98G [00:14<00:11, 382MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  56%|█ | 5.57G/9.98G [00:14<00:11, 384MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  56%|█ | 5.61G/9.98G [00:14<00:11, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  57%|█▏| 5.65G/9.98G [00:14<00:11, 384MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  57%|█▏| 5.69G/9.98G [00:14<00:11, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  57%|█▏| 5.74G/9.98G [00:14<00:10, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  58%|█▏| 5.78G/9.98G [00:15<00:10, 389MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  58%|█▏| 5.82G/9.98G [00:15<00:10, 393MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  59%|█▏| 5.86G/9.98G [00:15<00:10, 396MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  59%|█▏| 5.90G/9.98G [00:15<00:10, 394MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  60%|█▏| 5.95G/9.98G [00:15<00:10, 396MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  60%|█▏| 5.99G/9.98G [00:15<00:10, 392MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  60%|█▏| 6.03G/9.98G [00:15<00:10, 392MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  61%|█▏| 6.07G/9.98G [00:15<00:09, 395MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  61%|█▏| 6.11G/9.98G [00:15<00:09, 397MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  62%|█▏| 6.16G/9.98G [00:15<00:09, 397MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  62%|█▏| 6.20G/9.98G [00:16<00:09, 397MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  63%|█▎| 6.24G/9.98G [00:16<00:09, 399MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  63%|█▎| 6.28G/9.98G [00:16<00:09, 398MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  63%|█▎| 6.32G/9.98G [00:16<00:09, 401MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  64%|█▎| 6.36G/9.98G [00:16<00:09, 399MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  64%|█▎| 6.41G/9.98G [00:16<00:09, 394MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  65%|█▎| 6.45G/9.98G [00:16<00:14, 246MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  65%|█▎| 6.49G/9.98G [00:17<00:12, 276MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  65%|█▎| 6.53G/9.98G [00:17<00:11, 304MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  66%|█▎| 6.57G/9.98G [00:17<00:10, 329MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  66%|█▎| 6.62G/9.98G [00:17<00:09, 346MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  67%|█▎| 6.66G/9.98G [00:17<00:09, 360MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  67%|█▎| 6.70G/9.98G [00:17<00:08, 372MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  68%|█▎| 6.74G/9.98G [00:17<00:08, 382MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  68%|█▎| 6.78G/9.98G [00:17<00:08, 385MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  68%|█▎| 6.83G/9.98G [00:17<00:08, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  69%|█▍| 6.87G/9.98G [00:17<00:07, 394MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  69%|█▍| 6.91G/9.98G [00:18<00:07, 392MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  70%|█▍| 6.95G/9.98G [00:18<00:07, 394MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  70%|█▍| 6.99G/9.98G [00:18<00:07, 396MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  71%|█▍| 7.04G/9.98G [00:18<00:07, 394MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  71%|█▍| 7.08G/9.98G [00:18<00:07, 394MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  71%|█▍| 7.12G/9.98G [00:18<00:07, 395MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  72%|█▍| 7.16G/9.98G [00:18<00:07, 393MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  72%|█▍| 7.20G/9.98G [00:18<00:07, 394MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  73%|█▍| 7.25G/9.98G [00:18<00:06, 394MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  73%|█▍| 7.29G/9.98G [00:19<00:06, 393MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  73%|█▍| 7.33G/9.98G [00:19<00:06, 394MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  74%|█▍| 7.37G/9.98G [00:19<00:06, 395MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  74%|█▍| 7.41G/9.98G [00:19<00:06, 393MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  75%|█▍| 7.46G/9.98G [00:19<00:06, 395MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  75%|█▌| 7.50G/9.98G [00:19<00:06, 397MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  76%|█▌| 7.54G/9.98G [00:19<00:06, 394MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  76%|█▌| 7.58G/9.98G [00:19<00:06, 396MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  76%|█▌| 7.62G/9.98G [00:19<00:05, 397MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  77%|█▌| 7.67G/9.98G [00:19<00:05, 393MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  77%|█▌| 7.71G/9.98G [00:20<00:05, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  78%|█▌| 7.75G/9.98G [00:20<00:05, 390MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  78%|█▌| 7.79G/9.98G [00:20<00:05, 391MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  79%|█▌| 7.83G/9.98G [00:20<00:05, 393MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  79%|█▌| 7.89G/9.98G [00:20<00:05, 405MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  80%|█▌| 7.94G/9.98G [00:20<00:04, 420MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  80%|█▌| 7.99G/9.98G [00:20<00:04, 431MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  81%|█▌| 8.04G/9.98G [00:20<00:04, 442MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  81%|█▌| 8.10G/9.98G [00:20<00:04, 445MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  82%|█▋| 8.15G/9.98G [00:21<00:04, 447MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  82%|█▋| 8.20G/9.98G [00:21<00:03, 453MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  83%|█▋| 8.25G/9.98G [00:21<00:03, 455MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  83%|█▋| 8.30G/9.98G [00:21<00:03, 453MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  84%|█▋| 8.36G/9.98G [00:21<00:03, 452MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  84%|█▋| 8.41G/9.98G [00:21<00:03, 449MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  85%|█▋| 8.46G/9.98G [00:21<00:03, 453MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  85%|█▋| 8.51G/9.98G [00:21<00:03, 458MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  86%|█▋| 8.57G/9.98G [00:22<00:03, 459MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  86%|█▋| 8.62G/9.98G [00:22<00:02, 462MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  87%|█▋| 8.67G/9.98G [00:22<00:02, 458MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  87%|█▋| 8.72G/9.98G [00:22<00:02, 458MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  88%|█▊| 8.78G/9.98G [00:22<00:02, 461MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  88%|█▊| 8.83G/9.98G [00:22<00:02, 459MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  89%|█▊| 8.88G/9.98G [00:22<00:02, 462MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  90%|█▊| 8.93G/9.98G [00:22<00:02, 464MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  90%|█▊| 8.99G/9.98G [00:22<00:02, 464MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  91%|█▊| 9.04G/9.98G [00:23<00:02, 467MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  91%|█▊| 9.09G/9.98G [00:23<00:01, 463MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  92%|█▊| 9.14G/9.98G [00:23<00:01, 458MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  92%|█▊| 9.20G/9.98G [00:23<00:01, 456MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  93%|█▊| 9.25G/9.98G [00:23<00:01, 458MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  93%|█▊| 9.30G/9.98G [00:23<00:01, 460MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  94%|█▉| 9.35G/9.98G [00:23<00:01, 464MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  94%|█▉| 9.41G/9.98G [00:23<00:01, 464MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  95%|█▉| 9.46G/9.98G [00:23<00:01, 465MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  95%|█▉| 9.51G/9.98G [00:24<00:01, 463MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  96%|█▉| 9.56G/9.98G [00:24<00:00, 466MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  96%|█▉| 9.62G/9.98G [00:24<00:00, 462MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  97%|█▉| 9.67G/9.98G [00:24<00:00, 460MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  97%|█▉| 9.72G/9.98G [00:24<00:00, 455MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  98%|█▉| 9.77G/9.98G [00:24<00:00, 392MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  98%|█▉| 9.83G/9.98G [00:24<00:00, 411MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  99%|█▉| 9.88G/9.98G [00:24<00:00, 423MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors: 100%|█▉| 9.93G/9.98G [00:25<00:00, 430MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors: 100%|██| 9.98G/9.98G [00:25<00:00, 397MB/s]\u001b[A\n",
      "Downloading shards:  50%|████████████▌            | 1/2 [00:25<00:25, 25.31s/it]\n",
      "Downloading (…)of-00002.safetensors:   0%|          | 0.00/3.50G [00:00<?, ?B/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   1%|  | 52.4M/3.50G [00:00<00:08, 419MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   3%|  | 94.4M/3.50G [00:00<00:08, 415MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   4%|   | 136M/3.50G [00:00<00:08, 410MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   5%|▏  | 178M/3.50G [00:00<00:08, 413MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   7%|▏  | 231M/3.50G [00:00<00:07, 416MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   8%|▏  | 273M/3.50G [00:00<00:07, 414MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:   9%|▎  | 315M/3.50G [00:00<00:07, 415MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  10%|▎  | 367M/3.50G [00:00<00:07, 417MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  12%|▎  | 409M/3.50G [00:00<00:07, 417MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  13%|▍  | 461M/3.50G [00:01<00:07, 417MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  14%|▍  | 503M/3.50G [00:01<00:07, 413MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  16%|▍  | 545M/3.50G [00:01<00:07, 413MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  17%|▌  | 587M/3.50G [00:01<00:07, 415MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  18%|▌  | 629M/3.50G [00:01<00:06, 415MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  19%|▌  | 671M/3.50G [00:01<00:06, 415MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  20%|▌  | 713M/3.50G [00:01<00:06, 416MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  22%|▋  | 755M/3.50G [00:01<00:06, 416MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  23%|▋  | 797M/3.50G [00:01<00:06, 416MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  24%|▋  | 839M/3.50G [00:02<00:06, 415MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  25%|▊  | 881M/3.50G [00:02<00:10, 257MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  26%|▊  | 923M/3.50G [00:02<00:08, 290MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  28%|▊  | 965M/3.50G [00:02<00:07, 319MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  29%|▌ | 1.01G/3.50G [00:02<00:07, 343MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  30%|▌ | 1.05G/3.50G [00:02<00:06, 362MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  31%|▌ | 1.09G/3.50G [00:02<00:06, 377MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  32%|▋ | 1.13G/3.50G [00:02<00:06, 387MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  34%|▋ | 1.17G/3.50G [00:03<00:05, 395MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  35%|▋ | 1.22G/3.50G [00:03<00:05, 401MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  36%|▋ | 1.26G/3.50G [00:03<00:05, 405MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  37%|▋ | 1.30G/3.50G [00:03<00:05, 407MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  39%|▊ | 1.35G/3.50G [00:03<00:05, 413MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  40%|▊ | 1.41G/3.50G [00:03<00:04, 422MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  42%|▊ | 1.46G/3.50G [00:03<00:04, 435MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  43%|▊ | 1.51G/3.50G [00:03<00:04, 446MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  45%|▉ | 1.56G/3.50G [00:03<00:04, 453MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  46%|▉ | 1.61G/3.50G [00:04<00:04, 453MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  48%|▉ | 1.67G/3.50G [00:04<00:04, 456MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  49%|▉ | 1.72G/3.50G [00:04<00:03, 459MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  51%|█ | 1.77G/3.50G [00:04<00:03, 463MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  52%|█ | 1.82G/3.50G [00:04<00:03, 460MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  54%|█ | 1.88G/3.50G [00:04<00:03, 460MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  55%|█ | 1.93G/3.50G [00:04<00:03, 460MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  57%|█▏| 1.98G/3.50G [00:04<00:03, 460MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  58%|█▏| 2.03G/3.50G [00:04<00:03, 458MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  60%|█▏| 2.09G/3.50G [00:05<00:03, 454MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  61%|█▏| 2.14G/3.50G [00:05<00:02, 454MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  63%|█▎| 2.19G/3.50G [00:05<00:02, 455MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  64%|█▎| 2.24G/3.50G [00:05<00:02, 455MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  66%|█▎| 2.30G/3.50G [00:05<00:02, 458MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  67%|█▎| 2.35G/3.50G [00:05<00:02, 456MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  69%|█▎| 2.40G/3.50G [00:05<00:02, 458MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  70%|█▍| 2.45G/3.50G [00:05<00:02, 458MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  72%|█▍| 2.51G/3.50G [00:05<00:02, 461MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  73%|█▍| 2.56G/3.50G [00:06<00:02, 460MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  75%|█▍| 2.61G/3.50G [00:06<00:01, 462MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  76%|█▌| 2.66G/3.50G [00:06<00:01, 462MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  78%|█▌| 2.72G/3.50G [00:06<00:01, 460MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  79%|█▌| 2.77G/3.50G [00:06<00:01, 463MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  81%|█▌| 2.82G/3.50G [00:06<00:01, 462MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  82%|█▋| 2.87G/3.50G [00:06<00:01, 461MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  84%|█▋| 2.93G/3.50G [00:06<00:01, 462MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  85%|█▋| 2.98G/3.50G [00:06<00:01, 464MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  87%|█▋| 3.03G/3.50G [00:07<00:01, 464MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  88%|█▊| 3.08G/3.50G [00:07<00:00, 461MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  90%|█▊| 3.14G/3.50G [00:07<00:00, 458MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  91%|█▊| 3.19G/3.50G [00:07<00:00, 456MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  93%|█▊| 3.24G/3.50G [00:07<00:00, 458MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  94%|█▉| 3.29G/3.50G [00:07<00:00, 460MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  96%|█▉| 3.34G/3.50G [00:07<00:00, 456MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  97%|█▉| 3.40G/3.50G [00:07<00:00, 457MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors:  99%|█▉| 3.45G/3.50G [00:08<00:00, 455MB/s]\u001b[A\n",
      "Downloading (…)of-00002.safetensors: 100%|██| 3.50G/3.50G [00:08<00:00, 430MB/s]\u001b[A\n",
      "Downloading shards: 100%|█████████████████████████| 2/2 [00:33<00:00, 16.75s/it]\n",
      "Loading checkpoint shards: 100%|██████████████████| 2/2 [00:01<00:00,  1.10it/s]\n",
      "/data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages/transformers/utils/hub.py:374: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.\n",
      "  warnings.warn(\n",
      "Downloading (…)neration_config.json: 100%|██████| 188/188 [00:00<00:00, 130kB/s]\n",
      "/data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages/transformers/modeling_attn_mask_utils.py:94: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal:\n",
      "/data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages/transformers/modeling_attn_mask_utils.py:137: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if past_key_values_length > 0:\n",
      "/data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages/transformers/models/llama/modeling_llama.py:140: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if seq_len > self.max_seq_len_cached:\n",
      "/data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages/transformers/models/llama/modeling_llama.py:392: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):\n",
      "/data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages/transformers/models/llama/modeling_llama.py:399: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):\n",
      "/data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages/transformers/models/llama/modeling_llama.py:409: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
      "The Llama-2-7b-hf merged ONNX model has been successfully created with the TorchScript exporter!\n",
      "Optimizing models...\n",
      "Removed 392 Cast nodes with output type same as input\n",
      "Fused SimplifiedLayerNormalization: 65\n",
      "Fused SkipSimplifiedLayerNormalization: 64\n",
      "Fused RotaryEmbedding: 64\n",
      "Removed 1023 nodes\n",
      "Fused MultiHeadAttention: 32\n",
      "Removed 1088 nodes\n",
      "Remove reshape node /model/Reshape_1 since its input shape is same as output: [4]\n",
      "Remove reshape node /model/Reshape_2 since its input shape is same as output: [4]\n",
      "Removed 2 nodes\n",
      "postprocess: remove Reshape count: 0\n",
      "opset version: 13\n",
      "Sort graphs in topological order\n",
      "Model saved to ./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32_opt.onnx\n",
      "The ONNX model at ./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32.onnx has been successfully optimized and saved at ./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32_opt.onnx!\n",
      "\u001b[33mRemoved ./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32.onnx and ./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32.onnx.data\u001b[0m\n",
      "The Llama-2-7b-hf ONNX model has been successfully optimized with the ORT transformer optimizer script!\n",
      "Converting to float16...\n",
      "Removed 64 Cast nodes with output type same as input\n",
      "Removed 94 nodes\n",
      "Sort graphs in topological order\n",
      "Model saved to ./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp16.onnx\n",
      "The ONNX model at ./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32_opt.onnx has been converted to float16 and saved at ./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp16.onnx!\n",
      "\u001b[33mRemoved ./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32_opt.onnx and ./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32_opt.onnx.data\u001b[0m\n",
      "The Llama-2-7b-hf ONNX model has been successfully converted to float16!\n",
      "Verifying parity on all ONNX models created\n",
      "Arguments: Namespace(model_name='meta-llama/Llama-2-7b-hf', torch_model_directory='.', onnx_model_path='./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp16.onnx', execution_provider='cuda', verbose=False, use_past_kv=False, use_gqa=True, merged=True, precision='fp16', cache_dir='./cache_dir')\n",
      "world_size: 1\n",
      "Loading checkpoint shards: 100%|██████████████████| 2/2 [00:02<00:00,  1.27s/it]\n",
      "PyTorch took 0.6413478851318359 s\n",
      "2023-11-12 09:12:09.737007636 [W:onnxruntime:, session_state.cc:1162 VerifyEachNodeIsAssignedToAnEp] Some nodes were not assigned to the preferred execution providers which may or may not have an negative impact on performance. e.g. ORT explicitly assigns shape related ops to CPU to improve perf.\n",
      "2023-11-12 09:12:09.737041499 [W:onnxruntime:, session_state.cc:1164 VerifyEachNodeIsAssignedToAnEp] Rerunning with verbose output on a non-minimal build will show node assignments.\n",
      "ONNX Runtime took 1.381934642791748 s\n",
      "\u001b[33mAre PyTorch and ONNX Runtime results close? True\u001b[0m\n",
      "PyTorch took 0.03689312934875488 s\n",
      "2023-11-12 09:12:15.119724309 [W:onnxruntime:, session_state.cc:1162 VerifyEachNodeIsAssignedToAnEp] Some nodes were not assigned to the preferred execution providers which may or may not have an negative impact on performance. e.g. ORT explicitly assigns shape related ops to CPU to improve perf.\n",
      "2023-11-12 09:12:15.119767700 [W:onnxruntime:, session_state.cc:1164 VerifyEachNodeIsAssignedToAnEp] Rerunning with verbose output on a non-minimal build will show node assignments.\n",
      "ONNX Runtime took 0.013516426086425781 s\n",
      "\u001b[33mAre PyTorch and ONNX Runtime results close? True\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# Add/change/remove flags to run below command in your terminal for your scenario\n",
    "!{sys.executable} -m onnxruntime.transformers.models.llama.convert_to_onnx -m \"meta-llama/Llama-2-7b-hf\" --cache_dir \"./cache_dir\" --output \"./llama2-7b-fp16-gqa\" --precision \"fp16\" --execution_provider \"cuda\" --use_gqa"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also pick up an already exported + optimized ONNX model from [Microsoft's repository](https://github.com/microsoft/Llama-2-Onnx/tree/main-CUDA_CPU)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Run LLaMA-2 End-to-End\n",
    "\n",
    "Now that the model is exported to ONNX, you can run it end-to-end. For this example, you can use the LLaMA-2 7B FP16 CUDA model with GQA.\n",
    "\n",
    "Let's first import the necessary libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import LlamaConfig, LlamaTokenizer\n",
    "import numpy as np\n",
    "import onnxruntime as ort\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, you can set the main settings that you want to run your end-to-end scenario with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change the below settings to your desired scenario\n",
    "model_name = \"meta-llama/Llama-2-7b-hf\"  # Model name in Hugging Face\n",
    "onnx_model_path = \"./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp16.onnx\"  # Path to exported ONNX model on disk\n",
    "use_fp16 = True  # True when KV cache inputs/outputs are in float16\n",
    "use_buffer_share = True  # True when --use_gqa was passed during export\n",
    "\n",
    "prompt = [\"ONNX Runtime is \", \"I want to book a vacation to Hawaii. First, I need to \", \"A good workout routine is \", \"How are astronauts launched into space? \"] # List of prompts to use\n",
    "max_length = 64  # max(prompt length + generation length)\n",
    "\n",
    "device_id = 0\n",
    "device = torch.device(f\"cuda:{device_id}\")  # Change to torch.device(\"cpu\") if running on CPU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With your main settings finalized, you can import the model's configuration and tokenizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/kvaishnavi/anaconda3/envs/llama2/lib/python3.9/site-packages/transformers/tokenization_utils_base.py:1895: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1dadcc37df3e45cdba8e91dfd81e8d59",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)okenizer_config.json:   0%|          | 0.00/776 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d025a097ac404074a0953d2a180cd852",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading tokenizer.model:   0%|          | 0.00/500k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "300253b10d264ced85f0a237c3d0a523",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/414 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "72bb3d44395942a4b2c326514c1fc6a9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)/main/tokenizer.json:   0%|          | 0.00/1.84M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "config = LlamaConfig.from_pretrained(model_name, use_auth_token=True, cache_dir=cache_dir)\n",
    "tokenizer = LlamaTokenizer.from_pretrained(model_name, use_auth_token=True, cache_dir=cache_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you need to prepare your inputs for the model to understand. This involves several steps.\n",
    "1. Tokenize the prompts so that the model can understand the inputs.\n",
    "2. Pad the tokenized output so that each batch entry has the same length.\n",
    "3. Pre-allocate on-device memory to store the inputs and outputs.\n",
    "\n",
    "   a. In a typical transformer model, the present KV cache outputs are passed as the past KV cache inputs in the next iteration. But for most models, you need to pre-allocate separate on-device memory for the past KV cache inputs and present KV cache outputs. With GroupQueryAttention, you can share this memory. This allows you to pass the present KV cache outputs directly to the past KV cache inputs.\n",
    "\n",
    "   b. Because the on-device memory can be shared in GroupQueryAttention, you need to pre-allocate enough so that the model can run at any prompt length + generation length. Therefore, you should pre-allocate the on-device KV cache memory to have enough memory to hold the largest sequence length that the model can produce."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_initial_inputs_and_outputs(config, tokenizer, prompt, device, use_fp16, use_buffer_share):\n",
    "    tokenizer.pad_token = \"[PAD]\"  # Set pad token for tokenizer\n",
    "    encodings_dict = tokenizer.batch_encode_plus(prompt, padding=True)\n",
    "    torch_dtype = torch.float16 if use_fp16 else torch.float32\n",
    "\n",
    "    # Move inputs from tokenizer to on-device memory\n",
    "    input_ids = torch.tensor(encodings_dict[\"input_ids\"], device=device, dtype=torch.int64)\n",
    "    attention_mask = torch.tensor(encodings_dict[\"attention_mask\"], device=device, dtype=torch.int64)\n",
    "    position_ids = attention_mask.long().cumsum(-1) - 1\n",
    "    position_ids.masked_fill_(attention_mask == 0, 1)\n",
    "\n",
    "    inputs = {\n",
    "        \"input_ids\": input_ids.contiguous(),\n",
    "        \"attention_mask\": attention_mask.contiguous(),\n",
    "        \"position_ids\": position_ids.contiguous(),\n",
    "    }\n",
    "\n",
    "    # Pre-allocate on-device memory for past_key_values (past KV cache)\n",
    "    # Share on-device memory if use_buffer_share is True\n",
    "    batch_size, sequence_length = input_ids.shape\n",
    "    max_sequence_length = config.max_position_embeddings\n",
    "    num_heads, head_size = config.num_key_value_heads, config.hidden_size // config.num_attention_heads\n",
    "    for i in range(config.num_hidden_layers):\n",
    "        past_key = torch.zeros(batch_size, num_heads, max_sequence_length if use_buffer_share else 0, head_size, device=device, dtype=torch_dtype)\n",
    "        past_value = torch.zeros(batch_size, num_heads, max_sequence_length if use_buffer_share else 0, head_size, device=device, dtype=torch_dtype)\n",
    "        inputs.update({\n",
    "            f\"past_key_values.{i}.key\": past_key.contiguous(),\n",
    "            f\"past_key_values.{i}.value\": past_value.contiguous()\n",
    "        })\n",
    "    \n",
    "    # Pre-allocate on-device memory for logits\n",
    "    logits = torch.zeros(batch_size, sequence_length, config.vocab_size, device=device, dtype=torch_dtype)\n",
    "    outputs = {\n",
    "        \"logits\": logits.contiguous()\n",
    "    }\n",
    "\n",
    "    # Pre-allocate on-device memory for present KV cache if use_buffer_share is False\n",
    "    if not use_buffer_share:\n",
    "        for i in range(config.num_hidden_layers):\n",
    "            present_key = torch.zeros(batch_size, num_heads, sequence_length, head_size, device=device, dtype=torch_dtype)\n",
    "            present_value = torch.zeros(batch_size, num_heads, sequence_length, head_size, device=device, dtype=torch_dtype)\n",
    "            outputs.update({\n",
    "                f\"present.{i}.key\": present_key.contiguous(),\n",
    "                f\"present.{i}.value\": present_value.contiguous()\n",
    "            })\n",
    "\n",
    "    return inputs, outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs, outputs = get_initial_inputs_and_outputs(config, tokenizer, prompt, device, use_fp16, use_buffer_share)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the on-device memory has been allocated, you can load the LLaMA-2 ONNX model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-11-12 09:18:02.339043617 [W:onnxruntime:, session_state.cc:1162 VerifyEachNodeIsAssignedToAnEp] Some nodes were not assigned to the preferred execution providers which may or may not have an negative impact on performance. e.g. ORT explicitly assigns shape related ops to CPU to improve perf.\n",
      "2023-11-12 09:18:02.339077992 [W:onnxruntime:, session_state.cc:1164 VerifyEachNodeIsAssignedToAnEp] Rerunning with verbose output on a non-minimal build will show node assignments.\n"
     ]
    }
   ],
   "source": [
    "sess_options = ort.SessionOptions()\n",
    "ep = (\"CUDAExecutionProvider\", {\"device_id\": device_id})  # change to ep = \"CPUExecutionProvider\" for CPU\n",
    "model = ort.InferenceSession(onnx_model_path, sess_options=sess_options, providers=[ep])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the ONNX model is loaded, you will need a way to bind the inputs and outputs to the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_io_binding(model, inputs, outputs, use_fp16, use_buffer_share):\n",
    "    # Check that all model inputs will be provided\n",
    "    model_inputs = set(map(lambda model_input: model_input.name, model.get_inputs()))\n",
    "    user_inputs = set(inputs.keys())\n",
    "    missing_inputs = model_inputs - user_inputs\n",
    "    if len(missing_inputs):\n",
    "        print(f\"The following model inputs are missing: {missing_inputs}\")\n",
    "        raise Exception(\"There are missing inputs to the model. Please add them and try again.\")\n",
    "\n",
    "    # Remove unnecessary inputs from model inputs\n",
    "    unnecessary_inputs = user_inputs - model_inputs\n",
    "    if len(unnecessary_inputs):\n",
    "        for unnecessary_input in unnecessary_inputs:\n",
    "            print(f\"Removing unnecessary input '{unnecessary_input}' from user provided inputs\")\n",
    "            del inputs[unnecessary_input]\n",
    "\n",
    "    # Bind inputs/outputs to IO binding\n",
    "    io_binding = model.io_binding()\n",
    "    device = None\n",
    "    pt_to_np = {\n",
    "        \"torch.int64\": np.int64,\n",
    "        \"torch.float32\": np.float32,\n",
    "        \"torch.float16\": np.float16\n",
    "    }\n",
    "\n",
    "    for k, v in inputs.items():\n",
    "        io_binding.bind_input(\n",
    "            name=k,\n",
    "            device_type=v.device.type,\n",
    "            device_id=0 if v.device.type == \"cpu\" else v.device.index,\n",
    "            element_type=pt_to_np[repr(v.dtype)],\n",
    "            shape=tuple(v.shape),\n",
    "            buffer_ptr=v.data_ptr()\n",
    "        )\n",
    "        device = v.device\n",
    "\n",
    "    for output in model.get_outputs():\n",
    "        name = output.name\n",
    "        if use_buffer_share and \"present\" in name:\n",
    "            # Bind KV cache outputs to KV cache inputs\n",
    "            v = inputs[name.replace(\"present\", \"past_key_values\")]\n",
    "            io_binding.bind_output(\n",
    "                name=name,\n",
    "                device_type=v.device.type,\n",
    "                device_id=v.device.index,\n",
    "                element_type=np.float16,\n",
    "                shape=tuple(v.shape),\n",
    "                buffer_ptr=v.data_ptr()\n",
    "            )\n",
    "        else:\n",
    "            v = outputs[name]\n",
    "            io_binding.bind_output(\n",
    "                name=name,\n",
    "                device_type=device.type,\n",
    "                device_id=0 if device.type == \"cpu\" else device.index,\n",
    "                element_type=(np.float16 if use_fp16 else np.float32),\n",
    "                shape=tuple(v.shape),\n",
    "                buffer_ptr=v.data_ptr()\n",
    "            )\n",
    "\n",
    "    return io_binding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You are almost ready to run inference with your ONNX model. You need to store the token ids that are generated and keep track of each batch entry to see whether it has completed generation or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_token_ids = inputs[\"input_ids\"].clone()  # store prompt token ids + generated token ids for transcription at the end\n",
    "batch_size, sequence_length = all_token_ids.shape\n",
    "max_sequence_length = config.max_position_embeddings\n",
    "num_heads, head_size = config.num_key_value_heads, config.hidden_size // config.num_attention_heads\n",
    "\n",
    "current_length = sequence_length  # keep track of current length (prompt length + generation length)\n",
    "has_eos = torch.zeros(batch_size, device=device, dtype=torch.bool)  # keep track of each batch entry's status and whether it has reached end-of-sequence (EOS) or not"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you can run inference and generate tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "while current_length <= max_length:\n",
    "    # Run inference\n",
    "    io_binding = apply_io_binding(model, inputs, outputs, use_fp16, use_buffer_share)\n",
    "    io_binding.synchronize_inputs()\n",
    "    model.run_with_iobinding(io_binding)\n",
    "    io_binding.synchronize_outputs()\n",
    "\n",
    "    # Sample/choose next token with argmax (greedy search)\n",
    "    if outputs[\"logits\"].shape[1] > 1:\n",
    "        prompt_end_indices = inputs[\"attention_mask\"].sum(1) - 1\n",
    "        idxs = prompt_end_indices.unsqueeze(dim=1).repeat(1, config.vocab_size).view(batch_size, 1, config.vocab_size)\n",
    "        next_token_logits = torch.gather(outputs[\"logits\"], 1, idxs).squeeze()\n",
    "    else:\n",
    "        next_token_logits = outputs[\"logits\"][:, -1, :]\n",
    "    next_tokens = torch.argmax(next_token_logits, dim=-1)\n",
    "\n",
    "    # Check if we previously reached EOS token id or if generated token id is EOS token id\n",
    "    has_eos = has_eos | next_tokens == tokenizer.eos_token_id\n",
    "\n",
    "    # Determine which new tokens to add to list of all token ids\n",
    "    # Add EOS token ids for batch entries that ended early (ragged batching scenario where some batch entries ended early and some haven't)\n",
    "    tokens_to_add = next_tokens.masked_fill(has_eos, tokenizer.eos_token_id).reshape([batch_size, 1])\n",
    "    all_token_ids = torch.cat([all_token_ids, tokens_to_add], dim=-1)\n",
    "\n",
    "    # Return early if:\n",
    "    # 1) all batch entries have reached EOS token id or \n",
    "    # 2) we have reached the max length of a batch entry (prompt length + generation length) or\n",
    "    # 3) max sequence length that the model can support\n",
    "    current_length += 1\n",
    "    if torch.all(has_eos) or current_length > max_length or current_length > max_sequence_length:\n",
    "        break\n",
    "\n",
    "    # Update inputs for next inference run\n",
    "    inputs[\"input_ids\"] = tokens_to_add\n",
    "    inputs[\"position_ids\"] = torch.max(inputs[\"position_ids\"], dim=1)[0].reshape(batch_size, 1) + 1\n",
    "    inputs[\"attention_mask\"] = torch.cat([inputs[\"attention_mask\"], (~has_eos).to(torch.int64).reshape(batch_size, 1)], 1)\n",
    "\n",
    "    # Set logits to zeros for next inference run and re-use memory buffer\n",
    "    if outputs[\"logits\"].shape[1] != 1:\n",
    "        outputs[\"logits\"] = outputs[\"logits\"][:, :1, :].contiguous()\n",
    "    outputs[\"logits\"].zero_()\n",
    "\n",
    "    # If buffer sharing is off, pass the present KV cache from previous iteration as the past KV cache for next iteration\n",
    "    if not use_buffer_share:\n",
    "        for i in range(config.num_hidden_layers):\n",
    "            inputs[f\"past_key_values.{i}.key\"] = outputs[f\"present.{i}.key\"]\n",
    "            inputs[f\"past_key_values.{i}.value\"] = outputs[f\"present.{i}.value\"]\n",
    "\n",
    "        new_sequence_length = inputs[\"attention_mask\"].shape[1]\n",
    "        for i in range(config.num_hidden_layers):\n",
    "            present_key = torch.zeros(batch_size, num_heads, new_sequence_length, head_size, device=device, dtype=torch_dtype)\n",
    "            present_value = torch.zeros(batch_size, num_heads, new_sequence_length, head_size, device=device, dtype=torch_dtype)\n",
    "            outputs.update({\n",
    "                f\"present.{i}.key\": present_key.contiguous(),\n",
    "                f\"present.{i}.value\": present_value.contiguous()\n",
    "            })"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once generation is complete, you can batch decode all of the token ids to see what the model produced."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['ONNX Runtime is 100% open source and free to use.\\nONNX Runtime is a cross-platform runtime that can be used to run ONNX models on a variety of platforms.\\nONNX Runtime is a cross-',\n",
       " \"I want to book a vacation to Hawaii. First, I need to 1) find a good travel agent, 2) find a good hotel, and 3) find a good flight.\\nI've been to Hawaii before, so I know what I like. I'm looking for\",\n",
       " 'A good workout routine is 30 minutes of cardio and 30 minutes of strength training.\\nA good workout routine is 30 minutes of cardio and 30 minutes of strength training. This is the best way to get in shape',\n",
       " 'How are astronauts launched into space? 1. How are astronauts launched into space? 2. How do astronauts get to the moon? 3. How do astronauts get to the moon? 4. How do astronauts get to']"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.batch_decode(all_token_ids, skip_special_tokens=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations! You have successfully run an end-to-end example using LLaMA-2 in ONNX Runtime. For your convenience, the above code blocks to run the end-to-end example are combined into one code block."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import LlamaConfig, LlamaTokenizer\n",
    "import numpy as np\n",
    "import onnxruntime as ort\n",
    "import torch\n",
    "\n",
    "\n",
    "def get_initial_inputs_and_outputs(config, tokenizer, prompt, device, use_fp16, use_buffer_share):\n",
    "    tokenizer.pad_token = \"[PAD]\"\n",
    "    encodings_dict = tokenizer.batch_encode_plus(prompt, padding=True)\n",
    "    torch_dtype = torch.float16 if use_fp16 else torch.float32\n",
    "\n",
    "    input_ids = torch.tensor(encodings_dict[\"input_ids\"], device=device, dtype=torch.int64)\n",
    "    attention_mask = torch.tensor(encodings_dict[\"attention_mask\"], device=device, dtype=torch.int64)\n",
    "    position_ids = attention_mask.long().cumsum(-1) - 1\n",
    "    position_ids.masked_fill_(attention_mask == 0, 1)\n",
    "    \n",
    "    inputs = {\n",
    "        \"input_ids\": input_ids.contiguous(),\n",
    "        \"attention_mask\": attention_mask.contiguous(),\n",
    "        \"position_ids\": position_ids.contiguous(),\n",
    "    }\n",
    "\n",
    "    batch_size, sequence_length = input_ids.shape\n",
    "    max_sequence_length = config.max_position_embeddings\n",
    "    num_heads, head_size = config.num_key_value_heads, config.hidden_size // config.num_attention_heads\n",
    "    for i in range(config.num_hidden_layers):\n",
    "        past_key = torch.zeros(batch_size, num_heads, max_sequence_length if use_buffer_share else 0, head_size, device=device, dtype=torch_dtype)\n",
    "        past_value = torch.zeros(batch_size, num_heads, max_sequence_length if use_buffer_share else 0, head_size, device=device, dtype=torch_dtype)\n",
    "        inputs.update({\n",
    "            f\"past_key_values.{i}.key\": past_key.contiguous(),\n",
    "            f\"past_key_values.{i}.value\": past_value.contiguous()\n",
    "        })\n",
    "\n",
    "    logits = torch.zeros(batch_size, sequence_length, config.vocab_size, device=device, dtype=torch_dtype)\n",
    "    outputs = {\n",
    "        \"logits\": logits.contiguous()\n",
    "    }\n",
    "    if not use_buffer_share:\n",
    "        for i in range(config.num_hidden_layers):\n",
    "            present_key = torch.zeros(batch_size, num_heads, sequence_length, head_size, device=device, dtype=torch_dtype)\n",
    "            present_value = torch.zeros(batch_size, num_heads, sequence_length, head_size, device=device, dtype=torch_dtype)\n",
    "            outputs.update({\n",
    "                f\"present.{i}.key\": present_key.contiguous(),\n",
    "                f\"present.{i}.value\": present_value.contiguous()\n",
    "            })\n",
    "\n",
    "    return inputs, outputs\n",
    "\n",
    "\n",
    "def apply_io_binding(model, inputs, outputs, use_fp16, use_buffer_share):\n",
    "    # Check that all model inputs will be provided\n",
    "    model_inputs = set(map(lambda model_input: model_input.name, model.get_inputs()))\n",
    "    user_inputs = set(inputs.keys())\n",
    "    missing_inputs = model_inputs - user_inputs\n",
    "    if len(missing_inputs):\n",
    "        print(f\"The following model inputs are missing: {missing_inputs}\")\n",
    "        raise Exception(\"There are missing inputs to the model. Please add them and try again.\")\n",
    "\n",
    "    # Remove unnecessary inputs from model inputs\n",
    "    unnecessary_inputs = user_inputs - model_inputs\n",
    "    if len(unnecessary_inputs):\n",
    "        for unnecessary_input in unnecessary_inputs:\n",
    "            print(f\"Removing unnecessary input '{unnecessary_input}' from user provided inputs\")\n",
    "            del inputs[unnecessary_input]\n",
    "\n",
    "    # Bind inputs/outputs to IO binding\n",
    "    io_binding = model.io_binding()\n",
    "    device = None\n",
    "\n",
    "    for k, v in inputs.items():\n",
    "        io_binding.bind_input(\n",
    "            name=k,\n",
    "            device_type=v.device.type,\n",
    "            device_id=0 if v.device.type == \"cpu\" else v.device.index,\n",
    "            element_type=pt_to_np[repr(v.dtype)],\n",
    "            shape=tuple(v.shape),\n",
    "            buffer_ptr=v.data_ptr()\n",
    "        )\n",
    "        device = v.device\n",
    "\n",
    "    for output in model.get_outputs():\n",
    "        name = output.name\n",
    "        if use_buffer_share and \"present\" in name:\n",
    "            # Bind KV cache outputs to KV cache inputs\n",
    "            v = inputs[name.replace(\"present\", \"past_key_values\")]\n",
    "            io_binding.bind_output(\n",
    "                name=name,\n",
    "                device_type=v.device.type,\n",
    "                device_id=v.device.index,\n",
    "                element_type=np.float16,\n",
    "                shape=tuple(v.shape),\n",
    "                buffer_ptr=v.data_ptr()\n",
    "            )\n",
    "        else:\n",
    "            v = outputs[name]\n",
    "            io_binding.bind_output(\n",
    "                name=name,\n",
    "                device_type=device.type,\n",
    "                device_id=0 if device.type == \"cpu\" else device.index,\n",
    "                element_type=(np.float16 if use_fp16 else np.float32),\n",
    "                shape=tuple(v.shape),\n",
    "                buffer_ptr=v.data_ptr()\n",
    "            )\n",
    "\n",
    "    return io_binding\n",
    "\n",
    "def main():\n",
    "    # User settings\n",
    "    model_name = \"meta-llama/Llama-2-7b-hf\"\n",
    "    onnx_model_path = \"./llama2-7b-fp16-gqa/rank_0_Llama-2-7b-hf_decoder_merged_model_fp16.onnx\"\n",
    "    use_fp16 = True  # True when KV cache inputs/outputs are in float16\n",
    "    use_buffer_share = True  # True when --use_gqa was passed during export\n",
    "\n",
    "    prompt = [\"ONNX Runtime is \", \"I want to book a vacation to Hawaii. First, I need to \", \"A good workout routine is \", \"How are astronauts launched into space? \"]\n",
    "    max_length = 64  # max(prompt length + generation length)\n",
    "\n",
    "    device_id = 0\n",
    "    device = torch.device(f\"cuda:{device_id}\")  # Change to torch.device(\"cpu\") if running on CPU\n",
    "\n",
    "    config = LlamaConfig.from_pretrained(model_name, use_auth_token=True, cache_dir=cache_dir)\n",
    "    tokenizer = LlamaTokenizer.from_pretrained(model_name, use_auth_token=True, cache_dir=cache_dir)\n",
    "    torch_dtype = torch.float16 if use_fp16 else torch.float32\n",
    "\n",
    "    # Get model and its initial inputs/outputs\n",
    "    inputs, outputs = get_initial_inputs_and_outputs(config, tokenizer, prompt, device, use_fp16, use_buffer_share)\n",
    "\n",
    "    sess_options = ort.SessionOptions()\n",
    "    ep = (\"CUDAExecutionProvider\", {\"device_id\": device_id})  # change to ep = \"CPUExecutionProvider\" for CPU\n",
    "    model = ort.InferenceSession(onnx_model_path, sess_options=sess_options, providers=[ep])\n",
    "\n",
    "    all_token_ids = inputs[\"input_ids\"].clone()\n",
    "    batch_size, sequence_length = all_token_ids.shape\n",
    "    num_heads, head_size = config.num_key_value_heads, config.hidden_size // config.num_attention_heads\n",
    "\n",
    "    current_length = sequence_length\n",
    "    has_eos = torch.zeros(batch_size, device=device, dtype=torch.bool)\n",
    "\n",
    "    while current_length <= max_length:\n",
    "        # Run inference\n",
    "        io_binding = apply_io_binding(model, inputs, outputs, use_fp16, use_buffer_share)\n",
    "        io_binding.synchronize_inputs()\n",
    "        model.run_with_iobinding(io_binding)\n",
    "        io_binding.synchronize_outputs()\n",
    "\n",
    "        # Sample with argmax (greedy search)\n",
    "        if outputs[\"logits\"].shape[1] > 1:\n",
    "            prompt_end_indices = inputs[\"attention_mask\"].sum(1) - 1\n",
    "            idxs = prompt_end_indices.unsqueeze(dim=1).repeat(1, config.vocab_size).view(batch_size, 1, config.vocab_size)\n",
    "            next_token_logits = torch.gather(outputs[\"logits\"], 1, idxs).squeeze()\n",
    "        else:\n",
    "            next_token_logits = outputs[\"logits\"][:, -1, :]\n",
    "        next_tokens = torch.argmax(next_token_logits, dim=-1)\n",
    "\n",
    "        # Check if we previously reached EOS token id or if generated token id is EOS token id\n",
    "        has_eos = has_eos | next_tokens == tokenizer.eos_token_id\n",
    "\n",
    "        # Determine which new tokens to add to list of all token ids\n",
    "        # Add EOS token ids for batch entries that ended early (ragged batching scenario where some batch entries ended early and some haven't)\n",
    "        tokens_to_add = next_tokens.masked_fill(has_eos, tokenizer.eos_token_id).reshape([batch_size, 1])\n",
    "        all_token_ids = torch.cat([all_token_ids, tokens_to_add], dim=-1)\n",
    "\n",
    "        # Return early if all batch entries have reached EOS token id\n",
    "        current_length += 1\n",
    "        if torch.all(has_eos) or current_length > max_length:\n",
    "            break\n",
    "\n",
    "        # Update inputs for next inference run\n",
    "        inputs[\"input_ids\"] = tokens_to_add\n",
    "        inputs[\"position_ids\"] = torch.max(inputs[\"position_ids\"], dim=1)[0].reshape(batch_size, 1) + 1\n",
    "        inputs[\"attention_mask\"] = torch.cat([inputs[\"attention_mask\"], (~has_eos).to(torch.int64).reshape(batch_size, 1)], 1)\n",
    "\n",
    "        # Set logits to zeros for next inference run and re-use memory buffer\n",
    "        if outputs[\"logits\"].shape[1] != 1:\n",
    "            outputs[\"logits\"] = outputs[\"logits\"][:, :1, :].contiguous()\n",
    "        outputs[\"logits\"].zero_()\n",
    "\n",
    "        if not use_buffer_share:\n",
    "            for i in range(config.num_hidden_layers):\n",
    "                inputs[f\"past_key_values.{i}.key\"] = outputs[f\"present.{i}.key\"]\n",
    "                inputs[f\"past_key_values.{i}.value\"] = outputs[f\"present.{i}.value\"]\n",
    "\n",
    "            new_sequence_length = inputs[\"attention_mask\"].shape[1]\n",
    "            for i in range(config.num_hidden_layers):\n",
    "                present_key = torch.zeros(batch_size, num_heads, new_sequence_length, head_size, device=device, dtype=torch_dtype)\n",
    "                present_value = torch.zeros(batch_size, num_heads, new_sequence_length, head_size, device=device, dtype=torch_dtype)\n",
    "                outputs.update({\n",
    "                    f\"present.{i}.key\": present_key.contiguous(),\n",
    "                    f\"present.{i}.value\": present_value.contiguous()\n",
    "                })\n",
    "\n",
    "    # Batch decoding at end of generation\n",
    "    print(tokenizer.batch_decode(all_token_ids, skip_special_tokens=True))\n",
    "\n",
    "pt_to_np = {\n",
    "    \"torch.int64\": np.int64,\n",
    "    \"torch.float32\": np.float32,\n",
    "    \"torch.float16\": np.float16\n",
    "}\n",
    "main()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that you can also disable buffer sharing when using GroupQueryAttention. Inference will still work but performance will be worse."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Without buffer sharing, you can also run your model using [Hugging Face's Optimum](https://github.com/huggingface/optimum). Here's how you can use Optimum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import LlamaConfig, LlamaTokenizer\n",
    "from optimum.onnxruntime import ORTModelForCausalLM\n",
    "import torch\n",
    "\n",
    "# User settings\n",
    "model_name = \"meta-llama/Llama-2-7b-hf\"\n",
    "onnx_model_dir = \"./llama2-7b-fp16-gqa/\"\n",
    "cache_dir = \"./cache_dir\"\n",
    "\n",
    "device_id = 0\n",
    "device = torch.device(f\"cuda:{device_id}\")  # Change to torch.device(\"cpu\") if running on CPU\n",
    "\n",
    "ep = \"CUDAExecutionProvider\"  # change to CPUExecutionProvider if running on CPU\n",
    "ep_options = {\"device_id\": device_id}\n",
    "\n",
    "prompt = [\"ONNX Runtime is \", \"I want to book a vacation to Hawaii. First, I need to \", \"A good workout routine is \", \"How are astronauts launched into space? \"]\n",
    "max_length = 64  # max(prompt length + generation length)\n",
    "\n",
    "config = LlamaConfig.from_pretrained(model_name, use_auth_token=True, cache_dir=cache_dir)\n",
    "config.save_pretrained(onnx_model_dir)  # Save config file in ONNX model directory\n",
    "tokenizer = LlamaTokenizer.from_pretrained(model_name, use_auth_token=True, cache_dir=cache_dir)\n",
    "tokenizer.pad_token = \"[PAD]\"\n",
    "\n",
    "model = ORTModelForCausalLM.from_pretrained(\n",
    "    onnx_model_dir,\n",
    "    use_auth_token=True,\n",
    "    use_io_binding=True,\n",
    "    provider=ep,\n",
    "    provider_options={\"device_id\": device_id}  # comment out if running on CPU\n",
    ")\n",
    "inputs = tokenizer(prompt, return_tensors=\"pt\", padding=True).to(device)\n",
    "\n",
    "print(\"-------------\")\n",
    "generate_ids = model.generate(**inputs, do_sample=False, max_length=max_length)\n",
    "transcription = tokenizer.batch_decode(generate_ids, skip_special_tokens=True)\n",
    "print(transcription)\n",
    "print(\"-------------\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-notebook",
   "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.9.18"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}