demo/BERT/notebooks/benchmark.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "81fd5338-bf4e-45b2-8c8c-a6e8a7da07b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2021 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# ==============================================================================\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4896018-e56e-464f-9edc-a10057ea45d9",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# BERT TensorRT Benchmarking: FP16 vs. INT8 QAT+Sparsity\n",
    "\n",
    "In this notebook, we benchmark different BERT Large TensorRT engines at different batch sizes.\n",
    "\n",
    "**Notes**: \n",
    " - Prior to running this notebook, run [BERT-TRT-FP16.ipynb](BERT-TRT-FP16.ipynb) and [BERT-TRT-INT8-QAT-sparse.ipynb](BERT-TRT-INT8-QAT-sparse.ipynb) to generate the TensorRT engines.\n",
    " - This benchmarking focuses on the compute part using synthetic inputs, without taking into account pre and post processing time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c89592ab-e50b-47c9-9b5b-6b5067d0b22a",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import sys\n",
    "sys.path.append('/workspace/TensorRT/demo/BERT')\n",
    "\n",
    "import tensorrt as trt;\n",
    "TRT_VERSION = trt.__version__\n",
    "\n",
    "import time\n",
    "import argparse\n",
    "import ctypes\n",
    "import numpy as np\n",
    "import tensorrt as trt\n",
    "import pycuda.driver as cuda\n",
    "import pycuda.autoinit\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "import ipywidgets as widgets\n",
    "from ipywidgets import IntProgress\n",
    "from ipywidgets import Button, Layout\n",
    "from IPython.display import display\n",
    "\n",
    "import helpers.tokenization as tokenization\n",
    "import helpers.data_processing as dp\n",
    "\n",
    "TRT_LOGGER = trt.Logger(trt.Logger.ERROR)\n",
    "\n",
    "class DeviceBuffer(object):\n",
    "    def __init__(self, shape, dtype=trt.int32):\n",
    "        self.buf = cuda.mem_alloc(trt.volume(shape) * dtype.itemsize)\n",
    "\n",
    "    def binding(self):\n",
    "        return int(self.buf)\n",
    "\n",
    "    def free(self):\n",
    "        self.buf.free()\n",
    "\n",
    "doc_stride = 128\n",
    "max_query_length = 64\n",
    "max_seq_length = 384\n",
    "\n",
    "vocab_file = \"models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt\"\n",
    "tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file, do_lower_case=True)\n",
    "\n",
    "parser = argparse.ArgumentParser(description='BERT Inference Benchmark')\n",
    "parser.add_argument(\"-e\", \"--engine\", help='Path to BERT TensorRT engine', default='')\n",
    "parser.add_argument('-b', '--batch-size', default=[], action=\"append\", help='Batch size(s) to benchmark. Can be specified multiple times for more than one batch size. This script assumes that the engine has been built with one optimization profile for each batch size, and that these profiles are in order of increasing batch size.', type=int)\n",
    "parser.add_argument('-s', '--sequence-length', default=384, help='Sequence length of the BERT model', type=int)\n",
    "parser.add_argument('-i', '--iterations', default=1000, help='Number of iterations to run when benchmarking each batch size.', type=int)\n",
    "parser.add_argument('-w', '--warm-up-runs', default=10, help='Number of iterations to run prior to benchmarking.', type=int)\n",
    "parser.add_argument('-r', '--random-seed', required=False, default=12345, help='Random seed.', type=int)\n",
    "args, _ = parser.parse_known_args()\n",
    "args.batch_size = args.batch_size or [1]\n",
    "\n",
    "# Import necessary plugins for BERT TensorRT\n",
    "ctypes.CDLL(\"libnvinfer_plugin.so\", mode=ctypes.RTLD_GLOBAL)\n",
    "\n",
    "### INT8 TRT model\n",
    "def run_benchmark_INT8(b):\n",
    "    engine_path = \"engines_%s/megatron_large_seqlen384_int8qat_sparse.engine\"%TRT_VERSION\n",
    "    with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime, runtime.deserialize_cuda_engine(f.read()) as engine, engine.create_execution_context() as context:\n",
    "        with output:\n",
    "            #output.clear_output()\n",
    "            args.batch_size = [int(batchsize_selector.value)]\n",
    "\n",
    "            # Allocate buffers large enough to store the largest batch size\n",
    "            max_input_shape = (args.sequence_length * max(args.batch_size), )\n",
    "            max_output_shape = (args.sequence_length * max(args.batch_size), 2, 1, 1)\n",
    "            buffers = [\n",
    "                DeviceBuffer(max_input_shape),\n",
    "                DeviceBuffer(max_input_shape),\n",
    "                DeviceBuffer((max(args.batch_size) + 1, )),\n",
    "                DeviceBuffer((args.sequence_length, )),\n",
    "                DeviceBuffer(max_output_shape)\n",
    "            ]\n",
    "\n",
    "            # Prepare random input\n",
    "            pseudo_vocab_size = 30522\n",
    "            pseudo_type_vocab_size = 2\n",
    "            np.random.seed(args.random_seed)\n",
    "            test_word_ids = np.random.randint(0, pseudo_vocab_size, (args.sequence_length * max(args.batch_size)), dtype=np.int32)\n",
    "            test_segment_ids = np.random.randint(0, pseudo_type_vocab_size, (args.sequence_length * max(args.batch_size)), dtype=np.int32)\n",
    "            test_cu_seq_lens = np.arange(0, args.sequence_length * max(args.batch_size) + 1, args.sequence_length, dtype=np.int32)\n",
    "\n",
    "            # Copy input h2d\n",
    "            cuda.memcpy_htod(buffers[0].buf, test_word_ids.ravel())\n",
    "            cuda.memcpy_htod(buffers[1].buf, test_segment_ids.ravel())\n",
    "            cuda.memcpy_htod(buffers[2].buf, test_cu_seq_lens.ravel())\n",
    "\n",
    "            bench_times = {}\n",
    "            stream = cuda.Stream()\n",
    "\n",
    "            for idx, batch_size in enumerate(sorted(args.batch_size)):\n",
    "                for idx in range(engine.num_optimization_profiles):\n",
    "                    profile_shape = engine.get_profile_shape(profile_index = idx, binding = idx * engine.num_io_tensors)\n",
    "                    if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size:\n",
    "                        context.set_optimization_profile_async(idx, stream.handle)\n",
    "                        binding_idx_offset = idx * engine.num_io_tensors\n",
    "                        break\n",
    "\n",
    "                # Each profile has unique bindings\n",
    "                bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]\n",
    "                input_shape = (batch_size, args.sequence_length)\n",
    "                for binding in range(3):\n",
    "                    tensor_name = engine.get_tensor_name(binding)\n",
    "                    context.set_input_shape(tensor_name, input_shape)\n",
    "                assert context.all_binding_shapes_specified\n",
    "\n",
    "                for i in range(engine.num_io_tensors):\n",
    "                    context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])\n",
    "\n",
    "                # Inference\n",
    "                total_time = 0\n",
    "                start = cuda.Event()\n",
    "                end = cuda.Event()\n",
    "\n",
    "                # Warmup\n",
    "                for _ in range(args.warm_up_runs):\n",
    "                    context.execute_async_v3(stream_handle=stream.handle)\n",
    "                    stream.synchronize()\n",
    "\n",
    "                # Timing loop\n",
    "                times = []\n",
    "                progress_bar.value = 0\n",
    "                for _ in range(iteration_selector.value):\n",
    "                    start.record(stream)\n",
    "                    context.execute_async_v3(stream_handle=stream.handle)\n",
    "                    end.record(stream)\n",
    "                    stream.synchronize()\n",
    "                    times.append(end.time_since(start))\n",
    "                    progress_bar.value +=1\n",
    "\n",
    "                # Compute average time, 95th percentile time and 99th percentile time.\n",
    "                bench_times[batch_size] = times\n",
    "\n",
    "            [b.free() for b in buffers]\n",
    "\n",
    "            for batch_size, times in bench_times.items():\n",
    "                total_time = sum(times)\n",
    "                avg_time = total_time / float(len(times))\n",
    "                times.sort()\n",
    "                percentile95 = times[int(len(times) * 0.95)]\n",
    "                percentile99 = times[int(len(times) * 0.99)]\n",
    "                print(\"BERT TRT INT8: Running {:} iterations with Batch Size: {:}\\n\\tTotal Time: {:.2f} ms \\tAverage Time: {:.2f} ms\\t95th Percentile Time: {:.2f} ms\\t99th Percentile Time: {:.2f}\".format(args.iterations, batch_size, total_time, avg_time, percentile95, percentile99))\n",
    "\n",
    "\n",
    "### FP16 TRT model\n",
    "def run_benchmark_FP16(b):\n",
    "    engine_path = \"engines_%s/bert_large_384.engine\"%TRT_VERSION\n",
    "    with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime, runtime.deserialize_cuda_engine(f.read()) as engine, engine.create_execution_context() as context:\n",
    "        with output:\n",
    "            #output.clear_output()\n",
    "            args.batch_size = [int(batchsize_selector.value)]\n",
    "\n",
    "            # Allocate buffers large enough to store the largest batch size\n",
    "            max_input_shape = (max(args.batch_size), args.sequence_length)\n",
    "            max_output_shape = (max(args.batch_size), args.sequence_length, 2, 1, 1)\n",
    "            buffers = [\n",
    "                DeviceBuffer(max_input_shape),\n",
    "                DeviceBuffer(max_input_shape),\n",
    "                DeviceBuffer(max_input_shape),\n",
    "                DeviceBuffer(max_output_shape)\n",
    "            ]\n",
    "\n",
    "            # Prepare random input\n",
    "            pseudo_vocab_size = 30522\n",
    "            pseudo_type_vocab_size = 2\n",
    "            np.random.seed(args.random_seed)\n",
    "            test_word_ids = np.random.randint(0, pseudo_vocab_size, (max(args.batch_size), args.sequence_length), dtype=np.int32)\n",
    "            test_segment_ids = np.random.randint(0, pseudo_type_vocab_size, (max(args.batch_size), args.sequence_length), dtype=np.int32)\n",
    "            test_input_mask = np.ones((max(args.batch_size), args.sequence_length), dtype=np.int32)\n",
    "\n",
    "            # Copy input h2d\n",
    "            cuda.memcpy_htod(buffers[0].buf, test_word_ids.ravel())\n",
    "            cuda.memcpy_htod(buffers[1].buf, test_segment_ids.ravel())\n",
    "            cuda.memcpy_htod(buffers[2].buf, test_input_mask.ravel())\n",
    "\n",
    "            bench_times = {}\n",
    "            stream = cuda.Stream()\n",
    "\n",
    "            for idx, batch_size in enumerate(sorted(args.batch_size)):\n",
    "                for idx in range(engine.num_optimization_profiles):\n",
    "                    profile_shape = engine.get_profile_shape(profile_index = idx, binding = idx * engine.num_io_tensors)\n",
    "                    if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size:\n",
    "                        context.set_optimization_profile_async(idx, stream.handle)\n",
    "                        binding_idx_offset = idx * engine.num_io_tensors\n",
    "                        break\n",
    "\n",
    "                # Each profile has unique bindings\n",
    "                bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]\n",
    "                input_shape = (batch_size, args.sequence_length)\n",
    "                for binding in range(3):\n",
    "                    tensor_name = engine.get_tensor_name(binding)\n",
    "                    context.set_input_shape(tensor_name, input_shape)\n",
    "                assert context.all_binding_shapes_specified\n",
    "\n",
    "                for i in range(engine.num_io_tensors):\n",
    "                    context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])\n",
    "\n",
    "                # Inference\n",
    "                total_time = 0\n",
    "                start = cuda.Event()\n",
    "                end = cuda.Event()\n",
    "                stream = cuda.Stream()\n",
    "\n",
    "                # Warmup\n",
    "                for _ in range(args.warm_up_runs):\n",
    "                    context.execute_async_v3(stream_handle=stream.handle)\n",
    "                    stream.synchronize()\n",
    "\n",
    "                # Timing loop\n",
    "                times = []\n",
    "                progress_bar.value = 0\n",
    "                for _ in range(iteration_selector.value):\n",
    "                    start.record(stream)\n",
    "                    context.execute_async_v3(stream_handle=stream.handle)\n",
    "                    end.record(stream)\n",
    "                    stream.synchronize()\n",
    "                    times.append(end.time_since(start))\n",
    "                    progress_bar.value +=1\n",
    "\n",
    "                # Compute average time, 95th percentile time and 99th percentile time.\n",
    "                bench_times[batch_size] = times\n",
    "\n",
    "            [b.free() for b in buffers]\n",
    "\n",
    "            for batch_size, times in bench_times.items():\n",
    "                total_time = sum(times)\n",
    "                avg_time = total_time / float(len(times))\n",
    "                times.sort()\n",
    "                percentile95 = times[int(len(times) * 0.95)]\n",
    "                percentile99 = times[int(len(times) * 0.99)]\n",
    "                print(\"BERT TRT FP16: Running {:} iterations with Batch Size: {:}\\n\\tTotal Time: {:.2f} ms \\tAverage Time: {:.2f} ms\\t95th Percentile Time: {:.2f} ms\\t99th Percentile Time: {:.2f}\".format(args.iterations, batch_size, total_time, avg_time, percentile95, percentile99))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c72b6e9d-3181-4c22-8161-c1a1a9444e2d",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fcd71564172944a198f873066199c7e1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "RadioButtons(description='Engine:', options=('GPU - TensorRT FP16', 'GPU - TensorRT INT8'), value='GPU - Tenso…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8b51334850f84db4834a1d0035556971",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "RadioButtons(description='Batch size:', options=('1', '32', '64', '128'), value='1')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "44b1a0bd02834393b7d7204c18395578",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "IntSlider(value=500, continuous_update=False, description='Iterations:', max=1000, min=100)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f988290aab5243fb9d43a8581d104d64",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Button(description='Run benchmark', style=ButtonStyle()),), layout=Layout(align_items='center',…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "03b48cb93c834fe89f2180dc1562a4c7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3a3756b664bc482b843283c2e6cff3fa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "IntProgress(value=0, description='Progress:', layout=Layout(height='50px', width='100%'), max=1000, style=Prog…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# UI elements\n",
    "engine_selector = widgets.RadioButtons(\n",
    "    options=['GPU - TensorRT FP16',\n",
    "             'GPU - TensorRT INT8'],\n",
    "    description='Engine:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "batchsize_selector = widgets.RadioButtons(\n",
    "    options=['1', '32', '64', '128'],\n",
    "    description='Batch size:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "iteration_selector = widgets.IntSlider(\n",
    "    value=500,\n",
    "    min=100,\n",
    "    max=1000,\n",
    "    step=1,\n",
    "    description='Iterations:',\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='d'\n",
    ")\n",
    "\n",
    "button = widgets.Button(description=\"Run benchmark\")\n",
    "output = widgets.Output()\n",
    "#box = widgets.HBox(children=[button],layout=box_layout)\n",
    "\n",
    "def run_benchmark(b):\n",
    "    args.iterations = iteration_selector.value\n",
    "    progress_bar.max = iteration_selector.value\n",
    "    with output:\n",
    "        if engine_selector.value=='GPU - TensorRT FP16':\n",
    "            run_benchmark_FP16(b)\n",
    "        elif engine_selector.value=='GPU - TensorRT INT8':\n",
    "            run_benchmark_INT8(b)\n",
    "        \n",
    "button.on_click(run_benchmark)\n",
    "display(engine_selector, batchsize_selector, iteration_selector)\n",
    "\n",
    "from IPython.display import display\n",
    "box_layout = widgets.Layout(display='flex',\n",
    "                flex_flow='column',\n",
    "                align_items='center',\n",
    "                width='100%')\n",
    "box = widgets.HBox(children=[button],layout=box_layout)\n",
    "display(box, output)\n",
    "\n",
    "progress_bar = widgets.IntProgress(\n",
    "    value=0,\n",
    "    min=0,\n",
    "    max=1000,\n",
    "    description='Progress:',\n",
    "    bar_style='',\n",
    "    style={'bar_color': 'green'},\n",
    "    orientation='horizontal', \n",
    "    layout=Layout(width='100%', height='50px')\n",
    ")\n",
    "display(progress_bar)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1317f3b9-6bd0-4f7b-9956-b33e5948a898",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}