DLBoost_vs_CUDA_on_BERT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Intel CPU/DLBoost vs. nvidia GPU/CUDA on BERT in Tensorflow 2.x\n",
    "\n",
    "Prepared by Forrest Sheng Bao\n",
    "\n",
    "For NLP class at Iowa State University\n",
    "\n",
    "Adapted from https://www.tensorflow.org/official_models/fine_tuning_bert\n",
    "\n",
    "# Introduction\n",
    "To gain ground in the deep learing (DL) revolution, Intel released [DL Boost](https://www.intel.com/content/www/us/en/artificial-intelligence/deep-learning-boost.html), a set of new instructions and software tools, such that their CPUs won't look too slow in DL training. Intel is also maintaining [a tailored version of Tensorflow](https://software.intel.com/content/www/us/en/develop/articles/intel-optimization-for-tensorflow-installation-guide.html) for their CPUs, including those with DL Boost.  But how does it really turn out? \n",
    "\n",
    "To compare the DL power between Intel's CPUs with DLBoost and nVidia's consumer-level GPUs, I will finetune the BERT base model on the MRPC task. \n",
    "Same Python code, same machine just different (co)-processors to run the code. \n",
    "\n",
    "\n",
    "**Configurations**: \n",
    "* Intel CPU: i9-10980XE, 18 cores, 165W, $1000 retail price\n",
    "* nVidia GPU: RTX 3090, 350W, $1500 retail price (although it could be $3000)\n",
    "* RAM 64GB\n",
    "\n",
    "All packages are binary installed via `pip`. \n",
    "\n",
    "# Conclusion (for those who cannot wait)\n",
    "\n",
    "Intel lost the game completely! \n",
    "\n",
    "**training speed**:\n",
    "* 3 seconds per step in Intel-optimized Tensorflow on an Intel CPU (10980XE) with DL Boost\n",
    "* 140 milliseconds per step in generic Tensorflow on RTX 3090. \n",
    "\n",
    "The Intel CPU is 50% power-hungry but only 5% capable compared to RTX 3090. \n",
    "In other words, for the same amount of energy consumed, the Intel CPU only gets 1/10 of the work done. \n",
    "If you further factor in the price tag (speed per Watt per dollar), the Intel CPU is still only 1/3 worthy than RTX 3090. \n",
    "\n",
    "**Takeaway**: If you are an DL researcher, do not wait your time and money on using Intel CPUs. Follow the crowd. Use GPUs. \n",
    "\n",
    "\n",
    "# Setting up Intel-optimized Tensorflow \n",
    "\n",
    "Create an virtual environment if you do not want Intel's TF to mess with official TF that uses CUDA. \n",
    "When not detecting the GPU, the official/generic TF will also use CPU cores but that would be generic TF code which cannot fully unleash the power of Intel DL Boost. \n",
    "\n",
    "First,\n",
    "\n",
    "```bash\n",
    "pip install  tensorflow_hub tensorflow_datasets  tf-models-official\n",
    "```\n",
    "\n",
    "Due to dependencies, this will install the official Tensorflow which has no \n",
    "special optimization for Intel's CPUs. \n",
    "\n",
    "Second, let's uninstall the official Tensorflow:\n",
    "```bash\n",
    "pip uninstall tensorflow\n",
    "```\n",
    "\n",
    "Third, install the Intel-optimized Tensorflow. \n",
    "```bash\n",
    "pip install intel-tensorflow-avx512\n",
    "```\n",
    "\n",
    "Now, please select the proper virtual environment, kernel and/or Tensorflow \n",
    "version to run the experiment. "
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "source": [
    "import os,json\n",
    "\n",
    "import numpy as np\n",
    "# import matplotlib.pyplot as plt\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "import tensorflow_hub as hub\n",
    "import tensorflow_datasets as tfds\n",
    "tfds.disable_progress_bar()\n",
    "\n",
    "# Load BERT related \n",
    "# !pip install tensorflow_hub tensorflow_datasets\n",
    "# !pip3 install -q tf-models-official==2.4.0\n",
    "from official.modeling import tf_utils\n",
    "from official import nlp\n",
    "from official.nlp import bert\n",
    "import official.nlp.optimization\n",
    "import official.nlp.bert.bert_models\n",
    "import official.nlp.bert.configs\n",
    "import official.nlp.bert.run_classifier\n",
    "import official.nlp.bert.tokenization\n",
    "import official.nlp.data.classifier_data_lib\n",
    "import official.nlp.modeling.losses\n",
    "import official.nlp.modeling.models\n",
    "import official.nlp.modeling.networks"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "source": [
    "# Settings for Intel CPUs. Should have no impact if you use GPU. \n",
    "os.environ[\"KMP_BLOCKTIME\"] = \"1\"\n",
    "\n",
    "os.environ[\"KMP_SETTINGS\"] = \"1\"\n",
    "\n",
    "os.environ[\"KMP_AFFINITY\"]= \"granularity=fine,verbose,compact,1,0\"\n",
    "\n",
    "os.environ[\"OMP_NUM_THREADS\"]= \"36\"\n",
    "\n",
    "os.environ[\"TF_ENABLE_ONEDNN_OPTS\"]=\"1\"\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Loading the task data \n",
    "Here we use MRPC task as an example. "
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "source": [
    "glue, info = tfds.load('glue/mrpc', with_info=True,\n",
    "                       # It's small, load the whole dataset\n",
    "                       batch_size=-1)\n",
    "print (\"The splits are: \\n\\t\", list(glue.keys()))\n",
    "print (\"Each sample is organized as: \\n\", info.features)\n",
    "# print (\"Each sample has two labels in MRPC: \\n\", info.features['label'].names)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "The splits are: \n",
      "\t ['train', 'validation', 'test']\n",
      "Each sample is organized as: \n",
      " FeaturesDict({\n",
      "    'idx': tf.int32,\n",
      "    'label': ClassLabel(shape=(), dtype=tf.int64, num_classes=2),\n",
      "    'sentence1': Text(shape=(), dtype=tf.string),\n",
      "    'sentence2': Text(shape=(), dtype=tf.string),\n",
      "})\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Preparing the data"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Specifying the location of BERT model files\n",
    "from a Google Cloud location\n",
    "\n",
    "Here we use the base BERT model. "
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "source": [
    "gs_folder_bert = \"gs://cloud-tpu-checkpoints/bert/v3/uncased_L-12_H-768_A-12\"\n",
    "tf.io.gfile.listdir(gs_folder_bert)"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "['bert_config.json',\n",
       " 'bert_model.ckpt.data-00000-of-00001',\n",
       " 'bert_model.ckpt.index',\n",
       " 'vocab.txt']"
      ]
     },
     "metadata": {},
     "execution_count": 17
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Setting up and testing Tokenizer for BERT "
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "source": [
    "# Set up tokenizer to generate Tensorflow dataset\n",
    "tokenizer = bert.tokenization.FullTokenizer(\n",
    "    vocab_file=os.path.join(gs_folder_bert, \"vocab.txt\"),\n",
    "     do_lower_case=True)\n",
    "\n",
    "# Test \n",
    "print(\"Vocab size:\", len(tokenizer.vocab))\n",
    "tokens = tokenizer.tokenize(\"The multilaterial relationship is important between geopolitical organizations and multinational conglomerates\")\n",
    "print(tokens)\n",
    "ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "print(ids)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Vocab size: 30522\n",
      "['the', 'multi', '##late', '##rial', 'relationship', 'is', 'important', 'between', 'geo', '##pol', '##itical', 'organizations', 'and', 'multinational', 'conglomerate', '##s']\n",
      "[1996, 4800, 13806, 14482, 3276, 2003, 2590, 2090, 20248, 18155, 26116, 4411, 1998, 20584, 22453, 2015]\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Generating token IDs, input masks, and input types"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "source": [
    "def encode_sentence(s):\n",
    "   tokens = list(tokenizer.tokenize(s))\n",
    "#    tokens.append('[SEP]')\n",
    "   return tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "# TODO: how does this compare with TF1.x version and how we can parallelize to\n",
    "#       multiple CPU cores \n",
    "def bert_encode(glue, split):\n",
    "#   num_examples = len(glue[split][\"sentence1\"])\n",
    "\n",
    "  sentence1 = tf.ragged.constant([\n",
    "      encode_sentence(s)\n",
    "      for s in np.array(glue[split][\"sentence1\"])])\n",
    "  sentence2 = tf.ragged.constant([\n",
    "      encode_sentence(s)\n",
    "       for s in np.array(glue[split][\"sentence2\"])])\n",
    "\n",
    "  cls_column = [tokenizer.convert_tokens_to_ids(['[CLS]'])]*sentence1.shape[0]\n",
    "  sep_column = [tokenizer.convert_tokens_to_ids(['[SEP]'])]*sentence1.shape[0]\n",
    "  input_word_ids = tf.concat([cls_column, sentence1, sep_column, sentence2, sep_column], axis=-1)\n",
    "\n",
    "  input_mask = tf.ones_like(input_word_ids).to_tensor()\n",
    "\n",
    "  type_cls = tf.zeros_like(cls_column)\n",
    "  type_sep = tf.zeros_like(sep_column)\n",
    "  type_s1 = tf.zeros_like(sentence1)\n",
    "  type_s2 = tf.ones_like(sentence2)\n",
    "  input_type_ids = tf.concat(\n",
    "      [type_cls, type_s1, type_sep, type_s2, type_sep], axis=-1).to_tensor()\n",
    "\n",
    "  inputs = {\n",
    "      'input_word_ids': input_word_ids.to_tensor(),\n",
    "      'input_mask': input_mask,\n",
    "      'input_type_ids': input_type_ids}\n",
    "\n",
    "  return inputs\n",
    "\n",
    "task_data = {\n",
    "    split: {'inputs':bert_encode(glue, split), 'labels':glue[split]['label']} \n",
    "            for split in glue.keys()\n",
    "}\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Loading the BERT model\n",
    "`gs_folder_bert` was defined earlier when loading tokenzier"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "source": [
    "bert_config_file = os.path.join(gs_folder_bert, \"bert_config.json\")\n",
    "config_dict = json.loads(tf.io.gfile.GFile(bert_config_file).read())\n",
    "bert_config = bert.configs.BertConfig.from_dict(config_dict)\n",
    "\n",
    "bert_classifier, bert_encoder = bert.bert_models.classifier_model(\n",
    "    bert_config, num_labels=2)\n",
    "\n",
    "\n",
    "checkpoint = tf.train.Checkpoint(encoder=bert_encoder)\n",
    "checkpoint.read(\n",
    "    os.path.join(gs_folder_bert, 'bert_model.ckpt')).assert_consumed()"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "2021-10-02 18:33:57.583654: E tensorflow/core/platform/cloud/curl_http_request.cc:614] The transmission  of request 0x1415b1a0 (URI: https://storage.googleapis.com/cloud-tpu-checkpoints/bert%2Fv3%2Funcased_L-12_H-768_A-12%2Fbert_config.json) has been stuck at 0 of 0 bytes for 61 seconds and will be aborted. CURL timing information: lookup time: 0.000424 (No error), connect time: 0.014014 (No error), pre-transfer time: 0 (No error), start-transfer time: 0 (No error)\n",
      "2021-10-02 18:36:06.138441: E tensorflow/core/platform/cloud/curl_http_request.cc:614] The transmission  of request 0x14ec1150 (URI: https://storage.googleapis.com/cloud-tpu-checkpoints/bert%2Fv3%2Funcased_L-12_H-768_A-12%2Fbert_model.ckpt.data-00000-of-00001) has been stuck at 0 of 0 bytes for 61 seconds and will be aborted. CURL timing information: lookup time: 0.000378 (No error), connect time: 0.015309 (No error), pre-transfer time: 0 (No error), start-transfer time: 0 (No error)\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7f0cc9b3c8e0>"
      ]
     },
     "metadata": {},
     "execution_count": 20
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Fine-tuning of the model "
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "source": [
    "# Set up epochs and steps\n",
    "epochs = 3\n",
    "batch_size = 32\n",
    "eval_batch_size = 32\n",
    "\n",
    "train_data_size = len(task_data['train']['labels'])\n",
    "steps_per_epoch = int(train_data_size / batch_size)\n",
    "num_train_steps = steps_per_epoch * epochs\n",
    "warmup_steps = int(epochs * train_data_size * 0.1 / batch_size)\n",
    "\n",
    "# creates an optimizer with learning rate schedule\n",
    "optimizer = nlp.optimization.create_optimizer(\n",
    "    2e-5, num_train_steps=num_train_steps, num_warmup_steps=warmup_steps)\n",
    "\n",
    "\n",
    "metrics = [tf.keras.metrics.SparseCategoricalAccuracy('accuracy', dtype=tf.float32)]\n",
    "loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "\n",
    "bert_classifier.compile(\n",
    "    optimizer=optimizer,\n",
    "    loss=loss,\n",
    "    metrics=metrics)\n",
    "\n",
    "bert_classifier.fit(\n",
    "      task_data['train']['inputs'], task_data['train']['labels'], \n",
    "      validation_data=(task_data['validation']['inputs'], task_data['validation']['labels']),\n",
    "      batch_size=batch_size,\n",
    "      epochs=epochs)\n"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Epoch 1/3\n",
      "115/115 [==============================] - 393s 3s/step - loss: 0.6354 - accuracy: 0.6050 - val_loss: 0.5164 - val_accuracy: 0.7598\n",
      "Epoch 2/3\n",
      "115/115 [==============================] - 382s 3s/step - loss: 0.4595 - accuracy: 0.7816 - val_loss: 0.4377 - val_accuracy: 0.8088\n",
      "Epoch 3/3\n",
      "115/115 [==============================] - 381s 3s/step - loss: 0.3356 - accuracy: 0.8604 - val_loss: 0.4397 - val_accuracy: 0.8162\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7f0cc9b2c610>"
      ]
     },
     "metadata": {},
     "execution_count": 21
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "orig_nbformat": 4,
  "language_info": {
   "name": "python",
   "version": "3.8.10",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.10 64-bit ('inteltf': venv)"
  },
  "interpreter": {
   "hash": "542670f3cc0d5a156cc19a4956bd40e81355002e08cbdd7546ca5f6f6e22bbea"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}