ipu/examples/jax_ipu_pmap_quickstart.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bee6e159",
   "metadata": {},
   "source": [
    "# Multi-IPUs in JAX: `pmap` quickstart\n",
    "\n",
    "JAX experimental on IPUs supports multiple IPUs, and collective operations between them (with some limitations on the topology).\n",
    "\n",
    "This guide is directly inspired by JAX official documentation and multi devices examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3824dac4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found existing installation: jax 0.3.16+ipu\n",
      "Uninstalling jax-0.3.16+ipu:\n",
      "  Successfully uninstalled jax-0.3.16+ipu\n",
      "Found existing installation: jaxlib 0.3.15+ipu.sdk310\n",
      "Uninstalling jaxlib-0.3.15+ipu.sdk310:\n",
      "  Successfully uninstalled jaxlib-0.3.15+ipu.sdk310\n"
     ]
    }
   ],
   "source": [
    "# Install experimental JAX for IPUs (SDK 3.1) from Github releases.\n",
    "import sys\n",
    "!{sys.executable} -m pip uninstall -y jax jaxlib\n",
    "!{sys.executable} -m pip install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk310 -f https://graphcore-research.github.io/jax-experimental/wheels.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5bda979f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.config import config\n",
    "\n",
    "# Select how many IPUs will be visible.\n",
    "config.FLAGS.jax_ipu_device_count = 4\n",
    "\n",
    "# Simulating `pmap` on CPU devices instead.\n",
    "# os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'\n",
    "# config.FLAGS.jax_platforms = \"cpu,ipu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ae924cd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b82bb9aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Platform=ipu\n",
      "Number of devices=4\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[IpuDevice(id=0, num_tiles=1472, version=ipu2),\n",
       " IpuDevice(id=1, num_tiles=1472, version=ipu2),\n",
       " IpuDevice(id=2, num_tiles=1472, version=ipu2),\n",
       " IpuDevice(id=3, num_tiles=1472, version=ipu2)]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"Platform={jax.default_backend()}\")\n",
    "print(f\"Number of devices={jax.device_count()}\")\n",
    "ipu_devices = jax.devices(\"ipu\")\n",
    "ipu_devices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5018155d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "5c7f72da",
   "metadata": {},
   "source": [
    "# Basic `pmap`: pure map with no communication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "405ea0ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import pmap\n",
    "from jax import lax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "27173e7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_devices = len(jax.devices())\n",
    "N = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b62360de",
   "metadata": {},
   "outputs": [],
   "source": [
    "@pmap\n",
    "def square(x):\n",
    "    return x ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1d9a663f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: The compile time engine option debug.branchRecordTile is set to \"5887\" when creating the Engine. (At compile-tile it was set to 1471)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INPUT: [[ 0.  1.  2.]\n",
      " [ 3.  4.  5.]\n",
      " [ 6.  7.  8.]\n",
      " [ 9. 10. 11.]]\n",
      "OUTPUT: [[  0.   1.   4.]\n",
      " [  9.  16.  25.]\n",
      " [ 36.  49.  64.]\n",
      " [ 81. 100. 121.]] <class 'jaxlib.xla_extension.pmap_lib.ShardedDeviceArray'>\n"
     ]
    }
   ],
   "source": [
    "data = np.arange(N * num_devices, dtype=np.float32).reshape((num_devices, -1))\n",
    "# First call triggers compilation, which can take a bit of time.\n",
    "output = square(data)\n",
    "\n",
    "print(\"INPUT:\", data)\n",
    "print(\"OUTPUT:\", output, type(output))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "465e4295",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OUTPUT: [[0.0000e+00 1.0000e+00 1.6000e+01]\n",
      " [8.1000e+01 2.5600e+02 6.2500e+02]\n",
      " [1.2960e+03 2.4010e+03 4.0960e+03]\n",
      " [6.5610e+03 1.0000e+04 1.4641e+04]]\n"
     ]
    }
   ],
   "source": [
    "# Second call is fast, with code already compiled and loaded on IPU devices.\n",
    "output = square(output)\n",
    "print(\"OUTPUT:\", output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5c367a8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9c601b7a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[IpuDevice(id=0, num_tiles=1472, version=ipu2),\n",
       " IpuDevice(id=1, num_tiles=1472, version=ipu2),\n",
       " IpuDevice(id=2, num_tiles=1472, version=ipu2),\n",
       " IpuDevice(id=3, num_tiles=1472, version=ipu2)]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Buffers of output sharded array living on different IPUs.\n",
    "[b.device() for b in output.device_buffers]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "80b807e1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DeviceArray([ 81., 256., 625.], dtype=float32),\n",
       " IpuDevice(id=3, num_tiles=1472, version=ipu2))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Can move on the fly to a different IPU.\n",
    "b = jax.device_put(output[1], ipu_devices[3])\n",
    "b, b.device()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2dd69ea6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[ 0.,  1., 16.]], dtype=float32)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output[0:1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85a2ad2e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "1e73c6d8",
   "metadata": {},
   "source": [
    "# Collective communication operations\n",
    "\n",
    "JAX on IPU collective operations are implemented using Graphcore GCL library (https://docs.graphcore.ai/projects/poplar-user-guide/en/latest/gcl.html). Similarly to TPUs, some restrictions on the IPU mesh topology apply."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bfeb770",
   "metadata": {},
   "source": [
    "## Single `pmap` reduction across all devices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f6c2ed1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "@partial(pmap, axis_name=\"i\")\n",
    "def normalize(x):\n",
    "    return x / lax.psum(x, axis_name=\"i\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "5614096d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: The compile time engine option debug.branchRecordTile is set to \"2943\" when creating the Engine. (At compile-tile it was set to 1471)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ShardedDeviceArray([[0.        , 0.2       , 0.2857143 ],\n",
       "                    [1.        , 0.8       , 0.71428573]], dtype=float32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# First call: compile & load on IPU devices.\n",
    "output = normalize(data[0:2])\n",
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f9f1c488",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1., 1., 1.], dtype=float32)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Proper normalization across IPUs!\n",
    "np.sum(np.asarray(output), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78a10225",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "c8ed3cfe",
   "metadata": {},
   "source": [
    "## `pmap` reduction across different replica groups\n",
    "\n",
    "A typical usecase of multiple `pmap` axes is combining data parallelism and tensor parallelism."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d4b96a55",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: The compile time engine option debug.branchRecordTile is set to \"5887\" when creating the Engine. (At compile-tile it was set to 1471)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OUTPUTS: (ShardedDeviceArray([[0.  , 0.25],\n",
      "                    [1.  , 0.75]], dtype=float32), ShardedDeviceArray([[0. , 1. ],\n",
      "                    [0.4, 0.6]], dtype=float32), ShardedDeviceArray([[0.        , 0.16666667],\n",
      "                    [0.33333334, 0.5       ]], dtype=float32))\n"
     ]
    }
   ],
   "source": [
    "@partial(pmap, axis_name='rows')\n",
    "@partial(pmap, axis_name='cols')\n",
    "def f(x):\n",
    "    row_normed = x / lax.psum(x, 'rows')\n",
    "    col_normed = x / lax.psum(x, 'cols')\n",
    "    doubly_normed = x / lax.psum(x, ('rows', 'cols'))\n",
    "    return row_normed, col_normed, doubly_normed\n",
    "\n",
    "x = np.arange(4., dtype=np.float32).reshape((2, 2))\n",
    "outputs = f(x)\n",
    "\n",
    "print(\"OUTPUTS:\", repr(outputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37a5152a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "630915d5",
   "metadata": {},
   "source": [
    "# Manual data sharding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "fa5b605e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_devices0 = ipu_devices[:2]\n",
    "sub_devices1 = ipu_devices[2:]\n",
    "\n",
    "indata0 = jax.device_put_sharded([v for v in data[:2]], sub_devices0)\n",
    "indata1 = jax.device_put_sharded([v for v in data[2:]], sub_devices1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "060f3316",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "107cec6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_fn(x):\n",
    "    return x / lax.psum(x, 'i')\n",
    "\n",
    "def normalize_fn2(x):\n",
    "    return x / lax.pmean(x, 'i')\n",
    "\n",
    "normalize0 = pmap(normalize_fn, axis_name='i', devices=sub_devices0)\n",
    "normalize1 = pmap(normalize_fn2, axis_name='i', devices=sub_devices1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "bfec8c81",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: The compile time engine option debug.branchRecordTile is set to \"2943\" when creating the Engine. (At compile-tile it was set to 1471)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(ShardedDeviceArray([[0.        , 0.2       , 0.2857143 ],\n",
       "                     [1.        , 0.8       , 0.71428573]], dtype=float32),\n",
       " ShardedDeviceArray([[0.8       , 0.8235294 , 0.84210527],\n",
       "                     [1.2       , 1.1764706 , 1.1578947 ]], dtype=float32))"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out0 = normalize0(indata0)\n",
    "out1 = normalize1(indata1)\n",
    "\n",
    "out0, out1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "4cd200e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([IpuDevice(id=0, num_tiles=1472, version=ipu2),\n",
       "  IpuDevice(id=1, num_tiles=1472, version=ipu2)],\n",
       " [IpuDevice(id=2, num_tiles=1472, version=ipu2),\n",
       "  IpuDevice(id=3, num_tiles=1472, version=ipu2)])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[b.device() for b in out0.device_buffers], [b.device() for b in out1.device_buffers]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebf451a5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e68205aa",
   "metadata": {},
   "source": [
    "## Collective `permute` between IPUs\n",
    "\n",
    "Potentially useful for compiling a pipeline on a Transformer model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b11c3ab7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax._src.lib import xla_bridge\n",
    "device_count = jax.device_count()\n",
    "\n",
    "def send_right(x, axis_name):\n",
    "    left_perm = [(i, (i + 1) % device_count) for i in range(device_count)]\n",
    "    return lax.ppermute(x, perm=left_perm, axis_name=axis_name)\n",
    "\n",
    "def send_left(x, axis_name):\n",
    "    left_perm = [((i + 1) % device_count, i) for i in range(device_count)]\n",
    "    return lax.ppermute(x, perm=left_perm, axis_name=axis_name)\n",
    "\n",
    "def update_board(board):\n",
    "    left = board[:-2]\n",
    "    right = board[2:]\n",
    "    center = board[1:-1]\n",
    "    return lax.bitwise_xor(left, lax.bitwise_or(center, right))\n",
    "\n",
    "@partial(pmap, axis_name='i')\n",
    "def step(board_slice):\n",
    "    left, right = board_slice[:1], board_slice[-1:]\n",
    "    right, left = send_left(left, 'i'), send_right(right, 'i')\n",
    "    enlarged_board_slice = jnp.concatenate([left, board_slice, right])\n",
    "    return update_board(enlarged_board_slice)\n",
    "\n",
    "def print_board(board):\n",
    "    print(''.join('*' if x else ' ' for x in np.asarray(board).ravel()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "c9f6a659",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                    *                   \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: The compile time engine option debug.branchRecordTile is set to \"5887\" when creating the Engine. (At compile-tile it was set to 1471)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                   ***                  \n",
      "                  **  *                 \n",
      "                 ** ****                \n",
      "                **  *   *               \n",
      "               ** **** ***              \n",
      "              **  *    *  *             \n",
      "             ** ****  ******            \n",
      "            **  *   ***     *           \n",
      "           ** **** **  *   ***          \n",
      "          **  *    * **** **  *         \n",
      "         ** ****  ** *    * ****        \n",
      "        **  *   ***  **  ** *   *       \n",
      "       ** **** **  *** ***  ** ***      \n",
      "      **  *    * ***   *  ***  *  *     \n",
      "     ** ****  ** *  * *****  *******    \n",
      "    **  *   ***  **** *    ***      *   \n",
      "   ** **** **  ***    **  **  *    ***  \n",
      "  **  *    * ***  *  ** *** ****  **  * \n",
      " ** ****  ** *  ******  *   *   *** ****\n",
      " *  *   ***  ****     **** *** **   *   \n"
     ]
    }
   ],
   "source": [
    "board = np.zeros(40, dtype=bool)\n",
    "board[board.shape[0] // 2] = True\n",
    "reshaped_board = board.reshape((device_count, -1))\n",
    "\n",
    "print_board(reshaped_board)\n",
    "for _ in range(20):\n",
    "    reshaped_board = step(reshaped_board)\n",
    "    print_board(reshaped_board)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ae93851",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}