ddp_apex_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed Mixed-precision Training with PyTorch and NVIDIA `Apex`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial goes over the most important parts of distributed training in PyTorch.\n",
    "\n",
    "For more details, please refer to: https://github.com/richardkxu/distributed-pytorch\n",
    "\n",
    "For the full ImageNet training script please refer to `imagenet_ddp_apex.py` in the above Git repo."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is `Apex`?\n",
    "A Pytorch extension with NVIDIA-maintained utilities to streamline mixed precision and distributed training. It contains the full features of the built-in PyTorch Distributed Data Parallel (DDP) package. Additionally, it integrates better with NVIDIA GPUs and provides mixed-precision training acceleration.\n",
    "\n",
    "`Apex` uses NVIDIA NVIDIA Collective Communications Library (NCCL) as the distributed backend. NCCL handles the communication for GPUs within and across multiple nodes. It curruntly has the best performance and integration with NVIDIA GPUs. \n",
    "\n",
    "Most deep learning frameworks, including PyTorch, train using 32-bit floating point (FP32) arithmetic by default. However, using FP32 for all operations is not essential to achieve full accuracy for many state-of-the-art deep neural networks (DNNs). In mixed precision training, majority of the network uses FP16 arithmetic, while automatically casting potentially unstable operations to FP32.\n",
    "\n",
    "Key points:\n",
    "- Ensuring that weight updates are carried out in FP32.\n",
    "- Loss scaling to prevent underflowing gradients.\n",
    "- A few operations (e.g. large reductions) left in FP32.\n",
    "- Everything else (the majority of the network) executed in FP16."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why `Apex`?\n",
    "\n",
    "- comes with all the distributed training features of the built-in PyTorch DDP\n",
    "- better performance than built-in DDP\n",
    "- reducing memory storage/bandwidth demands by 2x\n",
    "- use larger batch sizes\n",
    "- take advantage of NVIDIA Tensor Cores for matrix multiplications and convolutions\n",
    "- don't need to explicitly convert your model, or the input data, to half()."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How to use `Apex`?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's say we are using 2 computer nodes, each with 4 GPUs. Then:\n",
    "* world size = 8\n",
    "* on each node, local rank of each GPU will be 0-3\n",
    "* the global rank of each GPU will be 0-7\n",
    "* we will use pytorch's `torch.distributed.launch` module to spawn 8 processes, one for each GPU. More on this later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this will be 0-3 if you have 4 GPUs on curr node\n",
    "args.gpu = args.local_rank\n",
    "torch.cuda.set_device(args.gpu)\n",
    "torch.distributed.init_process_group(backend='nccl',\n",
    "                                     init_method='env://')\n",
    "# this is the total # of GPUs across all nodes\n",
    "# if using 2 nodes with 4 GPUs each, world size is 8\n",
    "args.world_size = torch.distributed.get_world_size()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You load your model, and initialize your optimizer like how you usually do in pytorch. However, you need to perform learning rate scaling to combat the difficulties of learning in a distributed context. If your global batch size is 4 times your normal learning rate, then you need to mulitple your learning rate by 4 to \"speed up\" training at the early stage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = models.resnet50(pretrained=True)\n",
    "model.cuda()\n",
    "# Scale init learning rate based on global batch size\n",
    "args.lr = args.lr * float(args.batch_size*args.world_size)/256.\n",
    "optimizer = torch.optim.SGD(model.parameters(), args.lr,\n",
    "                            momentum=args.momentum,\n",
    "                            weight_decay=args.weight_decay)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`apex.amp` is a tool to enable mixed precision training when using `apex`. \n",
    "\n",
    "Optimization level ranges from: \"O0\" to \"O3\". Typically we use either \"O1\" or \"O2\". More details: https://nvidia.github.io/apex/amp.html#opt-levels\n",
    "* O0: pure FP32 training\n",
    "* O1: GEMMs and convolutions are in FP16, model weights, softmax are in FP32\n",
    "* O2: “Almost FP16” Mixed Precision. O2 casts the model weights to FP16, patches the model’s forward method to cast input data to FP16, keeps batchnorms in FP32, maintains FP32 master weights, updates the optimizer’s param_groups so that the optimizer.step() acts directly on the FP32 weights\n",
    "* O3: pure FP16 training, but not stable. Only use as a speedy baseline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize Amp.  Amp accepts either values or strings for the optional override arguments,\n",
    "# for convenient interoperation with argparse.\n",
    "model, optimizer = amp.initialize(model, optimizer,\n",
    "                                  opt_level=args.opt_level,\n",
    "                                  keep_batchnorm_fp32=args.keep_batchnorm_fp32,)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wrap your model with `apex.parallel.DistributedDataParallel`. It is similar to `torch.nn.parallel.DistributedDataParallel`. It enables convenient multiprocess distributed training, optimized for NVIDIA's NCCL communication library.\n",
    "\n",
    "Define your loss function like before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = DDP(model, delay_allreduce=True)\n",
    "criterion = nn.CrossEntropyLoss().cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define your training and test dataset like before. However, you need to call `torch.utils.data.distributed.DistributedSampler` to makes sure that each process gets a different slice of the training data during distributed training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = datasets.ImageFolder(\n",
    "    traindir,\n",
    "    transforms.Compose([\n",
    "        transforms.RandomResizedCrop(crop_size),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        # transforms.ToTensor(), Too slow\n",
    "        # normalize,\n",
    "    ]))\n",
    "val_dataset = datasets.ImageFolder(valdir, transforms.Compose([\n",
    "        transforms.Resize(val_size),\n",
    "        transforms.CenterCrop(crop_size),\n",
    "    ]))\n",
    "train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)\n",
    "val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For dataloader, `args.batch_size` is the per GPU batch size. Notice that we turn off shuffling and use distributed data sampler. `args.workers` is the number of subprocesses per GPU you want for dataloading. **`args.workers` * number of GPU per node** should be <= **the number of CPU threads capable of your CPU**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = torch.utils.data.DataLoader(\n",
    "    train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),\n",
    "    num_workers=args.workers, sampler=train_sampler)\n",
    "\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "    val_dataset,\n",
    "    batch_size=args.batch_size, shuffle=False,\n",
    "    num_workers=args.workers,\n",
    "    sampler=val_sampler)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want to perform \"learning rate warmup\" to stablize training at early stages. The following is a regular step learning rate schedule with \"warmup\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def adjust_learning_rate(optimizer, epoch, step, len_epoch):\n",
    "    \"\"\"LR schedule that should yield 76% converged accuracy with batch size 256\"\"\"\n",
    "    factor = epoch // 30\n",
    "\n",
    "    if epoch >= 80:\n",
    "        factor = factor + 1\n",
    "\n",
    "    lr = args.lr*(0.1**factor)\n",
    "\n",
    "    \"\"\"Warmup\"\"\"\n",
    "    if epoch < 5:\n",
    "        lr = lr*float(1 + step + epoch*len_epoch)/(5.*len_epoch)\n",
    "\n",
    "    for param_group in optimizer.param_groups:\n",
    "        param_group['lr'] = lr\n",
    "\n",
    "    return lr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For pytorch, usually in your train function, you do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = model(input)\n",
    "loss = criterion(output, target)\n",
    "loss.backward()\n",
    "optimizer.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you should do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = model(input)\n",
    "loss = criterion(output, target)\n",
    "# Mixed-precision training requires that the loss is scaled in order\n",
    "# to prevent the gradients from underflow\n",
    "with amp.scale_loss(loss, optimizer) as scaled_loss:\n",
    "    scaled_loss.backward()\n",
    "optimizer.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bells and Whistles\n",
    "\n",
    "### How to prevent race condition when mutiple devices try to do `Tensorboard` logging or print to output file?\n",
    "If you are training distributedly on 2 nodes, each with 4 GPUs, then we will spawn 8 processes to run the same .py file, one process for each GPU. If two or more processes try to log or write on disk at the same time, race condition could happen. To prevent race condition, we usually only allow the process with global rank 0 to do all the logging and printing. Simple check `torch.distributed.get_rank() == 0` before you do the logging or write to `Tensorboard`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# only allow GPU0 to print training states to prevent double logging\n",
    "if torch.distributed.get_rank() == 0:\n",
    "    writer.add_scalar('Loss/train', train_losses, epoch + 1)\n",
    "    writer.add_scalar('Loss/val', val_losses, epoch + 1)\n",
    "    writer.add_scalar('Top1/train', train_top1, epoch + 1)\n",
    "    writer.add_scalar('Top1/val', val_top1, epoch + 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### How to save model checkpoints?\n",
    "Following the same idea above, we only allow the process with global rank 0 to save model checkpoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.distributed.get_rank() == 0:\n",
    "    is_best = val_top1 > best_prec1\n",
    "    best_prec1 = max(val_top1, best_prec1)\n",
    "    save_checkpoint({\n",
    "        'epoch': epoch + 1,\n",
    "        'arch': args.arch,\n",
    "        'state_dict': model.state_dict(),\n",
    "        'best_prec1': best_prec1,\n",
    "        'optimizer': optimizer.state_dict(),\n",
    "    }, is_best, writer.log_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How to run the program?\n",
    "To run your programe on 2 nodes with 4 GPU each, you will need to open 2 terminals and run slightly different command on each node.\n",
    "\n",
    "Node 0:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr=\"192.168.100.11\" --master_port=8888 imagenet_ddp_apex.py -a resnet50 --b 208 --workers 20 --opt-level O2 /home/shared/imagenet/raw/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- torch distributed launch module: https://github.com/pytorch/pytorch/blob/master/torch/distributed/launch.py\n",
    "- --nproc_per_node: number of GPUs on the current node, each process is bound to a single GPU\n",
    "- ----node_rank: rank of the current node, should be an int between 0 and --world-size - 1\n",
    "- --master_addr: IP address for the master node of your choice. type str\n",
    "- --master_port: open port number on the master node. type int. if you don't know, use 8888\n",
    "- --workers: # of data loading workers for the current node. this is different from the processes that run the programe on each GPU. the total # of processes = # of data loading workers + # of GPUs (one process to run each GPU)\n",
    "- -b: per GPU batch size, for a 16 GB GPU, 224 is the max batch size. Need to be a multiple of 8 to make use of Tensor Cores. If you are using tensorboard logging, you need to assign a slightly smaller batch size!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Node 1:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=1 --master_addr=\"192.168.100.11\" --master_port=8888 imagenet_ddp_apex.py -a resnet50 --b 208 --workers 20 --opt-level O2 /home/shared/imagenet/raw/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}