examples/question_answering/pretrained-BERT-SQuAD-deep-dive-aml.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) Microsoft Corporation. All rights reserved.\n",
    "\n",
    "Licensed under the MIT License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Question Answering: Fine-Tune BERT on AzureML (PyTorch)\n",
    "**BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding** [\\[1\\]](#References)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Impressions](https://PixelServer20190423114238.azurewebsites.net/api/impressions/nlp/examples/question_answering/pretrained_BERT_SQuAD_deep_dive_aml.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook contains an end-to-end walkthrough using [Azure Machine Learning service (AzureML)](https://azure.microsoft.com/en-us/services/machine-learning-service/) to fine-tune the pretrained [PyTorch implementation](https://github.com/huggingface/pytorch-pretrained-BERT) of [Google's TensorFlow repository for the BERT model](https://github.com/google-research/bert) developed by Hugging Face. \n",
    "\n",
    "**Note: To learn how to do pre-training on your own, please reference the [AzureML-BERT repo](https://github.com/microsoft/AzureML-BERT) created by Microsoft.**\n",
    "\n",
    "This notebook will walk through the following:\n",
    "- Download the SQuAD dataset on a remote compute and store the dataset in Azure Blob storage\n",
    "- Fine-tune BERT with distributed PyTorch by Horovod for SQuAD dataset using GPU clusters provided by AzureML\n",
    "- Further fine-tune BERT with AzureML's hyperparameter tuning "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is BERT?\n",
    "\n",
    "[BERT (Bidirectional Encoder Representations from Transformers)](https://arxiv.org/abs/1810.04805) is a powerful pre-trained language model by presenting state-of-the-art results in a wide variety of NLP tasks, including Question Answering (SQuAD v1.1), Natural language Inference (MNLI), Text Classification, Name Entity Recognition, etc., by only a few epochs of fine tuning on task specific datasets. The key technical innovation of BERT is applying the bidirectional training of Transformer, which is a popular attention model that learns contextual relations between words (or sub-words) in a text."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How to fine-tune BERT for QA\n",
    "\n",
    "The figure below shows how BERT can be fine tuned for Question and Answering (QA) tasks. BERT plugs the question-passage pairs in SQuAD dataset as the inputs, and the `[SEP]` representation is a special separator token for separating questions/answers. At the output layer, it outputs `Start/End` to denote the answer in the paragraph.\n",
    "\n",
    "<img src=\"https://nlpbp.blob.core.windows.net/images/bertqa.PNG\" height=400 width=400>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is the SQuAD dataset?\n",
    "\n",
    "\"Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span, from the corresponding reading passage, or the question might be unanswerable\". [\\[2\\]](#References)\n",
    "\n",
    "\"SQuAD 1.1, the previous version of the SQuAD dataset, contains 100,000+ question-answer pairs on 500+ articles\". [\\[2\\]](#References) More details from [https://rajpurkar.github.io/SQuAD-explorer/](https://rajpurkar.github.io/SQuAD-explorer/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Table of Contents\n",
    "\n",
    "<ol start=\"0\">\n",
    "  <li> [Prerequisites: Global Settings](#0.-Prerequisites:-Global-Settings)</li>\n",
    "  <li> [Data Loading](#1.-Data-Loading)\n",
    "  <ul style=\"list-style: none;\"><li>[1.1 Default AzureML datastore](#1.1-Default-AzureML-datastore)</li>\n",
    "        <li>[1.2 Download training dataset - SQuAD](#1.2-Download-training-dataset---SQuAD)</li>\n",
    "        <li>[1.3 Upload to Azure blob storage](#1.3-Upload-to-Azure-blob-storage)</li>\n",
    "    </ul>\n",
    "  </li><br>\n",
    "  <li> [Fine tuning BERT with Distributed Training by Horovod](#2.-Fine-tuning-BERT-with-Distributed-Training-by-Horovod)\n",
    "    <ul style=\"list-style: none;\">\n",
    "        <li> [2.1 Create or Attach Existing AmlCompute](#2.1-Create-a-GPU-remote-compute-target)</li>\n",
    "        <li> [2.2 Access to a Project Directory](#2.2-Create-a-project-directory)  </li>\n",
    "        <li> [2.3 Train Model on the Remote Compute](#2.3-Prepare-your-training-script) </li> \n",
    "        <li> [2.4 Create a PyTorch estimator for fine tuning](#2.4-Create-a-PyTorch-estimator-for-fine-tuning) </li> \n",
    "        <li> [2.5 Create an experiment](#2.5-Create-an-experiment)  </li>\n",
    "        <li> [2.6 Submit and Monitor your run](#2.6-Submit-and-Monitor-your-run)  </li>\n",
    "      </ul>\n",
    "    </li><br>\n",
    "    <li>[Fine Tuning BERT with Hyperparameter Tuning](#3-Fine-Tuning-BERT-with-Hyperparameter-Tuning)\n",
    "        <ul style=\"list-style: none;\">\n",
    "            <li> [3.1 Start a hyperparameter sweep](#3.1-Start-a-hyperparameter-sweep)</li>\n",
    "            <li> [3.2 Monitor HyperDrive runs](#3.2-Monitor-HyperDrive-runs)</li>\n",
    "            <li> [3.3 Find and register the best model](#3.3-Find-and-register-the-best-model)</li>\n",
    "        </ul>\n",
    "    </li><br>\n",
    "    <li>[References](#References)</li>\n",
    "</ol>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Prerequisites: Global Settings\n",
    "You will need to do the following to be successful with the rest of the notebook:\n",
    "- Have an existing Azure subscription. You can get started for free [here](https://azure.microsoft.com/free/)\n",
    "- Understand the [architecture and terms](https://docs.microsoft.com/azure/machine-learning/service/concept-azure-machine-learning-architecture) introduced by Azure Machine Learning service (AzureML)\n",
    "\n",
    "- Make sure the [AzureML Python SDK](https://pypi.org/project/azureml-sdk/) is installed with notebooks and contrib add ons.\n",
    "```\n",
    "conda create -n azureml -y Python=3.6\n",
    "source activate azureml\n",
    "pip install --upgrade azureml-sdk[notebooks,contrib] \n",
    "conda install ipywidgets\n",
    "jupyter nbextension install --py --user azureml.widgets\n",
    "jupyter nbextension enable azureml.widgets --user --py\n",
    "```\n",
    "- Import the required packages\n",
    "- Set Environment Variables\n",
    "- Connect to an [Azure Machine Learning service workspace](https://docs.microsoft.com/en-us/azure/machine-learning/service/how-to-manage-workspace#create-a-workspace)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the following cell to make sure you have installed all the packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../../\")\n",
    "import urllib, os\n",
    "from utils_nlp.azureml import azureml_utils\n",
    "import math\n",
    "import json \n",
    "import pandas as pd\n",
    "import papermill as pm\n",
    "#package for flattening json in pandas df\n",
    "from pandas.io.json import json_normalize\n",
    "import shutil\n",
    "import scrapbook as sb\n",
    "# Check core SDK version number\n",
    "import azureml.core\n",
    "from azureml.core import Datastore\n",
    "from azureml.core import Experiment\n",
    "from azureml.core.compute import ComputeTarget, AmlCompute\n",
    "from azureml.core.compute_target import ComputeTargetException\n",
    "from azureml.core.runconfig import MpiConfiguration\n",
    "from azureml.telemetry import set_diagnostics_collection\n",
    "from azureml.train.dnn import PyTorch\n",
    "from azureml.train.hyperdrive import *\n",
    "from azureml.widgets import RunDetails\n",
    "\n",
    "print(\"SDK version:\", azureml.core.VERSION)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following parameters are set as variables to be used throughout the notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Model configuration\n",
    "DATA_FOLDER = './squad'\n",
    "PROJECT_FOLDER = './transformers'\n",
    "EXPERIMENT_NAME = 'NLP-QA-BERT-deepdive'\n",
    "BERT_MODEL = 'bert-large-uncased'\n",
    "TARGET_GRADIENT_STEPS = 16\n",
    "INIT_GRADIENT_STEPS = 2\n",
    "MAX_SEQ_LENGTH = 384\n",
    "NUM_TRAIN_EPOCHS = 2.0\n",
    "NODE_COUNT = 2\n",
    "TRAIN_SCRIPT_PATH = 'bert_run_squad_azureml.py'\n",
    "MAX_TOTAL_RUNS = 8\n",
    "MAX_CONCURRENT_RUNS = 4\n",
    "BERT_UTIL_PATH = '../../utils_nlp/azureml/azureml_bert_util.py'\n",
    "EVALUATE_SQAD_PATH = '../../utils_nlp/eval/evaluate_squad.py'\n",
    "\n",
    "# Azure resources\n",
    "subscription_id = \"YOUR_SUBSCRIPTION_ID\"\n",
    "resource_group = \"YOUR_RESOURCE_GROUP_NAME\"  \n",
    "workspace_name = \"YOUR_WORKSPACE_NAME\"  \n",
    "workspace_region = \"YOUR_WORKSPACE_REGION\" #Possible values eastus, eastus2 and so on.\n",
    "AZUREML_CONFIG_PATH = \"./.azureml\"\n",
    "AZUREML_VERBOSE = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Initialize workspace**\n",
    "\n",
    "The following cell looks to set up the connection to your [Azure Machine Learning service Workspace](https://docs.microsoft.com/azure/machine-learning/service/concept-azure-machine-learning-architecture#workspace). You can choose to connect to an existing workspace or create a new one. \n",
    "\n",
    "**To access an existing workspace:**\n",
    "1. If you have a `config.json` file, you do not need to provide the workspace information; you will only need to update the `config_path` variable that is defined above which contains the file.\n",
    "2. Otherwise, you will need to supply the following:\n",
    "    * The name of your workspace\n",
    "    * Your subscription id\n",
    "    * The resource group name\n",
    "\n",
    "**To create a new workspace:**\n",
    "\n",
    "Set the following information:\n",
    "* A name for your workspace\n",
    "* Your subscription id\n",
    "* The resource group name\n",
    "* [Azure region](https://azure.microsoft.com/en-us/global-infrastructure/regions/) to create the workspace in, such as `eastus2`. \n",
    "\n",
    "This will automatically create a new resource group for you in the region provided if a resource group with the name given does not already exist. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(AZUREML_CONFIG_PATH):\n",
    "    ws = azureml_utils.get_or_create_workspace(config_path=AZUREML_CONFIG_PATH)\n",
    "else:\n",
    "    ws = azureml_utils.get_or_create_workspace(\n",
    "        subscription_id=subscription_id,\n",
    "        resource_group=resource_group,\n",
    "        workspace_name=workspace_name,\n",
    "        workspace_region=workspace_region,\n",
    "    )\n",
    "\n",
    "if AZUREML_VERBOSE:\n",
    "    print('Workspace name: ' + ws.name, \n",
    "          'Azure region: ' + ws.location, \n",
    "          'Subscription id: ' + ws.subscription_id, \n",
    "          'Resource group: ' + ws.resource_group, sep='\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Diagnostics**\n",
    "\n",
    "Opt-in diagnostics for better experience, quality, and security of future releases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Turning diagnostics collection on. \n"
     ]
    }
   ],
   "source": [
    "set_diagnostics_collection(send_diagnostics=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Data Loading"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we will\n",
    "1. Connect to the default AzureML datastore\n",
    "2. Download and load the dataset\n",
    "3. Upload the training set to the default blob storage of the workspace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_folder = DATA_FOLDER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Default AzureML datastore\n",
    "\n",
    "To make data accessible for remote training, AzureML provides a convenient way to do so via a [Datastore](https://docs.microsoft.com/azure/machine-learning/service/how-to-access-data). The datastore provides a mechanism for you to upload/download data to Azure Storage, and interact with it from your remote compute targets.\n",
    "\n",
    "Each workspace is associated with a default Azure Blob datastore named `'workspaceblobstore'`. We use this default datastore to collect the SQuAD training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AzureBlob maidaipberteas6144514557 azureml-blobstore-cf97de17-8d21-437f-8b4c-298560f34ecd $AZUREML_DATAREFERENCE_workspaceblobstore\n"
     ]
    }
   ],
   "source": [
    "ds = ws.get_default_datastore()\n",
    "if AZUREML_VERBOSE:\n",
    "    print(ds.datastore_type, ds.account_name, ds.container_name, ds.as_mount())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Download training dataset - SQuAD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The SQuAD dataset can be downloaded with the following links and should be saved in a blob storage.\n",
    "- [train-v1.1.json](https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json)\n",
    "- [dev-v1.1.json](https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('./squad\\\\dev-v1.1.json', <http.client.HTTPMessage at 0x1569b645f28>)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.makedirs('./squad', exist_ok=True)\n",
    "urllib.request.urlretrieve('https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json', filename=os.path.join(data_folder, 'train-v1.1.json'))\n",
    "urllib.request.urlretrieve('https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json', filename=os.path.join(data_folder, 'dev-v1.1.json'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The SQuAD dataset contains question-answer pairs on 500+ articles. For each observation in the training set, we have a **context, question, and text**. An example shown as below: [source](https://towardsdatascience.com/building-a-question-answering-system-part-1-9388aadff507)\n",
    "<img src=\"https://nlpbp.blob.core.windows.net/images/squad.png\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The structure of an example in the training data.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>context</th>\n",
       "      <th>qas</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>[{'answers': [{'answer_start': 515, 'text': 'S...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                             context  \\\n",
       "0  Architecturally, the school has a Catholic cha...   \n",
       "\n",
       "                                                 qas  \n",
       "0  [{'answers': [{'answer_start': 515, 'text': 'S...  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#load json object\n",
    "with open(os.path.join(data_folder, 'train-v1.1.json')) as train_file:\n",
    "    train = json.load(train_file)\n",
    "for paragraph in train['data'][0]['paragraphs']:\n",
    "    answer_question = paragraph['qas']\n",
    "    context = paragraph['context']\n",
    "    paragraph = paragraph\n",
    "    break\n",
    "print(\"The structure of an example in the training data.\")\n",
    "json_normalize(paragraph).head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The structure of question answer pairs in the above example.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>answers</th>\n",
       "      <th>id</th>\n",
       "      <th>question</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[{'answer_start': 515, 'text': 'Saint Bernadet...</td>\n",
       "      <td>5733be284776f41900661182</td>\n",
       "      <td>To whom did the Virgin Mary allegedly appear i...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[{'answer_start': 188, 'text': 'a copper statu...</td>\n",
       "      <td>5733be284776f4190066117f</td>\n",
       "      <td>What is in front of the Notre Dame Main Building?</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>[{'answer_start': 279, 'text': 'the Main Build...</td>\n",
       "      <td>5733be284776f41900661180</td>\n",
       "      <td>The Basilica of the Sacred heart at Notre Dame...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>[{'answer_start': 381, 'text': 'a Marian place...</td>\n",
       "      <td>5733be284776f41900661181</td>\n",
       "      <td>What is the Grotto at Notre Dame?</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>[{'answer_start': 92, 'text': 'a golden statue...</td>\n",
       "      <td>5733be284776f4190066117e</td>\n",
       "      <td>What sits on top of the Main Building at Notre...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                             answers  \\\n",
       "0  [{'answer_start': 515, 'text': 'Saint Bernadet...   \n",
       "1  [{'answer_start': 188, 'text': 'a copper statu...   \n",
       "2  [{'answer_start': 279, 'text': 'the Main Build...   \n",
       "3  [{'answer_start': 381, 'text': 'a Marian place...   \n",
       "4  [{'answer_start': 92, 'text': 'a golden statue...   \n",
       "\n",
       "                         id                                           question  \n",
       "0  5733be284776f41900661182  To whom did the Virgin Mary allegedly appear i...  \n",
       "1  5733be284776f4190066117f  What is in front of the Notre Dame Main Building?  \n",
       "2  5733be284776f41900661180  The Basilica of the Sacred heart at Notre Dame...  \n",
       "3  5733be284776f41900661181                  What is the Grotto at Notre Dame?  \n",
       "4  5733be284776f4190066117e  What sits on top of the Main Building at Notre...  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"The structure of question answer pairs in the above example.\")\n",
    "json_normalize(answer_question).head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Upload to Azure blob storage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code will upload the SQuAD dataset to the path ./squad on the default datastore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Uploading an estimated of 2 files\n",
      "Target already exists. Skipping upload for squad\\dev-v1.1.json\n",
      "Target already exists. Skipping upload for squad\\train-v1.1.json\n",
      "Uploaded 0 files\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "$AZUREML_DATAREFERENCE_972d18f476b34d26a1ffd6a11b473114"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds.upload(src_dir='./squad', target_path='./squad', show_progress=AZUREML_VERBOSE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Fine-tuning BERT with Distributed Training by Horovod\n",
    "We can reference the datastore to access the SQuAD dataset and start fine-tuning the model by exploring the power of distributed training on AzureML GPU clusters.\n",
    "\n",
    "Once you've created your workspace and set up your development environment, training a model in Azure Machine Learning involves the following steps:\n",
    "1. Create a GPU remote compute target\n",
    "2. Create a project directory\n",
    "3. Prepare your training script\n",
    "4. Create an Estimator object\n",
    "5. Submit the estimator to an experiment object under the workspace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Create a GPU remote compute target\n",
    "\n",
    "We need to create a GPU [compute target](https://docs.microsoft.com/azure/machine-learning/service/concept-azure-machine-learning-architecture#compute-target) to perform the fine-tuning. In this example, we create an AmlCompute cluster as our training compute resource. Please find the information of Azure VM size in below table.\n",
    "\n",
    "    \n",
    "|    VM Size    \t| CPU \t|   GPU   \t| Storage (SSD) \t| GPU memory \t| InfiniBand  \t|\n",
    "|:-------------:\t|:---:\t|:-------:\t|:-------------:\t|:----------:\t|:----------:\t|\n",
    "|  Standard_NC6 \t|  6  \t| 1 x K80 \t|    340 GiB    \t|    8 GiB   \t|      No   \t|\n",
    "| Standard_NC12 \t|  12 \t| 2 x K80 \t|    680 GiB    \t|   16 GiB   \t|      No   \t|\n",
    "| Standard_NC24 \t|  24 \t| 4 x K80 \t|    1440 GiB   \t|   32 GiB   \t|      No   \t|\n",
    "| Standard_NC24r \t|  24 \t| 4 x K80 \t|    1440 GiB   \t|   32 GiB   \t|      Yes   \t|\n",
    "| Standard_NC6s_v3 \t|  6  \t| 1 x V100 \t|    736 GiB    \t|   16 GiB   \t|      No   \t|\n",
    "| Standard_NC12s_v3 |  12 \t| 2 x V100 \t|    1474 GiB   \t|   32 GiB   \t|      No   \t|\n",
    "| Standard_NC24s_v3 |  24 \t| 4 x V100 \t|    2948 GiB   \t|   64 GiB   \t|      No   \t|\n",
    "| Standard_NC24rs_v3|  24 \t| 4 x V100 \t|    2948 GiB   \t|   64 GiB   \t|      Yes   \t|\n",
    "\n",
    "This code creates a cluster for you if it does not already exist in your workspace.\n",
    "\n",
    "***We strongly recommend to use NCv3-series (NVIDIA Tesla V100) to fine-tune with SQuAD dataset. You will need to request quota of NCv3-series for your AzureML subscription.***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found existing compute target.\n",
      "{'currentNodeCount': 0, 'targetNodeCount': 0, 'nodeStateCounts': {'preparingNodeCount': 0, 'runningNodeCount': 0, 'idleNodeCount': 0, 'unusableNodeCount': 0, 'leavingNodeCount': 0, 'preemptedNodeCount': 0}, 'allocationState': 'Steady', 'allocationStateTransitionTime': '2019-07-22T22:38:04.496000+00:00', 'errors': None, 'creationTime': '2019-07-12T19:59:45.933132+00:00', 'modifiedTime': '2019-07-12T20:00:01.793458+00:00', 'provisioningState': 'Succeeded', 'provisioningStateTransitionTime': None, 'scaleSettings': {'minNodeCount': 0, 'maxNodeCount': 4, 'nodeIdleTimeBeforeScaleDown': 'PT120S'}, 'vmPriority': 'Dedicated', 'vmSize': 'STANDARD_NC24RS_V3'}\n"
     ]
    }
   ],
   "source": [
    "# choose a name for your cluster\n",
    "cluster_name = \"bertncrs24\"\n",
    "\n",
    "try:\n",
    "    gpu_compute_target = ComputeTarget(workspace=ws, name=cluster_name)\n",
    "    print('Found existing compute target.')\n",
    "except ComputeTargetException:\n",
    "    print('Creating a new compute target...')\n",
    "    compute_config = AmlCompute.provisioning_configuration(vm_size='STANDARD_NC24rs_v3',\n",
    "                                                           max_nodes=4)\n",
    "\n",
    "    # create the cluster\n",
    "    gpu_compute_target = ComputeTarget.create(ws, cluster_name, compute_config)\n",
    "\n",
    "    gpu_compute_target.wait_for_completion(show_output=True)\n",
    "\n",
    "# use get_status() to get a detailed status for the current AmlCompute. \n",
    "print(gpu_compute_target.get_status().serialize())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Create a project directory\n",
    "Create a directory that will contain all the necessary code from your local machine that you will need access to on the remote resource. This includes the training script and any additional files your training script depends on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "project_folder = PROJECT_FOLDER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make a local clone of the original [PyTorch reimplementation](https://github.com/huggingface/pytorch-pretrained-BERT) repository"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "!git clone -b v0.4.0 https://github.com/huggingface/transformers.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Prepare your training script\n",
    "Let us prepare the training script to run the fine-tuning script `run_squad.py` from [the Hugging Face repository](https://github.com/huggingface/pytorch-pretrained-BERT/blob/master/examples/run_squad.py). Please refer to the [repo](https://github.com/huggingface/pytorch-pretrained-BERT#fine-tuning-with-bert-running-the-examples) for more details about the script. \n",
    "\n",
    "The original `run_squad.py` script uses the PyTorch distributed launch utility to launch multiple processes across nodes and GPUs. Here we use a [modified version](https://github.com/microsoft/AzureML-BERT/blob/master/finetune/run_squad_azureml.py) of this file provided by the [AzureML-BERT repo](https://github.com/microsoft/AzureML-BERT) from Microsoft to be able to launch multiple processes across nodes and GPUs in with an AzureML built-in MPI backend."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's retrieve and copy the training script [bert_run_squad_azureml.py](.\\bert_run_squad_azureml.py), evaluation script for SQuAD v1.1 [evaluate-v1.1.py](../../utils_nlp/eval/evaluate_squad.py) and the helper utility script for Horovod [azureml_bert_util.py](../../utils_nlp/azureml/azureml_bert_util.py) into our project directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "shutil.copy(EVALUATE_SQAD_PATH, project_folder)\n",
    "shutil.copy(BERT_UTIL_PATH, project_folder)\n",
    "shutil.copy(TRAIN_SCRIPT_PATH, project_folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Create a PyTorch estimator for fine-tuning\n",
    "Let us create a new PyTorch estimator to run the fine-tuning script `run_squad_azureml.py`. To use AzureML's tracking and metrics capabilities, we need to add a small amount of AzureML code inside the training script.\n",
    "\n",
    "In `run_squad_azureml.py`, we will log some metrics to our AzureML run. To do so, we will access the AzureML run object within the script:\n",
    "```Python\n",
    "from azureml.core.run import Run\n",
    "run = Run.get_context()\n",
    "```\n",
    "Further within `run_squad_azureml.py`, we log learning rate, training loss and prediction scores the model achieves as:\n",
    "```Python\n",
    "run.log('lr', np.float(args.learning_rate))\n",
    "...\n",
    "\n",
    "for step, batch in enumerate(tqdm(train_dataloader, desc=\"Iteration\")): \n",
    "    ...\n",
    "    run.log('train_loss', np.float(loss))\n",
    "\n",
    "..\n",
    "```\n",
    "These run metrics will become particularly important when we begin hyperparameter tuning our model in the \"Tune model hyperparameters\" section below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, AzureML PyTorch estimator can be defined as below. We use `azuremlsamples/bert:torch-1.0.0-apex-cuda9` as the base docker image with [dockerfile](./dockerfile). In this example, we use STANDARD_NC24rs_v3 which has 4 GPUs. Thus, we can set `process_count_per_node=4`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "mpiConfig=MpiConfiguration()\n",
    "mpiConfig.process_count_per_node=4\n",
    "\n",
    "estimator = PyTorch(source_directory=project_folder,\n",
    "                    compute_target=gpu_compute_target,\n",
    "                    script_params = {\n",
    "                          '--bert_model':BERT_MODEL,\n",
    "                          '--do_train' : '',\n",
    "                          '--do_predict': '',\n",
    "                          '--train_file': ds.path('squad/train-v1.1.json').as_mount(),\n",
    "                          '--predict_file': ds.path('squad/dev-v1.1.json').as_mount(),\n",
    "                          '--max_seq_length': MAX_SEQ_LENGTH,\n",
    "                          '--train_batch_size': 8,\n",
    "                          '--learning_rate': 6.8e-5,\n",
    "                          '--num_train_epochs': NUM_TRAIN_EPOCHS,\n",
    "                          '--doc_stride': 128,\n",
    "                          '--seed': 32,\n",
    "                          '--init_gradient_accumulation_steps':INIT_GRADIENT_STEPS,\n",
    "                          '--target_gradient_accumulation_steps':TARGET_GRADIENT_STEPS,\n",
    "                          '--accumulation_warmup_proportion':0.25,\n",
    "                          '--output_dir': './outputs',\n",
    "                          '--loss_scale':256,\n",
    "                    },\n",
    "                    custom_docker_image='azuremlsamples/bert:torch-1.0.0-apex-cuda9',\n",
    "                    entry_script='bert_run_squad_azureml.py',\n",
    "                    node_count=NODE_COUNT,\n",
    "                    distributed_training=mpiConfig,\n",
    "                    framework_version='1.1',\n",
    "                    use_gpu=True)\n",
    "estimator._estimator_config.environment.python.user_managed_dependencies=True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note: You can try with `--bert_model:'bert-base-uncased'`to run a smaller bert model faster.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.5 Create an experiment\n",
    "Create an [Experiment](https://docs.microsoft.com/azure/machine-learning/service/concept-azure-machine-learning-architecture#experiment) to track all the runs in your workspace for this distributed PyTorch tutorial. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_name = EXPERIMENT_NAME\n",
    "experiment = Experiment(ws, name=experiment_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.6 Submit and Monitor your run\n",
    "AzureML can automatically create figures on the loss and time, which is eaiser to track the performance as in the following figure shown the train loss v.s. the number of iterations:\n",
    "![train_loss_bert](https://nlpbp.blob.core.windows.net/images/train_loss_bert.PNG)\n",
    "\n",
    "The Jupyter widget would be like this:\n",
    "![train_loss_bert](https://nlpbp.blob.core.windows.net/images/bert_widget.PNG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "edeffbf4c94047f48699af91c45a27b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "_UserRunWidget(widget_settings={'childWidgetDisplay': 'popup', 'send_telemetry': True, 'log_level': 'INFO', 's…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "run = experiment.submit(estimator)\n",
    "RunDetails(run).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = run.wait_for_completion(show_output=AZUREML_VERBOSE) # Block until complete"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Cancel the job**\n",
    "\n",
    "You can cancel the job manually to make sure you do not waste resources.\n",
    " ```python\n",
    "# Cancel the job with id.\n",
    "job_id = \"BERT-SQuAD_1562612876_bab5b3af\"\n",
    "run = get_run(experiment, job_id)\n",
    "\n",
    "# Cancel jobs.\n",
    "run.cancel()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To achieve an F1 score over **90.5** F1 score and an Exact-Match over **83.5** with the `SQuAD v1.1` dataset, it requires **2** epochs when fine-tuning with the `BERT large` model. Below please find the elapsed time using deferent Azure GPU VMs and configures. \n",
    "\n",
    "The default configuration in this notebook uses 2 `STANDARD_NC24rs_v3` (8 x V100) with `fp16` enabled. The training phase should take **22 mins** to complete 2 epochs. \n",
    "\n",
    "|  GPU counts \t|    1 GPU    \t|         2 GPU \t| 4 GPU      \t| 8 GPU      \t|\n",
    "|------------:\t|:-----------:\t|--------------:\t|------------\t|------------\t|\n",
    "| NCv3-series \t|     340 mins  |    180 mins \t    |    80 mins \t|   48 mins \t|\n",
    "| NCv3 with fp16|     140 mins  |    79 mins \t    |    38 mins \t|   22 mins \t|\n",
    "\n",
    "The performance with different VMs with `fp16` enabled (Duration = training time + preparing time):\n",
    "\n",
    "|  VM Size \t|  GPU counts|    Node counts|    Duration    \t|         F1 \t| EM     \t| Pretrained BERT      \t|\n",
    "|------------:\t|:-----------:\t|--------------:\t|------------\t|------------\t|------------\t|------------\t|\n",
    "| NC6sv3 |   4 |  4 |  31 mins  |    88.24 \t    |    80.59 \t|   Base \t|\n",
    "| NC6sv3 |   4 |  4 |  80 mins  |    90.78 \t    |    83.96 \t|   Large \t|\n",
    "| NC24rsv3 |  4 |   1 |  19 mins  |    86.18 \t    |    77.90 \t|   Base \t|\n",
    "| NC24rsv3 |  4 |   1 |  46 mins  |    90.53 \t    |    83.56 \t|   Large \t|\n",
    "| NC24rsv3 |  8 |   2 |  19 mins  |    87.47 \t    |    79.52 \t|   Base \t|\n",
    "| NC24rsv3 |  8 |   2 |  32 mins  |    90.57 \t    |    83.58 \t|   Large \t|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 Fine-Tuning BERT with Hyperparameter Tuning\n",
    "\n",
    "We would also like to optimize our hyperparameter, `learning rate`, using Azure Machine Learning's hyperparameter tuning capabilities."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Start a hyperparameter sweep\n",
    "First, we will define the hyperparameter space to sweep over. In this example we will use random sampling to try different configuration sets of hyperparameter to minimize our primary metric, the f1 score (`f1`). For simplicity, we tune the BERT base model with  `--bert_model':'bert-base-uncased` and  `node_count=1`.\n",
    "\n",
    "We can also try with `BayesianParameterSampling` with suggested `max_total_runs=20`.\n",
    "```Python\n",
    "param_sampling = BayesianParameterSampling( {\n",
    "         'learning_rate': uniform(5e-5, 9e-5),\n",
    "    }\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_sampling = RandomParameterSampling( {\n",
    "         'learning_rate': uniform(5e-5, 9e-5),\n",
    "    }\n",
    ")\n",
    "hyperdrive_config = HyperDriveConfig(estimator=estimator,\n",
    "                                         hyperparameter_sampling=param_sampling, \n",
    "                                         primary_metric_name='f1',\n",
    "                                         primary_metric_goal=PrimaryMetricGoal.MAXIMIZE,\n",
    "                                         max_total_runs=MAX_TOTAL_RUNS,\n",
    "                                         max_concurrent_runs=MAX_CONCURRENT_RUNS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, lauch the hyperparameter tuning job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start the HyperDrive run\n",
    "hyperdrive_run = experiment.submit(hyperdrive_config)\n",
    "RunDetails(hyperdrive_run).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Monitor HyperDrive runs\n",
    "We can monitor the progress of the runs with the following Jupyter widget. \n",
    "![](https://nlpbp.blob.core.windows.net/images/bert_tune.PNG)\n",
    "![](https://nlpbp.blob.core.windows.net/images/bert_tune2.PNG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = hyperdrive_run.wait_for_completion(show_output=AZUREML_VERBOSE) # Block until complete"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see the experiment progress from this notebook by using `azureml.widgets.RunDetails(hd_run).show()` or check from the Azure portal with the url link you can get by running `hd_run.get_portal_url()`.\n",
    "To load an existing Hyperdrive run, use `hd_run = hd.HyperDriveRun(exp, <user-run-id>, hyperdrive_run_config=hd_run_config)`. You also can cancel a run with `hd_run.cancel()`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Cancel the hyper drive run to save the resources**\n",
    " ```python\n",
    "# Cancel the hyper drive\n",
    "hyperdrive_run.cancel()\n",
    " ```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Find and register the best model\n",
    "Once all the runs complete, we can find the run that produced the model with the highest F1 score. The F1 score with default learning rate is **86.18** in [Submit and Monitor your run](#2.6-Submit-and-Monitor-your-run) . The best F1 score is **87.01** after tuning with `learning rate=0.000090` with random sampling. With Bayesian sampling, the best F1 score is **86.87** after tuning with `learning rate=0.0000896`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run(Experiment: BERT-SQuAD,\n",
      "Id: BERT-SQuAD_1562966635446_2,\n",
      "Type: azureml.scriptrun,\n",
      "Status: Completed)\n",
      "Best Run is:\n",
      "  F1 score: 87.01 \n",
      "  Learning rate: 0.000090\n"
     ]
    }
   ],
   "source": [
    "best_run = hyperdrive_run.get_best_run_by_primary_metric()\n",
    "best_run_metrics = best_run.get_metrics()\n",
    "print('Best Run is:\\n  F1 score: %.2f \\n  Learning rate: %f' % (best_run_metrics['f1'], best_run_metrics['lr']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Persist properties of the run so we can access the logged metrics later\n",
    "sb.glue(\"f1\", best_run_metrics['f1'])\n",
    "sb.glue(\"learning_rate\", best_run_metrics['lr'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "1. Devlin, Jacob and Chang, Ming-Wei and Lee, Kenton and Toutanova, Kristina, [*BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding*](https://arxiv.org/abs/1810.04805), ACL, 2018.\n",
    "2. Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, Percy Liang, [*SQuAD: 100,000+ Questions for Machine Comprehension of Text*](https://arxiv.org/abs/1606.05250), EMNLP, 2016. Dataset available at [https://rajpurkar.github.io/SQuAD-explorer/](https://rajpurkar.github.io/SQuAD-explorer/)."
   ]
  }
 ],
 "metadata": {
  "authors": [
   {
    "name": "minxia"
   }
  ],
  "celltoolbar": "Tags",
  "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.5"
  },
  "msauthor": "minxia"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}