pretrain/PyTorch/notebooks/BERT_Pretrain.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pretraining of the BERT model\n",
    "\n",
    "This notebook contains an end-to-end walkthrough of using Azure Machine Learning service and pretraining [BERT: Bidirectional Encoder Representations from Transformers](https://arxiv.org/abs/1810.04805) models.\n",
    "\n",
    "Methodology:\n",
    "- Intialize an AzureML workspace\n",
    "- Register a datastore\n",
    "- Create an experiment\n",
    "- Provision a compute target\n",
    "- Create an Estimator\n",
    "- Configure and Run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "If you are using an [Azure Machine Learning Notebook VM](https://docs.microsoft.com/en-us/azure/machine-learning/service/quickstart-run-cloud-notebook), you are all set. Otherwise, refer to the [configuration Notebook](https://github.com/Azure/MachineLearningNotebooks/blob/56e0ebc5acb9614fac51d8b98ede5acee8003820/configuration.ipynb) first if you haven't already to establish your connection to the AzureML Workspace. Prerequisites are:\n",
    "* Azure subscription\n",
    "* Azure Machine Learning Workspace\n",
    "* Azure Machine Learning SDK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Library import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Regular python libraries\n",
    "import os\n",
    "import requests\n",
    "import sys\n",
    "\n",
    "# AzureML libraries\n",
    "import azureml.core\n",
    "from azureml.core import Experiment, Workspace, Datastore, Run\n",
    "from azureml.core.compute import ComputeTarget, AmlCompute\n",
    "from azureml.core.compute_target import ComputeTargetException\n",
    "from azureml.core.conda_dependencies import CondaDependencies\n",
    "from azureml.core.container_registry import ContainerRegistry\n",
    "from azureml.core.runconfig import MpiConfiguration, RunConfiguration, DEFAULT_GPU_IMAGE\n",
    "from azureml.train.dnn import PyTorch\n",
    "from azureml.train.estimator import Estimator\n",
    "from azureml.widgets import RunDetails\n",
    "\n",
    "# Check core SDK version number\n",
    "print(\"SDK version:\", azureml.core.VERSION)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Workspace setup\n",
    "\n",
    "Initialize a Workspace object from the existing workspace you created in the Prerequisites step or create a new one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve the workspace\n",
    "ws = Workspace.setup()\n",
    "\n",
    "# Print the workspace attributes\n",
    "print('Workspace name: ' + ws.name, \n",
    "      'Workspace region: ' + ws.location, \n",
    "      'Subscription id: ' + ws.subscription_id, \n",
    "      'Resource group: ' + ws.resource_group, sep = '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Datastore registration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[BERT paper](https://arxiv.org/pdf/1810.04805) references `Wikipedia` and `BookCorpus` datasets for pretraining. This notebook is configured to use Wikipedia dataset only, but can be used with other datasets as well, including custom datasets. The preprocessed data should be available in a `Datastore` in AzureML `Workspace`. \n",
    "\n",
    "The Wikipedia corpus used for BERT pretraining is preprocessed following the [data prep instructions](https://github.com/microsoft/AzureML-BERT/blob/master/docs/dataprep.md) and uploaded to  https://bertonazuremlwestus2.blob.core.windows.net/public2/bert_data.tar.gz (70 GB). You need to extract the files and copy them to another Azure blob container and register it as a workspace to use it in the pretraining job. Additional details on the tar.gz file and the data transfer are available at [artifacts.md](https://github.com/microsoft/AzureML-BERT/blob/master/docs/artifacts.md).\n",
    "\n",
    "Alternatively, you can preprocess the raw data from scratch (instructions available at the [data prep notes](https://github.com/microsoft/AzureML-BERT/blob/master/docs/dataprep.md)), upload that to an Azure blob container and use it as the datastore for the job. \n",
    "\n",
    "Note: it is also possible to use datasets other than Wikipedia corpus with this implementation. \n",
    "\n",
    "The following code assumes that the data is already copied to an Azure blob container with the following directory structure. It is recommended to retain this directory structure to run this notebook without code updates. In case the directory structure is different, the constructor of PyTorch estimator where the datastore is mounted should be modified.\n",
    "\n",
    "       \n",
    "```\n",
    "bert_data\n",
    "│   bert-base.json\n",
    "│   bert-large.json\n",
    "│   bert-base-single-node.json\n",
    "│   bert-large-single-node.json\n",
    "│\n",
    "└───512\n",
    "│   │\n",
    "│   └───wiki_pretrain\n",
    "│       │   wikipedia_segmented_part_0.bin\n",
    "│       │   wikipedia_segmented_part_1.bin\n",
    "│       │   ...\n",
    "│       │   wikipedia_segmented_part_98.bin\n",
    "└───validation_512_only\n",
    "    │\n",
    "    └───validation_set.bin\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Register the datastore with the workspace\n",
    "ds = Datastore.register_azure_blob_container(workspace=ws, \n",
    "                                             datastore_name='BERT_Preprocessed_Data',\n",
    "                                             container_name='data',\n",
    "                                             account_name='<name goes here>', \n",
    "                                             account_key='<key goes here>'\n",
    "                                            )\n",
    "\n",
    "# Help from: https://docs.microsoft.com/en-us/azure/machine-learning/service/how-to-access-data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print the workspace attributes\n",
    "print('Datastore name: ' + ds.name, \n",
    "      'Container name: ' + ds.container_name, \n",
    "      'Datastore type: ' + ds.datastore_type, \n",
    "      'Workspace name: ' + ds.workspace.name, sep = '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an Experiment\n",
    "\n",
    "Experiment is a logical container in an Azure ML Workspace. It hosts run records which can include run metrics and output artifacts from your experiments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an experiment\n",
    "experiment_name = 'BERT-pretraining'\n",
    "experiment = Experiment(ws, name=experiment_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Provision a cluster\n",
    "\n",
    "### Introduction to AmlCompute\n",
    "\n",
    "Azure Machine Learning Compute is managed compute infrastructure that allows the user to easily create single to multi-node compute of the appropriate VM Family. It is created within your workspace region and is a resource that can be used by other users in your workspace. It autoscales by default to the max_nodes, when a job is submitted, and executes in a containerized environment packaging the dependencies as specified by the user.\n",
    "\n",
    "Since it is managed compute, job scheduling and cluster management are handled internally by Azure Machine Learning service.\n",
    "\n",
    "For more information on Azure Machine Learning Compute, please read [this](https://docs.microsoft.com/en-us/azure/machine-learning/service/how-to-set-up-training-targets#amlcompute).\n",
    "\n",
    "Note: As with other Azure services, there are limits on certain resources (for eg. AmlCompute quota) associated with the Azure Machine Learning service. Please read [this](https://docs.microsoft.com/en-us/azure/machine-learning/service/how-to-manage-quotas) on the default limits and how to request more quota."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a compute target\n",
    "BERT pretraining on Azure Machine Learning Service is supported on 16 x `Standard_NC24s_v3` or 8 x `Standard_ND40_v2` VMs. In the next step, you will create a 16 node (i.e. 64 GPUs) AMLCompute cluster of `Standard_NC24s_v3` GPU VMs, if it doesn't already exist in your workspace. The code to create a cluster with 8 `Standard_ND40_v2` VMs is commented out in the cell below.\n",
    "\n",
    "* vm_size: VM family of the nodes provisioned by AmlCompute. Simply choose from the supported_vmsizes() above\n",
    "* max_nodes: Maximum nodes to autoscale to while running a job on AmlCompute\n",
    "* min_nodes: Minimum number of nodes while running a job on AmlCompute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the compute cluster\n",
    "gpu_cluster_name = \"pretraincluster\" \n",
    "\n",
    "# Verify that the cluster doesn't exist already\n",
    "try:\n",
    "    gpu_compute_target = ComputeTarget(workspace=ws, name=gpu_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_NC24s_v3', min_nodes=0, max_nodes=16)\n",
    "    # compute_config = AmlCompute.provisioning_configuration(vm_size='Standard_NC40_v2', min_nodes=0, max_nodes=8)\n",
    "    \n",
    "    # create the cluster\n",
    "    gpu_compute_target = ComputeTarget.create(ws, gpu_cluster_name, compute_config)\n",
    "    gpu_compute_target.wait_for_completion(show_output=True)\n",
    "\n",
    "# Use the 'status' property to get a detailed status for the current cluster. \n",
    "#print(gpu_compute_target.status.serialize())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Estimator definition and run submission\n",
    "\n",
    "The estimator uses a custom docker image and train.py as the entry script for execution.\n",
    "\n",
    "For more information on Estimator, refer [here](https://docs.microsoft.com/en-us/azure/machine-learning/service/how-to-train-pytorch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the project folder\n",
    "project_folder = '..' # This is to allow the libraries stored under pytorch/ to be loaded\n",
    "\n",
    "## Using a public image published on Azure.\n",
    "image_name = 'mcr.microsoft.com/azureml/bert:pretrain-openmpi3.1.2-cuda10.0-cudnn7-ubuntu16.04'\n",
    "\n",
    "# Using MPI to execute a distributed run\n",
    "mpi = MpiConfiguration()\n",
    "# Standard_NC24s_v3 VM has 4 GPUs. !!!! update this appropriately if you use a different VM size !!!!\n",
    "mpi.process_count_per_node = 4 \n",
    "# !!!! use the following for Standard_NC40_v2 VM !!!!\n",
    "# mpi.process_count_per_node = 8\n",
    "\n",
    "# Define the Pytorch estimator\n",
    "estimator = PyTorch(source_directory=project_folder,\n",
    "                    # Compute configuration\n",
    "                    compute_target=gpu_compute_target,\n",
    "                    node_count=16, \n",
    "                    distributed_training=mpi,\n",
    "                    use_gpu=True,\n",
    "                    \n",
    "                    #Docker image\n",
    "                    use_docker=True,\n",
    "                    custom_docker_image=image_name,\n",
    "                    user_managed=True,\n",
    "                    \n",
    "                    # Training script parameters\n",
    "                    script_params = {\n",
    "                        # Required Params\n",
    "                        \"--config_file\": \"bert-large.json\",\n",
    "                        # bert_data is where pre-processed training data are\n",
    "                        '--train_path' : ds.path('bert_data/512/wiki_pretrain/').as_mount(), \n",
    "                        '--validation_path':ds.path('bert_data/validation_512_only/').as_mount(),\n",
    "                        # Optional Params\n",
    "                        \"--max_seq_length\": 512,\n",
    "                        \"--max_predictions_per_seq\": 80,\n",
    "                        \"--masked_lm_prob\": 0.15,\n",
    "                        \"--train_batch_size\": 64,\n",
    "                        '--seed': 42,\n",
    "                        '--accumulate_gradients': \"True\",\n",
    "                        '--gradient_accumulation_steps': 16,\n",
    "                        '--fp16': \"True\",\n",
    "                        '--loss_scale': 0,\n",
    "                        '--epochs' : 2,\n",
    "                        '--config_file_path' :ds.path('bert_data/').as_mount() ,\n",
    "                        '--output_dir':ds.path(f'bert_data/output/{experiment_name}/').as_mount(),\n",
    "                        '--best_cp_dir':ds.path(f'bert_data/best_cp/{experiment_name}/').as_mount(),\n",
    "                        '--latest_cp_dir':ds.path(f'bert_data/latest_cp/{experiment_name}/').as_mount(),\n",
    "                        '--backend':\"nccl\"\n",
    "                    },\n",
    "                    \n",
    "                    entry_script='train.py',\n",
    "                    inputs=[ds.path('bert_data/').as_mount()]\n",
    "                   )\n",
    "# path to the Python environment in the custom Docker image\n",
    "estimator._estimator_config.environment.python.interpreter_path = '/opt/miniconda/envs/amlbert/bin/python'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For single node (1 NC24s_v3 VM), multi-GPU runs for debugging purposes, use the following configuration:\n",
    "- node_count=1, '--config_file':`bert-base-single-node.json`, '--gradient_accumulation_steps': `64`,\"--train_batch_size\": `1024` (for bert-base)\n",
    "- node_count=1, '--config_file':`bert-large-single-node.json`,'--gradient_accumulation_steps': `256`,\"--train_batch_size\": `1024` (for bert-large)\n",
    "\n",
    "To resume from the latest checkpoint, use `load_training_checkpoint` parameter to pass the checkpoint directory. It will load the latest checkpoint from the directory.\n",
    "\n",
    "'--load_training_checkpoint':ds.path(f'bert_data/latest_cp/{experiment_name}/').as_mount(),"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit the run\n",
    "run = experiment.submit(estimator)\n",
    "RunDetails(run).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Python packages required to use Tensorboard with AzureML are azureml-tensorboard, tensorboardX and tensorboard\n",
    "# The Tensorboard constructor takes an array of runs, so be sure and pass it in as a single-element array here\n",
    "from azureml.tensorboard import Tensorboard\n",
    "tb = Tensorboard([run])\n",
    "tb.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#tb.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#run.cancel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Downloading log file to run the perf benchmarking script\n",
    "fetched_run = Run(experiment, run.id)\n",
    "for f in fetched_run.get_file_names():\n",
    "    if \"70_driver_log_rank_0.txt\" in f:\n",
    "        dest = os.path.join('outputs', f.split('/')[-1])\n",
    "        print('Downloading file {} to {}...'.format(f, dest))\n",
    "        fetched_run.download_file(f, dest)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get average throughput of the training run\n",
    "sys.path.append(os.path.abspath(os.path.join('..', '')))\n",
    "from benchmark import *\n",
    "get_perf_metrics(dest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "file_extension": ".py",
  "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"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3
 },
 "nbformat": 4,
 "nbformat_minor": 2
}