chapter07_distributed-learning/training-with-multiple-machines.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed training with multiple machines\n",
    "\n",
    "In the previous two tutorials, we saw \n",
    "that using multiple GPUs within a machine \n",
    "can accelerate training. \n",
    "The speedup, however, is limited \n",
    "by the number of GPUs installed in that machine.\n",
    "And it's rare to find a single machine with more than 16 GPUs nowadays. \n",
    "For some truly large-scale applications, \n",
    "this speedup might still be insufficient.\n",
    "For example, it could still take many days \n",
    "to train a state-of-the-art CNN on millions of images.\n",
    "\n",
    "In this tutorial, we'll discuss the key concepts you'll need \n",
    "in order to go from a program that does single-machine training\n",
    "to one that executes distributed training across multiple machines. \n",
    "We depict a typical distributed system in the following figure, where\n",
    "multiple machines are connected by network switches.\n",
    "\n",
    "![](../img/multi-machines.svg)\n",
    "\n",
    "Note that the way we used `copyto` to copy data from one GPU to another in the [multiple-GPU tutorial](../multiple-gpus-scratch.ipynb) does not work when our GPUs are sitting on different machines. To make use of the available resources here well need a better abstraction."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Key-value store\n",
    "\n",
    "MXNet provides a key-value store to synchronize data among devices. The following code initializes an `ndarray` associated with the key \"weight\" on a key-value store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== init \"weight\" ===\n",
      "[[ 0.54881352  0.59284461  0.71518934]\n",
      " [ 0.84426576  0.60276335  0.85794562]]\n",
      "<NDArray 2x3 @cpu(0)>\n"
     ]
    }
   ],
   "source": [
    "from mxnet import kv, nd\n",
    "store = kv.create('local')\n",
    "shape = (2, 3)\n",
    "x = nd.random_uniform(shape=shape)\n",
    "store.init('weight', x) \n",
    "print('=== init \"weight\" ==={}'.format(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After initialization, we can pull the value to multiple devices. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== pull \"weight\" to [gpu(0), gpu(1)] ===\n",
      "[\n",
      "[[ 0.54881352  0.59284461  0.71518934]\n",
      " [ 0.84426576  0.60276335  0.85794562]]\n",
      "<NDArray 2x3 @gpu(0)>, \n",
      "[[ 0.54881352  0.59284461  0.71518934]\n",
      " [ 0.84426576  0.60276335  0.85794562]]\n",
      "<NDArray 2x3 @gpu(1)>]\n"
     ]
    }
   ],
   "source": [
    "from mxnet import gpu\n",
    "ctx = [gpu(0), gpu(1)]\n",
    "y = [nd.zeros(shape, ctx=c) for c in ctx]\n",
    "store.pull('weight', out=y)\n",
    "print('=== pull \"weight\" to {} ===\\n{}'.format(ctx, y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also push new data value into the store. It will first sum the data on the same key and then overwrite the current value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== push to \"weight\" ===\n",
      "[\n",
      "[[ 1.  1.  1.]\n",
      " [ 1.  1.  1.]]\n",
      "<NDArray 2x3 @gpu(0)>, \n",
      "[[ 2.  2.  2.]\n",
      " [ 2.  2.  2.]]\n",
      "<NDArray 2x3 @gpu(1)>]\n",
      "=== pull \"weight\" ===\n",
      "[\n",
      "[[ 3.  3.  3.]\n",
      " [ 3.  3.  3.]]\n",
      "<NDArray 2x3 @gpu(0)>, \n",
      "[[ 3.  3.  3.]\n",
      " [ 3.  3.  3.]]\n",
      "<NDArray 2x3 @gpu(1)>]\n"
     ]
    }
   ],
   "source": [
    "z = [nd.ones(shape, ctx=ctx[i])+i for i in range(len(ctx))]\n",
    "store.push('weight', z)\n",
    "print('=== push to \"weight\" ===\\n{}'.format(z))\n",
    "store.pull('weight', out=y)\n",
    "print('=== pull \"weight\" ===\\n{}'.format(y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With `push` and `pull` we can replace the `allreduce` function defined in [multiple-gpus-scratch](P14-C02-multiple-gpus-scratch.ipynb) by\n",
    "\n",
    "```python\n",
    "def allreduce(data, data_name, store):\n",
    "    store.push(data_name, data)\n",
    "    store.pull(data_name, out=data)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed key-value store\n",
    "\n",
    "Not only can we synchronize data within a machine, with the key-value store we can facilitate inter-machine communication. To use it, one can create a distributed kvstore by using the following command: (Note: distributed key-value store requires `MXNet` to be compiled with the flag `USE_DIST_KVSTORE=1`, e.g. `make USE_DIST_KVSTORE=1`.)\n",
    "\n",
    "```python\n",
    "store = kv.create('dist')\n",
    "```\n",
    "\n",
    "Now if we run the code from the previous section on two machines at the same time, then the store will aggregate the two ndarrays pushed from each machine, and after that, the pulled results will be: \n",
    "\n",
    "```\n",
    "[[ 6.  6.  6.]\n",
    " [ 6.  6.  6.]]\n",
    "```\n",
    "\n",
    "In the distributed setting, `MXNet` launches three kinds of processes (each time, running `python myprog.py` will create a process). One is a *worker*, which runs the user program, such as the code in the previous section. The other two are the *server*, which maintains the data pushed into the store, and the *scheduler*, which monitors the aliveness of each node.\n",
    "\n",
    "It's up to users which machines to run these processes on. But to simplify the process placement and launching, MXNet provides a tool located at [tools/launch.py](https://github.com/dmlc/mxnet/blob/master/tools/launch.py). \n",
    "\n",
    "Assume there are two machines, A and B. They are ssh-able, and their IPs are saved in a file named `hostfile`. Then we can start one worker in each machine through: \n",
    "\n",
    "```\n",
    "$ mxnet_path/tools/launch.py -H hostfile -n 2 python myprog.py\n",
    "```\n",
    "\n",
    "It will also start a server in each machine, and the scheduler on the same machine we are currently on.\n",
    "\n",
    "![](img/dist_kv.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using `kvstore` in `gluon`\n",
    "\n",
    "As mentioned in [our section on training with multiple GPUs from scratch](multiple-gpus-scratch.ipynb#data-parallelism), to implement data parallelism we just need to specify \n",
    "\n",
    "- how to split data\n",
    "- how to synchronize gradients and weights\n",
    "\n",
    "We already see from [multiple-gpu-gluon](P14-C03-multiple-gpus-gluon.ipynb#put-all-things-together) that a `gluon` trainer can automatically aggregate the gradients among different GPUs. What it really does is having a key-value store with type `local` within it. Therefore, to change to multi-machine training we only need to pass a distributed key-value store, for example,\n",
    "\n",
    "```python\n",
    "store = kv.create('dist')\n",
    "trainer = gluon.Trainer(..., kvstore=store)\n",
    "```\n",
    "\n",
    "To split the data, however, we cannot directly copy the previous approach. One commonly used solution is to split the whole dataset into *k* parts at the beginning, then let the *i*-th worker only read the *i*-th part of the data.\n",
    "\n",
    "We can obtain the total number of workers by reading the attribute `num_workers` and the rank of the current worker from the attribute `rank`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total number of workers: 1\n",
      "my rank among workers: 0\n"
     ]
    }
   ],
   "source": [
    "print('total number of workers: %d'%(store.num_workers))\n",
    "print('my rank among workers: %d'%(store.rank))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this information, we can manually access the proper chunk of the input data. In addition, several data iterators provided by `MXNet` already support reading only part of the data. For example,\n",
    "\n",
    "```python\n",
    "from mxnet.io import ImageRecordIter\n",
    "data = ImageRecordIter(num_parts=store.num_workers, part_index=store.rank, ...)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "For whinges or inquiries, [open an issue on  GitHub.](https://github.com/zackchase/mxnet-the-straight-dope)"
   ]
  }
 ],
 "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.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}