{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "Copyright (c) 2022 Graphcore Ltd. All rights reserved." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "*Notebook autogenerated from walkthrough.py on 27-Sep-2022*" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# PopTorch Parallel Execution Using Pipelining" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "This PopTorch tutorial demonstrates how to train and test a neural network\n", "on the MNIST dataset by splitting the model over several IPUs and using pipelining\n", "to make efficient use of all the available IPUs.\n", "Pipelining is the recommended strategy\n", "when a model is too large to fit into the memory of a single IPU.\n", "Pipelining can be used by training and inference models.\n", "You will see both scenarios applied to the [MNIST simple application](../../../simple_applications/pytorch/mnist).\n", "\n", "There are two parts: a Jupyter notebook and a Python script. The notebook works through how you use the pipelining API to define pipelines in PopTorch. The `mnist_pipeline.py` script shows the effect of hyper-parameters on a pipelined model's execution schedule when the model is run on the IPU.\n", "In this tutorial you will:\n", "\n", "- Use the PopTorch pipelining API to execute a model over multiple IPUs in parallel.\n", "- Partition a model into blocks (`poptorch.Block`) and group them into pipeline stages (`poptorch.Stage`).\n", "- Learn how gradient accumulation and device iteration hyperparameters impact performance.\n", "- Debug a multi-IPU model using a sharded execution schedule (`poptorch.ShardedExecution`).\n", "- Learn how to offload the optimiser tensors to save memory on device.\n", "- Use the [Popvision Graph Analyser](https://docs.graphcore.ai/projects/graph-analyser-userguide/en/3.11.2/introduction.html) to look at execution profiles and understand how pipelining impacts the scheduling of a model across IPUs.\n", "\n", "If you are unfamiliar with PopTorch, you may want to check out [our tutorial introducing PopTorch](../basics)\n", "and learn how to convert a model from PyTorch to PopTorch to make it run on a Graphcore IPU.\n", "Additionally, some knowledge of data loading in poptorch and batching on IPU is helpful, these topics are covered in [our tutorial on efficient data loading](../efficient_data_loading)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## File structure\n", "\n", "- `mnist_pipeline.py` The main example script.\n", "- `walkthrough.ipynb` The Jupyter notebook for this tutorial.\n", "- `README.md` The markdown version of this tutorial.\n", "- `walkthrough.py` The tutorial annotated script (for notebook generation only).\n", "- `walkthrough_code_only.py` The tutorial code only (for notebook generation only).\n", "- `requirements.txt` The packages to be installed in order to run the code.\n", "- `tests/` The unit tests for `mnist_pipeline.py`." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Running on Paperspace\n", "\n", "The Paperspace environment lets you run this notebook with no set up. To improve your experience we preload datasets and pre-install packages, this can take a few minutes, if you experience errors immediately after starting a session please try restarting the kernel before contacting support. If a problem persists or you want to give us feedback on the content of this notebook, please reach out to through our community of developers using our [slack channel](https://www.graphcore.ai/join-community) or raise a [GitHub issue](https://github.com/gradient-ai/Graphcore-Pytorch/issues)." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "sst_ignore_md", "sst_ignore_code_only" ] }, "outputs": [], "source": [ "%pip install -r requirements.txt" ] }, { "cell_type": "code", "execution_count": null, "id": "3e9a05f5", "metadata": {}, "outputs": [], "source": [ "from examples_utils import notebook_logging\n", "%load_ext gc_logger" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ ">**Note**: In the Python script of this tutorial `mnist_pipeline.py`, the main function is executed under the scope:\n", ">\n", ">```python\n", ">if __name__ == '__main__':\n", ">```\n", ">\n", "> This is necessary to avoid [issues with asynchronous DataLoader](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/batching.html#poptorch-asynchronousdataaccessor).\n", "> Implications of the asynchronous data loader are covered in [our tutorial on efficient data loading](../efficient_data_loading)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Introduction to pipelined execution" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "PopTorch provides the pipelining API to split a model over several IPUs and parallelise its execution.\n", "This is the recommended strategy when a model does not fit into the memory of a single IPU.\n", "To pipeline a model:\n", "\n", "- the model must be partitioned into multiple stages;\n", "- the pipelined execution strategy must be selected, it will coordinate how the different stages are executed.\n", "This process is an example of [model parallelism](https://huggingface.co/docs/transformers/parallelism).\n", "\n", "A pipeline stage groups several consecutive layers of the full model to be executed on the same IPU.\n", "Each stage is placed on an IPU, when training, stages contain both the code for the forward and the backward passes.\n", "Activations corresponding to a batch of data will pass through all the stages during the forward pass, before going back through the stages in reverse order during the backward pass.\n", "By placing the stages on different IPUs, and by feeding enough mini-batches,\n", "all the stages will process a batch of data concurrently after a \"ramp-up\" period.\n", "\n", "![png](static/pipelining.png)\n", "\n", "The figure above illustrates the start of a pipelined training on 4 IPUs.\n", "Each block represents a batch of data being processed (numbered from 1 to 8).\n", "Blocks are marked with `F` when the IPU is executing the forward pass and `B` when it\n", "is executing the backward pass. We can see that after nine steps, the pipeline reaches a point\n", "where all the IPUs are used at the same time.\n", "In general, a pipeline reaches steady state and full utilisation after `2N+1` steps where `N` is the number of IPUs.\n", "For each stage, forward and backward passes are always executed on the same IPU.\n", "\n", "Beyond pipelined execution, other execution strategies are available which offer different performance and memory tradeoffs,\n", "these are beyond the scope of this tutorial but are covered in the [multi-IPU execution strategies section](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/overview.html#multi-ipu-execution-strategies) of the PopTorch documentation." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "For this tutorial, we need to import the following packages." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "import json\n", "import argparse\n", "from pkgutil import get_data\n", "import sys\n", "from tqdm import tqdm\n", "import torch\n", "import torch.nn as nn\n", "import torchvision\n", "import poptorch" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setting hyperparameters" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We are going to use the following hyperparameters.\n", "Let's set their default values." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "learning_rate = 0.001\n", "epochs = 3\n", "batch_size = 40\n", "test_batch_size = 8\n", "training_iterations = 10\n", "gradient_accumulation = 10\n", "inference_iterations = 100" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Preparing the data" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We need to get the MNIST dataset for training or inference.\n", "MNIST contains 70,000 images of handwritten digits: 60,000 for training and\n", "10,000 for testing. The images are grayscale with a resolution of 28x28 pixels.\n", "Two data loaders are employed to load training and testing data sets\n", "separately:\n", "the `train` argument is set to `True` for training, and `False`\n", "for testing. Furthermore, training and inference can load the data\n", "using a specific batch size.\n", "`torchvision.transforms.ToTensor()` converts the image into numbers,\n", "representing the brightness of their color between 0 and 255, separately in\n", "red, green and blue images.\n", "`torchvision.transforms.Normalize()`\n", "normalizes the tensor using the statistics from the MNIST dataset." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "train_dataset = torchvision.datasets.MNIST(\n", " \"mnist_data/\",\n", " train=True,\n", " download=True,\n", " transform=torchvision.transforms.Compose(\n", " [\n", " torchvision.transforms.ToTensor(),\n", " torchvision.transforms.Normalize((0.1307,), (0.3081,)),\n", " ]\n", " ),\n", ")\n", "\n", "test_dataset = torchvision.datasets.MNIST(\n", " \"mnist_data/\",\n", " train=False,\n", " download=True,\n", " transform=torchvision.transforms.Compose(\n", " [\n", " torchvision.transforms.ToTensor(),\n", " torchvision.transforms.Normalize((0.1307,), (0.3081,)),\n", " ]\n", " ),\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "[poptorch.DataLoader](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/pytorch_to_poptorch.html#preparing-your-data)\n", "is used to efficiently load data batches to IPU.\n", "More information about the use of `poptorch.Dataloader`\n", "can be found in the [PopTorch tutorial on efficient data loading](../efficient_data_loading).\n", "`poptorch.DataLoaderMode.Async` is the preferred mode as it can\n", "utilize several CPU threads to accelerate data processing. This can reduce\n", "the host and IPU communication time.\n", "The datasets will be downloaded, shuffled and transformed.\n", "A `poptorch.Options()` instance contains a set of default hyperparameters and options for the IPU.\n", "This is used both by the model and the PopTorch DataLoader.\n", "We change the default value of `deviceIterations`: 10 for training and 100 for inference.\n", "For training with pipelined execution, we also set `gradientAccumulation` to 10 (described in [Pipelined execution section](#pipelined-execution-parallel)).\n", "With these settings, both data loaders will pick 100 mini-batches of data per step." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "train_opts = poptorch.Options()\n", "train_opts.deviceIterations(training_iterations)\n", "train_opts.Training.gradientAccumulation(gradient_accumulation)\n", "\n", "training_data = poptorch.DataLoader(\n", " train_opts,\n", " train_dataset,\n", " batch_size=batch_size,\n", " shuffle=True,\n", " mode=poptorch.DataLoaderMode.Async,\n", " num_workers=16,\n", ")\n", "\n", "test_data = poptorch.DataLoader(\n", " poptorch.Options().deviceIterations(inference_iterations),\n", " test_dataset,\n", " batch_size=test_batch_size,\n", " shuffle=True,\n", " drop_last=True,\n", " mode=poptorch.DataLoaderMode.Async,\n", " num_workers=16,\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Model definition" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "A model needs to be created for the neural network to train on MNIST data.\n", "We start by defining a Module class\n", "called `ConvLayer` that is a basic convolution block for our model. It contains\n", "`Conv2D`, `MaxPool2d` and `ReLU` operations." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class ConvLayer(nn.Module):\n", " def __init__(self, in_channels, num_filters, kernel_size, pool_size):\n", " super().__init__()\n", " self.conv = nn.Conv2d(in_channels, num_filters, kernel_size=kernel_size)\n", " self.pool = nn.MaxPool2d(kernel_size=pool_size)\n", " self.relu = nn.ReLU()\n", "\n", " def forward(self, x):\n", " x = self.conv(x)\n", " x = self.pool(x)\n", " x = self.relu(x)\n", " return x" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To use pipelined execution we must split our model using [block annotations](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/overview.html#annotations).\n", "To provide flexibility, PopTorch lets us split our model into blocks and we can decide later how to group these blocks into\n", "multiple stages to place them on different IPUs." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Annotation for model partitioning" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To prepare a model for pipelining, we need to partition the computational graph into multiple `blocks`.\n", "They will be used later to build the `stages` that will be\n", "executed on different IPUs. The picture below illustrates this process.\n", "\n", "![jpeg](static/stages-poptorch.jpeg)\n", "\n", "As described in the [User Guide](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/overview.html#model-partitioning-using-blocks),\n", "there are three ways to define blocks:\n", "\n", "- using the [poptorch.Block](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/reference.html#poptorch.Block) scope;\n", "- using the [poptorch.BlockFunction](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/reference.html#poptorch.BlockFunction) decorator on the forward method of a `Module`;\n", "- wrapping layers with\n", "[poptorch.BeginBlock](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/reference.html#poptorch.BeginBlock).\n", "\n", "You can use a combination of these three annotation options.\n", "In this tutorial we will use the `poptorch.Block` scope.\n", "\n", "To do this, we modify the `forward` function of the `Network` class, wrapping each layer\n", "under the scope of a `poptorch.Block` with a name: `B1`, `B2`, `B3` and `B4`.\n", "\n", "We define a model called `Network`. It defines four layers using\n", "`ConvLayer`, `Linear` and `ReLU`, where `ConvLayer` was created above." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class Network(nn.Module):\n", " def __init__(self):\n", " super().__init__()\n", " self.layer1 = ConvLayer(1, 10, 5, 2)\n", " self.layer2 = ConvLayer(10, 20, 5, 2)\n", " self.layer3 = nn.Linear(320, 256)\n", " self.layer3_act = nn.ReLU()\n", " self.layer4 = nn.Linear(256, 10)\n", "\n", " def forward(self, x):\n", " with poptorch.Block(\"B1\"):\n", " x = self.layer1(x)\n", " with poptorch.Block(\"B2\"):\n", " x = self.layer2(x)\n", " with poptorch.Block(\"B3\"):\n", " x = x.view(-1, 320)\n", " x = self.layer3_act(self.layer3(x))\n", " with poptorch.Block(\"B4\"):\n", " x = self.layer4(x)\n", " return x\n", "\n", "\n", "model = Network()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Defining the training model" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Since the loss function is only needed for training,\n", "we create a Module class called `TraingModelWithLoss` to\n", "integrate the model and the loss function which is\n", "`CrossEntropyLoss` in this example.\n", "We add the loss to the same block as the last layer of the model.\n", "In the `forward` function of `TrainingModelWithLoss` below,\n", "the loss is annotated with `poptorch.Block` of `B4`." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class TrainingModelWithLoss(torch.nn.Module):\n", " def __init__(self, model):\n", " super().__init__()\n", " self.model = model\n", " self.loss = torch.nn.CrossEntropyLoss()\n", "\n", " def forward(self, args, labels=None):\n", " output = self.model(args)\n", " if labels is None:\n", " return output\n", " with poptorch.Block(\"B4\"):\n", " loss = self.loss(output, labels)\n", " return output, loss" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The training model is called with an instance of\n", "`TrainingModelWithLoss`, model options, and the [AdamW optimiser](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/reference.html#poptorch.optim.AdamW).\n", "The next section of this tutorial will explain how to set the model options for pipelining." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def get_training_model(opts, model):\n", " \"\"\"Wrap a model with the loss and the optimiser into a poptorch.trainingModel\"\"\"\n", " model_with_loss = TrainingModelWithLoss(model)\n", " training_model = poptorch.trainingModel(\n", " model_with_loss,\n", " opts,\n", " optimizer=poptorch.optim.AdamW(model.parameters(), lr=learning_rate),\n", " )\n", " return training_model" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Execution strategies" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now that we have split the model into blocks, we are ready to define stages and place them onto specific IPUs.\n", "Then we will apply an execution strategy which will schedule the stages on the IPUs.\n", "We will consider two execution strategies:\n", "\n", "- [poptorch.PipelinedExecution](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/reference.html#poptorch.PipelinedExecution)\n", "for efficient parallel utilisation;\n", "- and [poptorch.ShardedExecution](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/reference.html#poptorch.ShardedExecution) for debugging multi-IPU models.\n", "\n", "Note: This tutorial does not cover `poptorch.PhasedExecution` strategy, which is the third one available.\n", "\n", "The chosen strategy will be set in the model options, which are initialized by\n", "`poptorch.Options`." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "train_opts = poptorch.Options()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Pipelined execution (parallel)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Assigning blocks to stages and IPUs\n", "\n", "The blocks defined earlier now need to be assembled into stages and passed to the\n", "[poptorch.PipelinedExecution](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/reference.html#poptorch.PipelinedExecution)\n", "strategy object.\n", "This object defines a pipeline which enables parallel execution." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "A stage is created by instantiating a `poptorch.Stage` object and giving it one or more block names.\n", "These stages are then placed on specific IPUs identified by indices from 0 to 3.\n", "In this example, we set 4 stages with `poptorch.PipelinedExecution` to have one stage per IPU." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "pipelined_strategy = poptorch.PipelinedExecution(\n", " poptorch.Stage(\"B1\").ipu(0),\n", " poptorch.Stage(\"B2\").ipu(1),\n", " poptorch.Stage(\"B3\").ipu(2),\n", " poptorch.Stage(\"B4\").ipu(3),\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Next, we set pipelining as the execution strategy in our training options." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "train_opts.setExecutionStrategy(pipelined_strategy)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Note: Since IPUs are all set *in order* from 0 to 3, this shorter notation can also be used to, both, set the execution strategy and place the stages on the IPUs:\n", "\n", "```python\n", "train_opts.setExecutionStrategy(poptorch.PipelinedExecution(\"B1\", \"B2\", \"B3\", \"B4\"))\n", "```\n", "\n", "Instead of naming blocks, we could have placed them directly on specific IPUs by using the `ipu_id` argument in the `poptorch.Block` context manager.\n", "In that case PopTorch supports auto-staging which will automatically define one stage per block, placing each block on the specified IPU.\n", "`PipelinedExecution` and [AutoStage](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/overview.html#poptorch-autostage) are enabled by default in models which contain `poptorch.Block` contexts, decorators or `poptorch.BeginBlock` layers.\n", "Explicitly setting stages, as done in this tutorial, is required when `ipu_id` is not set." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Setting gradient accumulation and device iterations\n", "\n", "Parallelisation during pipelined execution requires both the\n", "[deviceIterations](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/batching.html#poptorch-options-deviceiterations) and [gradientAccumulation](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/batching.html#poptorch-options-training-gradientaccumulation)\n", "hyper-parameters to be set.\n", "\n", "- `gradientAccumulation` is used to push multiple mini-batches through the pipeline\n", "allowing IPUs to run multiple stages in parallel.\n", "During the backward pass, the gradients are accumulated over\n", "several mini-batches. The model parameters are only updated once the number of consecutive mini-batches set by\n", "`opts.Training.gradientAccumulation()` have been processed; their gradients are averaged.\n", "- `deviceIterations` defines the number of times the device should run over the data before returning to the user. This is equivalent to running the IPU in a loop over that specified number of iterations, with new mini-batches of data each time.\n", "However, increasing `deviceIterations` improves efficiency by reducing the number of interactions between the host CPU and the IPU.\n", "\n", "As a result, from the CPU point of view, the number of images used by a single call to the model will be equal to: `batch-size` * `gradientAccumulation` *`deviceIterations`.\n", "The implications of these parameters on I/O efficiency and data batching are covered in detail in the [tutorial on efficient data loading in PopTorch](../efficient_data_loading)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "train_opts.deviceIterations(training_iterations)\n", "train_opts.Training.gradientAccumulation(gradient_accumulation)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The following execution profile illustrates how these parameters affect the computation on IPUs.\n", "\n", "![png](static/pipeline-figure.png)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To generate your own profile, call the example script `mnist_pipeline.py` with\n", "the `--profile` argument, you will find complete usage instructions for this\n", "script [at the end of the tutorial](#how-to-run-the-example-script).\n", "This option will generate `profile.pop` files which can be opened with the\n", "[Popvision Graph Analyser](https://docs.graphcore.ai/projects/graph-analyser-userguide/en/3.11.2/user-guide.html).\n", "\n", "The execution report below illustrates the pipeline for 10 mini-batches\n", "on 4 IPUs. We can see the 4 stages, each on a different IPU.\n", "We highlighted the progression of the first batch:\n", "\n", "- Green for the forward pass.\n", "- Red for the backward pass.\n", "We see that after a few mini-batches, the stages are computing a forward pass\n", "directly followed by a backward pass.\n", "\n", "![png](static/pipelined-execution.png)\n", "\n", "Depending on your model and your performance target, you may need to balance\n", "the pipeline stages in terms of the amount of computation and the amount of\n", "memory demanded." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Efficient model partitioning (advanced)\n", "\n", "In order to make efficient use of multiple IPUs using pipelining, we consider the position of the start and end of each stage.\n", "We must have in mind that these partitions will be placed and executed on different IPUs and the way we split our model will directly impact the memory and performance on each stage.\n", "The following trade-offs must be considered:\n", "\n", "- Stages must fit in memory considering both live and not always live memory. [The PopVision Liveness report](https://docs.graphcore.ai/projects/graph-analyser-userguide/en/3.11.2/user-guide.html#viewing-a-liveness-report-image27) can help determine which Ops and tensors consume the most memory.\n", "- Stages must have similar execution times to avoid IPUs standing idle while another IPU completes its computation.\n", "The execution time of the longest stage determines the duration of one step. The [PopVision execution trace](https://docs.graphcore.ai/projects/graph-analyser-userguide/en/3.11.2/user-guide.html#viewing-an-execution-trace-image32) can be used to identify imbalanced stages.\n", "- Inter-IPU communication is slower than memory access within a single IPU and should be minimised.\n", "\n", "To achieve more efficient pipelines, you may want to have more `poptorch.Block` annotations to allow finer control of the partitioning.\n", "At the `Stage` level, you can reorganize which blocks are assigned to each stage using the following format:\n", "\n", "```python\n", "pipelined_strategy = poptorch.PipelinedExecution(\n", " poptorch.Stage(\"B1\", \"B2\").ipu(0), poptorch.Stage(\"B3\", \"B4\").ipu(1)\n", ")\n", "```\n", "\n", "This format places two blocks into two stages onto two IPUs instead of the four used earlier in the tutorial.\n", "\n", "More information on splitting models for pipelining are available in the\n", "[memory and performance optimisation guide](https://docs.graphcore.ai/projects/memory-performance-optimisation/en/3.1.0/optimising-performance.html#pipeline-execution-scheme), and\n", "the [TensorFlow pipelining technical note](https://docs.graphcore.ai/projects/tf-model-parallelism/en/3.1.0/index.html)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Sharded execution (sequential)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The [poptorch.ShardedExecution](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/reference.html#poptorch.ShardedExecution) option defines sequential execution.\n", "Unlike the pipelining mode, only one batch can be processed at a time so it is not a form of model parallelism.\n", "\n", "`ShardedExecution` can be useful for processing a single sample or\n", "debugging. Overall it has low efficiency since only one IPU is used at a time.\n", "\n", "In this example, 4 stages are specified to run sequentially on 4 IPUs." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "sharded_strategy = poptorch.ShardedExecution(\n", " poptorch.Stage(\"B1\").ipu(0),\n", " poptorch.Stage(\"B2\").ipu(1),\n", " poptorch.Stage(\"B3\").ipu(2),\n", " poptorch.Stage(\"B4\").ipu(3),\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Saving memory by offloading to remote buffers (advanced)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To fit a memory intensive model on IPUs, variables can be offloaded to the host to reduce the live memory.\n", "Activations, weights, the accumulator\n", "and the optimiser can either stay on the IPU or be copied to the host.\n", "It is most efficient to keep them on the IPU but has higher memory\n", "requirements. For larger models, offloading some variables\n", "to the host memory can accommodate other variables on device memory, but it will also\n", "increase the computation time of the weight update as more time is\n", "spent communicating with the host.\n", "In this code example, you can see how to set the storage on the device or host\n", "memory. We kept the default activations, weights, and accumulator on the device,\n", "while storing optimiser states to the host memory\n", "using [poptorch.Option.TensorLocations](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/reference.html#poptorch.options._TensorLocationOptions).\n", "The `AdamW` optimiser is used in this tutorial and its optimiser state variables\n", "do not need to be stored in the device memory during the forward and backward\n", "propagation of the model.\n", "These variables are only required during the weight update and so they are\n", "streamed onto the device during the weight update and then streamed back to\n", "the host memory after they have been updated." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "train_opts.TensorLocations.setOptimizerLocation(\n", " poptorch.TensorLocationSettings().useOnChipStorage(False)\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Training the model" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Once we have the training model, we run the training on the MNIST dataset,\n", "prepared in the first step.\n", "There are two nested `for` loops to traverse epochs and batches respectively.\n", "After calling `training_model()`, predictions and losses are returned.\n", "The accuracy and the mean of losses are calculated and displayed." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def train(training_model, training_data):\n", " nr_steps = len(training_data)\n", " for epoch in range(1, epochs + 1):\n", " print(f\"Epoch {epoch}/{epochs}\")\n", " bar = tqdm(training_data, total=nr_steps)\n", " for data, labels in bar:\n", " preds, losses = training_model(data, labels)\n", " mean_loss = torch.mean(losses).item()\n", " acc = accuracy(preds, labels)\n", " bar.set_description(f\"Loss:{mean_loss:0.4f} | Accuracy:{acc:0.2f}%\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The `accuracy()` function calculates the percentage of\n", "the correct predictions for labels." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def accuracy(predictions, labels):\n", " _, ind = torch.max(predictions, 1)\n", " labels = labels[-predictions.size()[0] :]\n", " accuracy = torch.sum(torch.eq(ind, labels)).item() / labels.size()[0] * 100.0\n", " return accuracy" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "When we call `training_model()`, the outputs are not returned for every\n", "mini-batch by default. This results in incomplete evaluation and inaccurate\n", "training accuracy. This behaviour can be changed by modifying\n", "[outputMode](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/reference.html#poptorch.OutputMode)\n", "in `poptorch.Options()`. For performance reasons, PopTorch uses\n", "`OutputMode.Final` by default. It only returns the final batch loss and\n", "predictions to the host machine. We can set this to `OutputMode.All` to be able\n", "to present the full information in this tutorial. In that case, `preds` and\n", "`losses` will be lists containing the results for the 100 mini-batches passed to\n", "the pipeline. However, this has an impact on the speed of training, due to the\n", "overhead of transferring more data to the host machine." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Warning: impacts performance\n", "train_opts.outputMode(poptorch.OutputMode.All)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We are ready to start training!" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "training_model = get_training_model(train_opts, model)\n", "train(training_model, training_data)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Release IPU resources:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "training_model.detachFromDevice()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Inference" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let's define our test function." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def test(inference_model, test_data):\n", " nr_steps = len(test_data)\n", " sum_acc = 0.0\n", " for data, labels in tqdm(test_data, total=nr_steps):\n", " output = inference_model(data)\n", " sum_acc += accuracy(output, labels)\n", " print(f\"Accuracy on test set: {sum_acc / len(test_data):0.2f}%\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "After training, we test it using the inference model.\n", "Since weights in `model.parameters()` will be copied implicitly,\n", "the inference model takes the parameters from the trained model.\n", "We will also reuse the PopTorch options from the training model except we will disable gradient accumulation since we are not training.\n", "The inference batch-size was set to 8. To improve performance, we increase `device_iterations` to 100, which pushes more mini-batches into the inference pipeline." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "inf_options = train_opts.clone()\n", "inf_options.Training.gradientAccumulation(1)\n", "inf_options.deviceIterations(inference_iterations)\n", "inference_model = poptorch.inferenceModel(model, inf_options)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Finally, run the inference." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "test(inference_model, test_data)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Release IPU resources:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "inference_model.detachFromDevice()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## How to run the example script" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "1. Prepare the environment.\n", "\n", " Install the package requirements:\n", "\n", "``\n", "python3 -m pip install -r requirements.txt\n", "``\n", "\n", "2. Run the program. Note that the PopTorch Python API only supports Python 3.\n", "Data will be automatically downloaded using Torchvision utils.\n", "\n", "``\n", "python3 mnist_pipeline.py\n", "``\n", "\n", "The program has a few command-line options:\n", "\n", "- `-h` Show usage information.\n", "- `--batch-size` Sets the mini-batch size for training. Default: 40.\n", "- `--device-iterations` Number of device iterations for training. Default: 10.\n", "- `--test-batch-size` Sets the batch size for inference. Default: 8.\n", "- `--epochs` Number of epochs to train for. Default: 3.\n", "- `--learning-rate` Learning rate of the optimiser. Default (tuned for AdamW): 0.001.\n", "- `--profile` Enable profiling.\n", "- `--profile-dir` Select a directory to save the profile (default: ./profile).\n", "- `--profile-dir` Select a directory to save the profile. Default: ./profile.\n", "- `--strategy` Sets the execution strategy (pipelined, sharded). Default: pipelined.\n", "- `--debug` Prints out the debug logging while running.\n", "- `--offload-optimiser` Set this option to offload the optimiser state.\n", "\n", "Among these input arguments, `--strategy` and `--offload-optimiser`\n", "are associated with the execution methods.\n", "\n", "3. Some example command lines using different command-line options:\n", "\n", "```sh\n", "python3 mnist_pipeline.py --learning-rate 0.002 --debug --epochs 5\n", "\n", "python3 mnist_pipeline.py --profile --profile-dir pipeline_profile\n", "\n", "python3 mnist_pipeline.py --strategy sharded --profile --profile-dir sharding_profile\n", "\n", "python3 mnist_pipeline.py --offload-optimiser\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Conclusion" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "In this tutorial we achieved model parallelism on multiple IPUs with the PopTorch pipelining API.\n", "The model was partitioned to define the pipeline stages. Then, the pipelined execution strategy was set in the `poptorch.Options()` object.\n", "We generated profiles and used the Popvision Graph Analyser to visualise the execution trace and the memory usage on IPUs.\n", "The graph profiles help to identify imbalances between stages.\n", "This information is key to choosing a model partitioning that will reduce the time IPUs spend idle before synchronisation.\n", "Additionally, memory usage on the IPU was reduced at the expense of execution time by offloading optimiser-specific tensors to the host.\n", "This trade-off is worth considering if you need to fit memory intensive models on the IPU.\n", "\n", "To continue your journey with the Graphcore IPU you may want to look at the\n", "Graphcore [examples\n", "repository](https://github.com/graphcore/examples/tree/v2.6.0) which contains\n", "implementations of standard deep learning models optimised for the IPU. You will\n", "find, amongst others: [convolutional neural\n", "networks](https://github.com/graphcore/examples/tree/v2.6.0/vision/cnns),\n", "[BERT](https://github.com/graphcore/examples/tree/v2.6.0/nlp/bert/pytorch),\n", "[mini\n", "DALL-E](https://github.com/graphcore/examples/tree/v2.6.0/multimodal/mini_dalle/pytorch),\n", "and [ViT](https://github.com/graphcore/examples/tree/v2.6.0/vision/vit/pytorch)\n", "models.\n", "\n", "If you are looking to port and optimise your own model, the\n", "[memory and performance optimisation guide](https://docs.graphcore.ai/projects/memory-performance-optimisation/en/3.1.0/index.html)\n", "will help you make the most of the hardware.\n", "Finally the\n", "[TensorFlow pipelining tech note](https://docs.graphcore.ai/projects/tf-model-parallelism/en/3.1.0/pipelining.html#optimising-the-pipeline)\n", "contains additional information on the core concepts required to optimise your pipelined model." ] } ], "metadata": { "language_info": { "name": "python" }, "traceability": { "sdk_version": "3.0.0+1145", "source_file": "walkthrough.py", "sst_version": "0.0.8", "timestamp": "2022-09-27T15:36" } }, "nbformat": 4, "nbformat_minor": 5 }