tutorials/using-pytorch-lightning/4.train-multi-node-ddp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-node distributed training with PyTorch Lightning\n",
    "\n",
    "description: multi-node, multi-gpu distributed pytorch lightning with distributeddataparallel (ddp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azureml.core import Workspace\n",
    "\n",
    "ws = Workspace.from_config()\n",
    "ws"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import git\n",
    "from pathlib import Path\n",
    "\n",
    "# get root of git repo\n",
    "prefix = Path(git.Repo(\".\", search_parent_directories=True).working_tree_dir)\n",
    "\n",
    "# training script\n",
    "source_dir = prefix.joinpath(\n",
    "    \"code\", \"train\", \"pytorch-lightning\", \"mnist-autoencoder\"\n",
    ")\n",
    "script_name = \"train-multi-node.py\"\n",
    "\n",
    "# environment file\n",
    "environment_file = prefix.joinpath(\"environments\", \"pt-lightning.yml\")\n",
    "\n",
    "# azure ml settings\n",
    "environment_name = \"pt-lightning\"\n",
    "experiment_name = \"pt-lightning-ddp-example\"\n",
    "cluster_name = \"gpu-K80-2\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create environment\n",
    "\n",
    "Define a conda environment YAML file with your training script dependencies and create an Azure ML environment. This notebook will use the same environment definition that was used for part 1 of the tutorial. The dependencies for this tutorial include **mlflow** and **azureml-mlflow**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azureml.core import Environment\n",
    "\n",
    "env = Environment.from_conda_specification(environment_name, environment_file)\n",
    "\n",
    "# specify a GPU base image\n",
    "env.docker.enabled = True\n",
    "env.docker.base_image = (\n",
    "    \"mcr.microsoft.com/azureml/openmpi3.1.2-cuda10.2-cudnn8-ubuntu18.04\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, you can just capture all your dependencies directly in a custom Docker image or Dockerfile, and create your environment from that. For more information, see [Train with custom image](https://docs.microsoft.com/azure/machine-learning/how-to-train-with-custom-image)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adapt training script to set required env vars\n",
    "\n",
    "In order to run multi-node Lightning jobs, Lightning requires the following environment variables to be set on each node in your cluster:\n",
    "\n",
    "* `MASTER_ADDR`: IP address of rank 0 node\n",
    "* `MASTER_PORT`: free port on rank 0 node\n",
    "* `NODE_RANK`: global rank of the node (from 0 to N, where N is the total number of nodes)\n",
    "\n",
    "Since Azure ML does not currently set these environment variables, we will write a utility script *azureml_env_adapter.py* that will set those environment variables using the OpenMPI environment variables that are set on each node. Import the `set_environment_variables()` method from the utility script into your training script, and call this method in the beginning of the training script (in this case inside the `cli_main()` method).\n",
    "\n",
    "In a future release, Azure ML will set these environment variables automatically for PyTorch jobs, at which point this adapter code will no longer be necessary. Once this is available, we will update this tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure and run training job\n",
    "Create a ScriptRunConfig to specify the training script & arguments, environment, and cluster to run on.\n",
    "\n",
    "Lightning supports several [distributed modes](https://pytorch-lightning.readthedocs.io/en/stable/multi_gpu.html#distributed-modes) for training. DistributedDataParallel (DDP) is recommended over DataParallel (DP) for training.\n",
    "\n",
    "For multi-node, specify the number of GPUs per node to train on (typically this will correspond to the number of GPUs in your cluster's SKU) and the distributed mode, in this case DistributedDataParallel (\"ddp\"). In addition, specify the number of nodes to use for distributed training. PyTorch Lightning expects these as arguments `--gpus`, `--accelerator` and `--num_nodes`, respectively. See their [Multi-GPU DistributedDataParallel](https://pytorch-lightning.readthedocs.io/en/stable/multi_gpu.html#distributed-data-parallel) training documentation for more information. Note that you do not need to define these flags manually in your training script as Lightning can add them automatically. The training script parses the command-line arguments and passes them to the [`Trainer()`](https://pytorch-lightning.readthedocs.io/en/stable/trainer.html?highlight=Trainer).\n",
    "\n",
    "### Azure ML distributed job configuration\n",
    "In order for Azure ML to launch the multi-node job, define an `MpiConfiguration` with a `node_count` count value that matches the value you specified to your training script's *--num_nodes* argument. For the MpiConfiguration, set `process_count_per_node=1` - this is already the default value, so we don't need to explicitly specify it again here. Note that even though we are running a multi-node, multi-GPU job, we are only specifying Azure ML to launch one process per node. This is because Lightning will handle launching the extra processes for each GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azureml.core import ScriptRunConfig, Experiment\n",
    "from azureml.core.runconfig import MpiConfiguration\n",
    "\n",
    "cluster = ws.compute_targets[cluster_name]\n",
    "\n",
    "num_nodes = 2\n",
    "src = ScriptRunConfig(\n",
    "    source_directory=source_dir,\n",
    "    script=script_name,\n",
    "    arguments=[\n",
    "        \"--max_epochs\",\n",
    "        50,\n",
    "        \"--gpus\",\n",
    "        2,\n",
    "        \"--accelerator\",\n",
    "        \"ddp\",\n",
    "        \"--num_nodes\",\n",
    "        num_nodes,\n",
    "    ],\n",
    "    compute_target=cluster,\n",
    "    environment=env,\n",
    "    distributed_job_config=MpiConfiguration(node_count=num_nodes),\n",
    ")\n",
    "\n",
    "run = Experiment(ws, experiment_name).submit(src)\n",
    "run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can monitor the progress of the run with a Jupyter widget. Like the run submission, the widget is asynchronous and provides live updates every 10-15 seconds until the job completes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azureml.widgets import RunDetails\n",
    "\n",
    "RunDetails(run).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run.wait_for_completion(show_output=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8",
   "language": "python",
   "name": "python3.8"
  },
  "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}