GPU_Workshop.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install --no-index cupy jax torchvision matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Best practices CPU -> GPU\n",
    "\n",
    "Over the last decade, GPUs have become an integral part of the High-Performance Computing (HPC) world. Originally designed to efficiently render images on a screen, it turns out that GPUs can also be used to speed up scientific computing workloads, sometimes by orders of magnitude!\n",
    "\n",
    "They are not, however, a silver bullet, in the sense that they will not magically accelerate just *anything* that you throw at them. In many cases, using a GPU may even *slow down* your workload. In other cases, a GPU may achieve pretty unimpressive performance gains over a CPU, sometimes no gain at all. Given that GPUs are **much** more expensive, and much less numerous than CPUs on the Digital Research Alliance of Canada's HPC Clusters, you would be better off using one or more CPUs instead in all of those cases.\n",
    "\n",
    "In general, users should be mindful of not wasting resources on the clusters. Whenever your code requests a GPUs, but does not make a reasonable use of them, you are **unnecessarily using up your group's priority**, which will impact your colleagues' ability to run their jobs! You might also be preventing other users, whose jobs actually require a GPU, from acessing these much needed resources!\n",
    "\n",
    "To avoid all these negative impacts, this workshop will introduce you to a set of best practices you can follow when deciding whether or not you need a GPU. There will be many factors that you will need to consider and we will provide concrete steps for you to reason about them. \n",
    "\n",
    "First, let's spend some time looking into the inner-workings of a GPU. Knowing how they work at a very high level will help you make sense out of our first set of recommendations.\n",
    "\n",
    "\n",
    "# What is a GPU and when should I use one?\n",
    "\n",
    "In this section, we will use Nvidia's [Ampere GPU architecture ](https://en.wikipedia.org/wiki/Ampere_(microarchitecture))  as a reference to introduce key concepts that will help you reason about what kinds of workload are best suited to run on GPUs. In general, most of these concepts will also apply more-or-less to other Nvidia GPU architectures and to GPUs designed by other vendors than Nvidia.\n",
    "\n",
    "So what is a GPU anyway?\n",
    "\n",
    "![cpu vs gpu](images/cpu-vs-gpu.001.png)\n",
    "\n",
    "Simply put, a GPU is a [Massively Parallel Processor (MPP)](https://en.wikipedia.org/wiki/Massively_parallel). While a CPU in an HPC node will typically have between 16 and 32 **compute cores**, a modern GPU has thousands of them. On both CPUs and GPUs, each one of these cores is a part of the chip that has the ability to perform a number of floating-point operation (Flop) per clock cycle. On modern CPUs, each core may be able to perform up to 16 Flops at a time. On GPUs, each core can perform one Flop at a time (two if we count fused operations).\n",
    "\n",
    "In the diagram above, the right-hand side represents a portion of a GPU component called a **Streaming Multiprocessor (SM)**. Each SM is made up of 4 of these. A modern A100 GPU has **108 SMs**, boasting a total of 8192 cores.\n",
    "\n",
    "So can we say that GPUs are *always faster* than CPUs because they can perform several thousands of Flops per clock cycle against a CPU's up to a few hundred?\n",
    "\n",
    "No. It is not that simple. \n",
    "\n",
    "It turns out that CPUs and GPUs are not designed for the same kinds of tasks. Each has a set of key characteristics that makes it best suited for different computational workloads. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 1 - Solving Linear Systems\n",
    "\n",
    "In this example, you will simultaneously solve N systems of equations of the type $Ax=b$ where $A$ is a $MxM$ matrix and $B$ is $Mx1$\n",
    "\n",
    "You will be given two functions: one to generate random matrices and one to solve the systems of equations. You will use the magic function <code>%timeit</code> to compare the running times for these functions with different combinations of the following parameters:\n",
    "\n",
    "- Using CPU or GPU\n",
    "- The number of CPUs to use\n",
    "- The number and size of the systems of equations to be solved\n",
    "\n",
    "You should be able to see how these parameters affect performance when using CPUs or a GPU to solve the systems.\n",
    "\n",
    "For the CPU part of teh example, you will use <code>numpy</code> - a Python library that contains optimized implementations of a myriad of numerical operations. For the GPU portion, you will use <code>cupy</code> - a library that aims to be an analogous of <code>numpy</code> but for GPUs.\n",
    "\n",
    "On the first cell, you will import all libraries needed for this example and you will choose the number of CPUs that should be used for your computations. The environment variable <code>OMP_NUM_THREADS</code> controls how many parallel threads should be spawned by routines written using [OpenMP](#) - an API for parallel programming. As it turns out, the low level C code that <code>numpy</code> calls to solve linear systems of equations and to generate random numbers are both examples of routines written using OpenMP!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import timeit\n",
    "\n",
    "N_CPUS=10\n",
    "\n",
    "os.environ[\"OMP_NUM_THREADS\"]= str(N_CPUS)\n",
    "\n",
    "import numpy as np\n",
    "import cupy as cp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next you will define functions to generate random data and to solve linear systems:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_data(n_systems,n_rows,n_cols,device):\n",
    "    lib = cp if device == \"gpu\" else np\n",
    "\n",
    "    rmg = lib.random.rand #random matrix generator\n",
    "    a = rmg(n_systems,n_rows,n_cols)\n",
    "    b = rmg(n_systems,n_rows,1)\n",
    "    return a,b\n",
    "\n",
    "def multiple_solve(a,b):\n",
    "    lib = cp if type(a) == cp.ndarray else np\n",
    "\n",
    "    results = lib.linalg.solve(a,b)\n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now try changing both the number and size of systems, as well as the device used to generate the data. The magic function <code>%timeit</code> will return the average run time of <code>generate_data</code> over several repetitions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit generate_data(n_systems=1, n_rows=10000, n_cols=10000, device=\"gpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You have now got a sense of the time it takes to generate random systems of equations on both CPU and GPU. Now let's store some random data in a variable so we can use it to time the execution time of our second function: <code>multiple_solve</code>. Try it out with different numbers and sizes of systems on both the CPU and the GPU. Notice how the difference in performance changes as the size of the problem grows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a,b = generate_data(1,10000,10000,\"gpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit multiple_solve(a,b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Last, but not least, let's see what happens when we generate data using the CPU, then copy the data to GPU before calling <code>multiple solve</code>. This should illustrate the benefits of generating random numbers directly on the GPU, instead of using a *bounce buffer* to load numbers generated on the CPU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "c,d = generate_data(10,10,10,\"cpu\")\n",
    "c_gpu, d_gpu = (cp.asarray(c),cp.asarray(d)) # This will load c and d into the GPU's memory\n",
    "multiple_solve(c_gpu,d_gpu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### CPU\n",
    "<br>\n",
    "<details>\n",
    "<summary>Key Charcteristics</summary>\n",
    "\n",
    "- Higher Clock Frequency\n",
    "\n",
    "- Low count of general purpose cores\n",
    "\n",
    "- Lower memory bandwidth\n",
    "\n",
    "- Designed to minimize memory access latency\n",
    "\n",
    "- Branching prediction\n",
    "\n",
    "<br>\n",
    "<details>\n",
    "<summary>They make CPUs good at:</summary>\n",
    "\n",
    "- Executing sequences of instructions that require access to different areas of memory at each step\n",
    "- Switching into different tasks that require different sets of instructions fast\n",
    "- Pipelining unordered sequences of instructions over small amounts of data (Out-of-order Execution)\n",
    "- Executing sequences that branch out into different instructions depending on the input data (If-Else)\n",
    "- Parallelizing individual instructions over very small batches of data at a time (SIMD)\n",
    "</details>\n",
    "</details>\n",
    "\n",
    "### GPU\n",
    "<br>\n",
    "<details>\n",
    "<summary>Key Charcteristics</summary>\n",
    "    \n",
    "- Lower clock frequency\n",
    "\n",
    "- High count of specialized cores\n",
    "\n",
    "- Higher memory bandwidth\n",
    "\n",
    "- Designed to maximize amount of data fetched per memory access\n",
    " \n",
    "- No branching prediction\n",
    "\n",
    "<details>\n",
    "<summary>They make GPUs good at:</summary>\n",
    "\n",
    "- Executing sequences of instructions that re-use the same area of memory at each step \n",
    "- Executing different tasks, that require the same sets of instructions, at the same time.\n",
    "- Pipelining ordered sequences of instructions over large amounts of data (In-order Execution)\n",
    "- Executing sequences that do not branch out depending on the input data.\n",
    "- Parallelizing individual instructions over relatively large batches of data at a time (SIMT)\n",
    "    \n",
    "</details>\n",
    "</details>\n",
    "\n",
    "\n",
    "![cpu-vs-cpu-bandwitdh](images/cpu-vs-gpu-bdwidth.png)\n",
    "\n",
    "### First set of best practices - when to use a GPU\n",
    "\n",
    "GPUs are the best option when:\n",
    "\n",
    "- What you are doing requires large volumes of data\n",
    "\n",
    "- What you are doing can be parallelized over thousands of \"workers\"\n",
    "\n",
    "- Your program does not branch out depending on its inputs\n",
    "\n",
    "- Your program re-uses large chunks of data often\n",
    "\n",
    "Now let's look at a few pseudo-code examples that illustrate these ideas.\n",
    "\n",
    "### Pseudo-Code Examples\n",
    "\n",
    "#### Example 1 - Embarassingly parallel problem\n",
    "<br>\n",
    "<details>\n",
    "<summary>Show/Hide</summary>\n",
    "\n",
    "```python\n",
    "\n",
    "for i in range(data.size):\n",
    "\n",
    "    data[i] + constant\n",
    "    \n",
    "```\n",
    "\n",
    "Is this code best suited to run on a CPU or a GPU?\n",
    "\n",
    "<br>\n",
    "<details>\n",
    "<summary>Answer:</summary>\n",
    "\n",
    "It depends on <code>data.size</code>.\n",
    "\n",
    "Assuming both <code>data</code> and <code>constant</code> are **available in the GPU's memory** when this program runs, then: \n",
    "\n",
    "1. CPU wins for **very small** arrays\n",
    "3. GPU wins for any thing larger than a few thousands of elements\n",
    "4. Performance gap increases **very quickly** as array size grows.\n",
    "</details>\n",
    "</details>\n",
    "\n",
    "#### Example 2 - Serial & branching problem\n",
    "<br>\n",
    "<details>\n",
    "<summary>Show/Hide</summary>\n",
    "    \n",
    "```python\n",
    "for i in range(1,data.size):\n",
    "\n",
    "    data[i] = data[i] + data[i-1] if other_data[i]%2 == 0 else another_data[i] + data[i-1]\n",
    "```\n",
    "\n",
    "Is this code best suited to run on a CPU or a GPU?\n",
    "\n",
    "<br>\n",
    "<details>\n",
    "<summary>Answer:</summary>\n",
    "\n",
    "The short answer here is **CPU**.\n",
    "\n",
    "Why does the CPU win?\n",
    "\n",
    "1. Each step depends on the value of the previous step. This is a serial loop.\n",
    "2. Assuming <code>other_data</code> and <code>another_data</code> are completely disjoint, each step fetches values from different memory regions.\n",
    "</details>\n",
    "</details>\n",
    "\n",
    "#### Example 3 - Communication-bound parallel problem\n",
    "<br>\n",
    "<details>\n",
    "<summary>Show/Hide</summary>\n",
    "\n",
    "```python\n",
    "for i in range(1,data.size):\n",
    "\n",
    "    data[i] = data[i] + constant\n",
    "    other_data[i] = other_data[i] + data[i]\n",
    "    more_data[i] = more_data[i] + other_data[i]\n",
    "\n",
    "output = sum(data + other_data + another_data) \n",
    "```\n",
    "\n",
    "Is this code best suited to run on a CPU or a GPU?\n",
    "\n",
    "<br>\n",
    "<details>\n",
    "<summary>Answer:</summary>\n",
    "\n",
    "As in the first example, it depends on <code>data.size</code>.\n",
    "\n",
    "Assuming both <code>data</code>, <code>other_data</code> and <code>more_data</code> are **available in the GPU's memory** when this program runs, then: \n",
    "\n",
    "1. CPU wins with **very small** arrays\n",
    "3. GPU wins for anything larger than a few thousands of elements.\n",
    "4. Performance gap increases **very quickly** as array size grows.\n",
    "</details>\n",
    "</details>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 2 - Linear Regression\n",
    "\n",
    "In this next example, you will use <code>JAX</code> to solve a linear regression problem. <code>JAX</code> is a numerical computing library that, similarly to <code>cupy</code>, aims to bring the <code>numpy</code> user experience to the world of GPUs. Unlike <code>cupy</code> however, JAX is geared towards Machine Learning research, and so it has many useful built-in features to fit models at scale - most notably [automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load cuda/11.4 to use JAX\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, you will define a function to generate random data for our linear regression problem. The function below takes in a mathematical function - preferably a linear function - and the number of desired samples as arguments. It then generates the desired number of samples by adding random noise to randomly chosen points along the line drawn by the mathematical function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_regression_data(function,n_samples):\n",
    "\n",
    "    samples = [(x,function(x) + np.random.normal(0,1) * 0.1) for x in np.sort(np.random.uniform(0,1,n_samples))]\n",
    "    X, y = zip(*samples)\n",
    "\n",
    "    X = np.array(X,dtype=jnp.float32) #Change the numpy array's data type into a JAX data type\n",
    "    X = np.c_[X,np.ones(X.shape[0])]\n",
    "    y = np.array(y,dtype=jnp.float32)\n",
    "    return X,y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by generating 50 noisy samples from the function $y = 2x + 5$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = lambda x: 2 * x + 5\n",
    "N_SAMPLES=50\n",
    "\n",
    "X,y = generate_regression_data(f,N_SAMPLES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our linear regression problem, you will be given only the blue dots and will attempt to recover the red line. In other words, given some data points (the blue dots), you will run an algorithm to **learn** the function that best describes the data (the red line):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_target = np.linspace(0,1,50).reshape(-1,1)\n",
    "\n",
    "plt.scatter(X[:,0],y,label='Samples')\n",
    "plt.plot(X_target,f(X_target),color='r',label='Target')\n",
    "\n",
    "plt.legend(loc='lower right')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will solve the linear regression using the iterative [least-squares approach](https://en.wikipedia.org/wiki/Least_squares). In the cell below, you are given a set of functions that implement this approach. In a nutshell - you will use [gradient descent](https://en.wikipedia.org/wiki/Gradient_descent) for <code>n_iterations</code> in an attempt to minimize the square of the distance between the blue dots and a linear function, that is initialized as $y = 0x + 0$.\n",
    "\n",
    "Notice how this is re-using the same dataset and how it is applying the same operations during each iteration over and over again? And think about the mathematical operations used in the function below - can they be computed in a parallelized manner?\n",
    "\n",
    "- Matrix multiplication of X and beta\n",
    "\n",
    "- Element-wise square of the element-wise difference between two vectors\n",
    "\n",
    "- Mean of a vector\n",
    "\n",
    "- A vector minus another vector multiplied by a constant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model(beta, X):\n",
    "    return (X @ beta).T\n",
    "\n",
    "def loss(beta, X, y):\n",
    "    return jnp.mean((model(beta, X) - y)**2)\n",
    "\n",
    "@jax.jit\n",
    "def update(beta, X, y, learning_rate):\n",
    "    return beta - learning_rate * jax.grad(loss)(beta, X, y)\n",
    "\n",
    "def fit_model(X, y, n_iterations, learning_rate=0.1):\n",
    "    beta = jnp.zeros([X.shape[1],1])\n",
    "    if y.ndim > 1:\n",
    "        y = y.reshape(-1)\n",
    "    for i in range(n_iterations):\n",
    "        beta = update(beta, X, y, learning_rate)\n",
    "\n",
    "    return beta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After <code>n_iterations</code> you should see the coefficients <code>beta</code> of the linear function converge to 2 and 5 - the true values of these coefficients that you had set when you first generated this random dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_fit = fit_model(X, y, 10000)\n",
    "print(f\"Estimated f(x) = {beta_fit[0]} * X + {beta_fit[1]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you plot this **learned function**, you will see that it is pretty close to that red line we had on the plot a couple cells above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(X[:,0],y,label='Samples')\n",
    "plt.plot(X_target, np.dot(np.c_[X_target, np.ones(X.shape[0])],beta_fit), color='g',label=\"Linear Fit\")\n",
    "plt.legend(loc=\"lower right\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's see what is the effect on performance when we increase the number of samples. First you will time the execution of <code>fit_model()</code> using the CPU. What happens to the execution time when the number of samples grow?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_SAMPLES=10000\n",
    "X,y = generate_regression_data(f,N_SAMPLES) #Try this with different values of N_SAMPLES like: 10000, 50000, 100000, 1000000...\n",
    "\n",
    "with jax.default_device(jax.devices(\"cpu\")[0]):\n",
    "\n",
    "    %timeit fit_model(X, y, 10000).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's see how long it takes to load data into the GPU using <code>JAX</code>:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit jax.device_put(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit jax.device_put(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's compare the execution time we obtained for the CPU with the execution time of the same function on the GPU. What can you say about the difference in performance as the number of samples gorw?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_SAMPLES=10000\n",
    "X,y = generate_regression_data(f,N_SAMPLES) #Try this with different values of N_SAMPLES like: 10000, 50000, 100000, 1000000...\n",
    "\n",
    "with jax.default_device(jax.devices(\"gpu\")[0]):\n",
    "    X_gpu = jax.device_put(X)\n",
    "    y_gpu = jax.device_put(y)\n",
    "\n",
    "    %timeit fit_model(X_gpu, y_gpu, 10000).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far you have solved a single-variable linear regression problem, that is, one where the matrix $X$ has dimensions $Nx2$ where $N$ is the number of samples and the second column of $X$ is filled with 1's. As it turns out, the one operation where you use $X$ in <code>fit_model</code> is a matrix multiplication, an operation that GPUs are able to massively parallelize. What do you think will happen if you add more columns to $X$? In other words, what will be the difference in performance between GPU an CPU in a **multivariate linear regression problem**?\n",
    "\n",
    "First you will use a slightly altered function to generate multivariate data (i.e, $X$ with more than 2 columns). This time, the coefficients <code>beta</code> of the linear function used to generate the data will be chosen randomly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_data_multi_regression(n_samples,n_variables):\n",
    "\n",
    "    beta = np.random.randint(1,10,(n_variables + 1,1))\n",
    "\n",
    "    f = lambda x: np.dot(x,beta)\n",
    "\n",
    "    X = np.array(np.random.rand(n_samples,n_variables))\n",
    "    X = np.c_[X,np.ones(X.shape[0])]\n",
    "\n",
    "    y = f(X) + np.random.normal(0,1,(n_samples,1))\n",
    "\n",
    "    X = np.array(X,dtype=jnp.float32)\n",
    "    y = np.array(y,dtype=jnp.float32)\n",
    "\n",
    "    return X,y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_SAMPLES=10000\n",
    "N_VARS=50\n",
    "\n",
    "X,y = generate_data_multi_regression(N_SAMPLES, N_VARS)#Try this with different values of N_SAMPLES and N_VARS\n",
    "\n",
    "with jax.default_device(jax.devices(\"cpu\")[0]):\n",
    "    %timeit fit_model(X, y, 10000, learning_rate=0.001).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with jax.default_device(jax.devices(\"gpu\")[0]):\n",
    "    X_gpu = jax.device_put(X)\n",
    "    y_gpu = jax.device_put(y)\n",
    "\n",
    "    %timeit fit_model(X_gpu, y_gpu, 10000, learning_rate=0.001).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# GPUs and Data Loading\n",
    "\n",
    "As previously mentioned, we are using Nvidia's [Ampere GPU architecture ](https://en.wikipedia.org/wiki/Ampere_(microarchitecture)) as a stand-in for generic GPUs in this material. In this architecture, as well as in older Nvidia architectures and comptemporary architectures from other vendors, GPUs and CPUs **do not share memory**. That is to say, each has its own memory and one cannot access data that is stored on the other's memory. It is also often the case in systems equipped with such GPUs that data cannot be loaded directly from disk to GPU memory. Data must first be loaded via CPU into the system's memory, and then from there it can be loaded into the GPU's memory - an approach known as a *bounce buffer*:\n",
    "\n",
    "![cpu-vs-cpu-bandwitdh](images/bounce_buffer.png)\n",
    "\n",
    "As it turns out, using a bounce buffer comes with a series of negative consequences in terms of performance, which include:\n",
    "\n",
    "- Latency due to multiple read/write operations, leaving the GPU's Streaming Multiprocessors idle while waiting for data from disk.\n",
    "\n",
    "- CPU might multi-task between data loading and running other processes.\n",
    "\n",
    "- Amount of system memory must be taken into consideration.\n",
    "\n",
    "This makes **loading data** into the GPU's memory a critical factor that influences the performance of a GPU program. Also worth noticing is the fact that the negative consequences apply the other way around too! If you need to copy data from GPU memory to system memory, you will also experience these performance setbacks.\n",
    "\n",
    "### Second set of best practices - Data Loading\n",
    "\n",
    "When designing pipelines to load data into GPU memory, always try to:\n",
    "\n",
    "- Load all the data you will need in one go, if possible, before starting your computations.\n",
    "\n",
    "- Load the maximum amount of data you can get away with if loading the whole dataset is not an option.\n",
    "\n",
    "- Use one or more additional CPUs to perform data loading if you must load data iteratively in batches. Avoid using the main process to load data.\n",
    "\n",
    "- Avoid moving data back and forth between GPU and CPU!\n",
    "\n",
    "- Generate random values directly on the GPU if your computatations require using them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 3 - Deep Learning\n",
    "\n",
    "In this final example you will scale up the ideas from the previous example to fit a more complicated model to a dataset. This time you will use <code>PyTorch</code> to fit a [Convolutional Neural Network](https://en.wikipedia.org/wiki/Convolutional_neural_network) to the CIFAR10 dataset - a collection of 60000 images representing 10 different animals and transportation methods.\n",
    "\n",
    "In the linear regression problem, the model being fit to the data was a linear function, which entailed performing two matrix multiplications per iteration of the least squares method. The convolutional neural network you will work with in this example consists of a sequence of two convolutions, followed by 400 concurrent matrix multiplications, then 120 concurrent matrix multiplications and finally 10 concurrent matrix multiplications. Both convolutions and matrix multiplications are operations that can be parallelized. There are also other operations in between each of these convolution and matrix multiplciation *layers*, which can all be parallelized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.datasets import CIFAR10\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "\n",
    "class Net(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "\n",
    "        self.conv1 = nn.Conv2d(3, 6, 5)\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = x.view(-1, 16 * 5 * 5)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "net = Net().cuda() # Load model on the GPU\n",
    "\n",
    "criterion = nn.CrossEntropyLoss().cuda() # Load the loss function on the GPU\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.001)\n",
    "\n",
    "transform_train = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Purely from the perspective of parallelizing a large number of computations, this looks like something that could benefit from a GPU. But, as you've seen in the previous example, another important part of deciding whether or not this should run on a GPU is the amount of data used in these computations. \n",
    "\n",
    "In Deep Learning, it is a common practice to use a variant of the gradient descent method called [mini-batch gradient descent](https://stats.stackexchange.com/questions/488017/understanding-mini-batch-gradient-descent). Instead of using the entire matrix of inputs $X$ at each iteration of the algorithm, you use a small subset of the rows of $X$. The amount of rows that make up this subset, called the **batch size**, is an important factor that affects whether or not as well as how quickly this approach converges.\n",
    "\n",
    "Next, you will use mini-batch gradient descent to fit our convolutional neural network to the CIFAR10 data. You will use the code below to load one batch of images at a time from disk, then time the execution of 3 iterations of this training algorithm. Try it out with different combinations of the three parameters below and see what the impact is on execution time and on GPU usage:\n",
    "\n",
    "1. **BATCH_SIZE**: this is the mini-batch gradient descent batch size, i.e., how many images should be loaded from disk and used by the trainig algorith at a time.\n",
    "2. **NUM_WORKERS**: this controls how many CPUs will be used to load batches of images *in parallel*. The idea is to use multiple processes, with one CPU each, to streamline loading images from disk *without* blocking the main process.\n",
    "3. **PRE_FETCH**: this controls how many batches of images should be loaded from disk *before* they are needed by the training algorithm. These batches will be stored in the system's memory and loaded on the GPU once its their turn to be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 4096\n",
    "NUM_WORKERS = 10\n",
    "PRE_FETCH = 2\n",
    "\n",
    "datadir = os.environ[\"CIFAR10_PATH\"] #f\"{os.getenv('SLURM_TMPDIR')}/data\"\n",
    "dataset_train = CIFAR10(root=datadir, train=True, download=False, transform=transform_train)\n",
    "\n",
    "train_loader = DataLoader(dataset_train, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, prefetch_factor=PRE_FETCH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, n_iterations, data):\n",
    "\n",
    "    for iteration in range(n_iterations):\n",
    "        for batch_idx, (inputs, targets) in data:\n",
    "\n",
    "            inputs = inputs.cuda()\n",
    "            targets = targets.cuda()\n",
    "\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, targets)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "%timeit train(net,n_iterations=3, data=enumerate(train_loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the example above, we used the <code>DataLoader</code> object to iretatively read bacthes of images from disk, which means 2 sets of read/write operations: <code>disk -> system memory</code> and <code>system memory -> GPU memory</code>. The <code>disk -> system memory</code> step is prone to causing bottlenecks, as reading from disk is usually much slower than reading from memory. What would the difference in performance be if we instead loaded the whole dataset to the system memory before starting our computations? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cifar10 = [(batch_idx,(inputs,targets)) for batch_idx, (inputs, targets) in enumerate(train_loader)]\n",
    "\n",
    "%timeit train(net,n_iterations=3, data=cifar10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The convolution neural network you used above was, relatively speaking, not that big. If you look at how much GPU memory is used when that model is loaded, you will see it is somewhere between 1 and 2 GB. There is more than enough space left to load the CIFAR10 dataset. Could we have sped-up the training algorithm even more simply by loading all images in GPU memory before calling <code>train()</code>?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cifar10 = [(batch_idx,(inputs.cuda(),targets.cuda())) for batch_idx, (inputs, targets) in enumerate(train_loader)]\n",
    "\n",
    "%timeit train(net,n_iterations=3, data=cifar10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's try this out with a much larger convolutional neural network - the [ResNet152]() model, which takes up about 20GB in memory and could thus justify keeping the dataset on disk instead of loading it on the GPU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models import resnet152\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "resnet = resnet152().cuda()\n",
    "criterion = nn.CrossEntropyLoss().cuda() # Load the loss function on the GPU\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit train(resnet,n_iterations=3,data=enumerate(train_loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Monitoring & Profiling - How do I know if I'm using a GPU correctly?\n",
    "\n",
    "\n",
    "So far we have seen two sets of best practices to keep in mind when reasoning about using GPUs in your computations. They are good pointers in theory, but even if you follow them to the letter it is still possible that things will not go as expected. Thus, the best way to make sure you are getting maximum performance out of a GPU is to monitor its usage and to profile your code.\n",
    "\n",
    "### Nvidia-smi\n",
    "\n",
    "We will start with monitoring tools. The first thing you should always verify is whether your code is using the GPU at all. The simplest way to answer this question is with the command <code>nvidia-smi</code>:\n",
    "\n",
    "\n",
    "![cpu-vs-cpu-bandwitdh](images/nvidia-smi.png)\n",
    "\n",
    "As you can see above, <code>nvidia-smi</code> will tell you whether or not any processes are currently using the GPU and how much GPU memory they have allocated. \n",
    "\n",
    "A common way of misreading the <code>utilisation %</code> field is to take it as the amount of compute capacity currently being used. That is **not** what that number is. Rather, <code>utilisation %</code> is the proportion of time where there was *at least one kernel* running on the GPU. A low % means that the GPU is either idle, or busy doing other things than running kernels, such as managing memory or scheduling tasks.\n",
    "\n",
    "Ideally, this % stays high for the duration of your job, but **that is not sufficient** to conclude that your code is not wasting GPU capacity. You could after all have a single kernel running 100% of the time using just a few cores, thereby massively wasting GPU capacity.\n",
    "\n",
    "### Nvtop\n",
    "\n",
    "Another way of checking whether or not your code is actually using the GPU is the command <code>nvtop</code>. This tool will display the same usage statistics as <code>nvidia-smi</code> as a line plot that changes over time. You will also get a breakdown of <code>utilisation %</code> per process if you have multiple processes sharing the same GPU:\n",
    "\n",
    "![cpu-vs-cpu-bandwitdh](images/nvtop.png)\n",
    "\n",
    "The line plot is useful to identify peaks and cliffs in <code>utilisation %</code> as well as any changes in memory allocation during the execution of your code. This can provide helpful clues for you to optimize your code to minimize idle time for example.\n",
    "\n",
    "Another interesting number provided by this tool is the **CPU % utilisation**. In general, CPU usage should be low in a GPU program. A high CPU utilisation % might indicate your program is not using the GPU optimally as a lot of work is being done by the CPU.\n",
    "\n",
    "### Cluster Portals\n",
    "\n",
    "Next we turn to a tool that stradles the line between monitoring and profiling: the Digital Research Alliance of Canada's cluster portals:\n",
    "\n",
    "\n",
    "![cpu-vs-cpu-bandwitdh](images/portal.png)\n",
    "\n",
    "\n",
    "Above you can see a different set of statistics pertaining to the same code example used to illustrate <code>nvidia-smi</code> <code>nvtop</code>:\n",
    "\n",
    "- **SM Activity:** This is the average % of time over all SMs where there was at least one set of concurrent operations active, no matter how many kernels that happens to be. Note that *active* does not mean *computing*. Other activities such as waiting for data and acessing memory also count here. Ideally this number should stay over 80% for the duration of your job, but that **does not necessarily mean** you are using 80% of the GPUs compute capacity. It means that either 80% of the time your were using 100% of the capcity and stayed idle 20% of the time, or that you used 80% capacity 100% of the time with 20% of vacant capacity.\n",
    "\n",
    "<br>\n",
    "\n",
    "- **SM Occupancy:**  This is the average, over time, of the ratio of active concurrent operations and the maximum number of concurrent operations supported by the GPU. A high % also **does not necessarily mean** you are using the GPU effectively. For our purposes however, we will say that this number being high is a strong indicator of good usage.\n",
    "\n",
    "<br>\n",
    "\n",
    "The other metrics we see above - <code>Tensor</code>, <code>FP64</code>, <code>FP32</code> and <code>FP16</code> - indicate what type of GPU core is being used and what % of time these specific parts of the card were active. Of note here is the <code>Tensor</code> metric. Tensor cores are specialized cores purpose built for speeding-up n-dimensional tensor operations like multiplications and convolutions. If your code makes heavy use of tensor operations, the Cuda compiler can automatically take advantage of these specialized cores if your program satisfies one of these conditions:\n",
    "\n",
    "- You use [mixed precision](https://en.wikipedia.org/wiki/Mixed-precision_arithmetic) in your program.\n",
    "- You explicitly use [Nvidia's TF32 data type](https://blogs.nvidia.com/blog/2020/05/14/tensorfloat-32-precision-format/) in your program.\n",
    "\n",
    "And additionally:\n",
    "- In a matrix multiplication, all your matrices' dimensions are multiples of 8.\n",
    "- In a convolution, the number of channels and the width are all multiples of 8.\n",
    "\n",
    "A high % usage of Tensor cores should correspond to good speed-ups relative to other GPU core types. One caveat is that the first two conditions require sacrificing precision in computations, so you have to be mindful of whether or not high precision is important for your application.\n",
    "\n",
    "Apart from the GPU metrics depicted above, you can also monitor a number of other job statistics on the Alliance's cluster portals such as: CPU and system memory usage; network metrics; IO metrics, disk usage and more.\n",
    "\n",
    "The alliance currently maintains cluster portals for the following clusters:\n",
    "\n",
    "- [Béluga](https://portail.beluga.calculquebec.ca)\n",
    "- [Narval](https://portail.narval.calculquebec.ca)\n",
    "- [Graham](https://dashboard.graham.sharcnet.ca)\n",
    "\n",
    "### Third set of best practices - Monitoring and Profiling\n",
    "\n",
    "- Always do a dry run of your job before requesting a GPU for long durations.\n",
    "- Make sure your program is even using a GPU in the first place.\n",
    "- Make sure your program is able to keep the GPU busy for most of the duration of your job.\n",
    "- Keep **SM Activity** above 80% most of the time.\n",
    "- Keep **SM Occupancy** above 50% most of the time.\n",
    "- Use Tensor cores when applicable.\n",
    "\n",
    "\n",
    "# When should I use multiple CPUs instead of a GPU?\n",
    "\n",
    "Let's assume you have a program that fits the bill of the first and second sets of best practices, but you can't for the love of all that is good get it to keep the GPU consistently busy. Maybe you have bursts of high usage followed by long periods of no usage at all. Maybe you can't get activity and/or occupancy above 50%.\n",
    "\n",
    "In all these cases, a GPU still beats a single CPU hands down... but how does it compare to multiple CPUs?\n",
    "\n",
    "While CPUs are not purpose-built to handle the type and scale of operations that a GPU does, they are still pretty good at it - especially when you can throw lots of them at the problem. In many cases, like linear algebra operators and fourier transforms, highly optimized parallel implementations of these operations for multiple CPUs have been around for ages.\n",
    "\n",
    "Furthermore, unless you are writing Cuda code directly, it is often possible to run the exact same code without any changes on a GPU or one/many CPUs.\n",
    "\n",
    "This, coupled with the fact that there are **many, many, many** more CPUs available in the Alliance's clusters than there are GPUs, make using multiple CPUs a great alternative.\n",
    "\n",
    "### Narval\n",
    "\n",
    "- CPUs: 75584\n",
    "- GPUs: 636\n",
    "\n",
    "### Béluga\n",
    "\n",
    "- CPUs: 32080\n",
    "- GPUs: 688\n",
    "\n",
    "### Cedar\n",
    "\n",
    "- CPUs: 90752\n",
    "- GPUs: 1352\n",
    "\n",
    "### Graham:\n",
    "\n",
    "- CPUs: 37664\n",
    "- GPUs: 536\n",
    "\n",
    "So even in cases where you **can** make an effective usage of a GPU, you might get, say, 16 CPUs allocated to you much faster and possibly beat the GPU's performance depending on the nature of your program!\n",
    "\n",
    "\n",
    "### Fourth set of best practices - Use multiple CPUs instead of a GPU\n",
    "\n",
    "- When dry-running your GPU code, dry run it with multiple CPUs too. See if there is really an advantage to using a GPU.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submitting Jobs\n",
    "\n",
    "So far we've been executing our code examples inside a Jupyter notebook. In real life, we encourage you to do so as a means of verifying that your code is able to use a GPU efficiently. In other words, Jupyter notebooks are a great option to **test your code** and keep an eye on resource usage as it runs.\n",
    "\n",
    "Using Jupyter notebooks, however, is **not recommended** on the Alliance's clusters if you do not need to run your code interactively. If your plan is to run a large workload that will run for a long duration without any intervention on your part, you should **submit a batch job** instead.\n",
    "\n",
    "Here is an example of a [job submission script](https://docs.alliancecan.ca/wiki/Running_jobs) for a job that requires a GPU:\n",
    "\n",
    "```bash\n",
    "#!/bin/bash\n",
    "\n",
    "#SBATCH ntasks=1\n",
    "#SBATCH gres=gpu:1\n",
    "#SBATCH cpus-per-task=4\n",
    "#SBATCH mem=8000M\n",
    "#SBATCH account=def-myself_gpu\n",
    "\n",
    "## YOUR JOB'S COMMANDS GO HERE ##\n",
    "```\n",
    "\n",
    "And here is an example of a job that does not requires multiple CPUs, but no GPU:\n",
    "\n",
    "```bash\n",
    "#!/bin/bash\n",
    "\n",
    "#SBATCH ntasks=1\n",
    "#SBATCH cpus-per-task=16\n",
    "#SBATCH mem=32000M\n",
    "#SBATCH account=def-myself_cpu\n",
    "\n",
    "## YOUR JOB'S COMMANDS GO HERE ##\n",
    "```\n",
    "\n",
    "**Important note:** the parameter <code>--mem</code> controls the amount of **system memory** your job is requesting. The amount of GPU memory you will get is fixed, ad depends on the type of GPU allocated to your job:\n",
    "\n",
    "![GPU Types](images/gpu_types.png)\n",
    "\n",
    "On Clusters where there's more than one type of GPU available, you can specify the one you want by adding the \"Slurm type specifier\" (third column on the image above) to the parameter <code>--gres=gpu</code>.\n",
    "\n",
    "For example, to specify that you want one V100 GPU with 32GB of memory on Cedar, the parameter should be set as <code>--gres=gpu:v100l:1</code> \n",
    "\n",
    "## Using Multiple CPUs\n",
    "\n",
    "In all code examples in this notebook, you experimented with using multiple CPUs for different purposes. On the first two cases, we compared execution performance on GPU versus an increasing number of CPUs. On the third, we used the <code>num_workers</code> parameter of the <code>DataLoader</code> object to tell our program to use multiple CPUs and multiple processes to load data from disk. In all these cases, we had a variable in our code where we \"hardcoded\" the desired number of CPUs.\n",
    "\n",
    "**That is not the recommended way to this when submitting jobs.**\n",
    "\n",
    "Instead, you should use a special environment variable that the cluster scheduler creates for you when you submit a job: <code>SLURM_CPUS_PER_TASK</code>. This variable will be set to the same value you pass to the parameter <code>--cpus-per-task</code>. \n",
    "\n",
    "On the first example, setting the variable <code>OMP_NUM_THREADS</code> would look like this:\n",
    "```python\n",
    "N_CPUS = os.environ[\"SLURM_CPUS_PER_TASK\"]\n",
    "\n",
    "os.environ[\"OMP_NUM_THREADS\"] = N_CPUS\n",
    "```\n",
    "\n",
    "On the third example, we would limit the number of threads that <code>torch</code> can spawn with:\n",
    "\n",
    "```python\n",
    "N_CPUS = os.environ[\"SLURM_CPUS_PER_TASK\"]\n",
    "torch.set_num_threads(N_CPUS)\n",
    "\n",
    "```\n",
    "\n",
    "Why is this important?\n",
    "\n",
    "As it turns out, many widely used Python libraries, including <code>torch</code> and <code>jax</code>, <code>tensorflow</code> will by default try to spawn a number of threads equal to the number of CPUs physically installed on the machine where the code is running. As we've seen above, that number will be somwehere between 32 and 64 depending on the cluster you choose. \n",
    "\n",
    "That means that even if you set <code>--cpus-per-task=4</code>, for example, these libraries would still try to spawn 32 - 64 threads. This would result in a scenario called **core oversubscription**, or, in other words, too many threads per CPU. This typically causes code execution to **slow down** significantly, sometimes even freezing it completely. It is thus important to have exactly one thread per CPU.\n",
    "\n",
    "This wraps up our discussion of the best practices when reasoning about using CPUs or a GPU. We hope this material will prove useful to you in designing your workloads in the future!\n",
    "\n",
    "If you have any questions, don't hesitate in writing to us. You can find information on how to communicate with us here [Technical_support](https://docs.alliancecan.ca/wiki/Technical_support) \n"
   ]
  },
  {
   "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}