P14-C04-training-with-multi-machines.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training with multi-machines\n",
    "\n",
    "On previous two tutorials we saw that using multiple GPUs within a machine can accelerate the training. The speedup, however, is limited by the number of GPUs installed in that machine, which is typically at most 16 nowadays. This amount of speedup is still not good enough for large scale applications, such as training a state-of-the-art CNN on millions of images.\n",
    "\n",
    "In this tutorial, we discuss the key ideas on how to turn a single machine training program into distributed training. A typical distributed system is illustrated in the following figure, where multiple machines are connected by network switches.\n",
    "\n",
    "![](img/multi-machines.svg)\n",
    "\n",
    "Note that the previous way of using `copyto` to copy data from one GPU to another does not work for GPUs sitting on a different machine. We need a better abstraction here."
   ]
  },
  {
   "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 codes initialize a `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 for data synchronization within a machine, the key-value store also supports 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 codes 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 *worker*, which runs the user program, such as the codes in the previous section. The other twos are *server*, which maintains the data pushed into the store, and *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": [
    "## Use `kvstore` in `gluon`\n",
    "\n",
    "As mentioned in [multiple-gpu-scratch](P14-C02-multiple-gpus-scratch.ipynb#data-parallelism), to implement data parallelism we just need to specify \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 reads the *i*-th part of the data.\n",
    "\n",
    "We can obtain the total number of workers by the attribute `num_workers` and the rank of the current worker by 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 seek to a proper position of the input data. In addition, several data iterators provided by `MXNet` already support reading a 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": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}