1. Introduction

PopTorch is a set of extensions for PyTorch to enable PyTorch models to run directly on Graphcore IPU hardware.

PopTorch supports both inference and training. To run a model on the IPU, you wrap your existing PyTorch model in either a PopTorch inference wrapper or a PopTorch training wrapper. You can provide further annotations to partition the model across multiple IPUs or execution phases.

You can wrap individual layers in an IPU helper to designate which IPU they should go on. Using the user-provided annotations, PopTorch will use PopART to parallelise the model over the given number of IPUs. Additional parallelism can be expressed via a replication factor which enables you to data-parallelise the model over more IPUs.

PopTorch has been designed to require as few changes as possible to your models in order to run them on IPU. However, it does have some differences from native PyTorch execution, to get the most out of IPU hardware.

  • In training mode, PopTorch has its own automatic differentiation (autograd) and will take the model (torch.nn.Module) with one or more losses built-in and an optimizer, and will perform the full forward and backward pass under the hood.

Listing 1.1 A simple example of training using PyTorch on the CPU
 1training_data = torch.utils.data.DataLoader(ExampleDataset(shape=[1],
 2                                                           length=20000),
 3                                            batch_size=10,
 4                                            shuffle=True,
 5                                            drop_last=True)
 6
 7model = ExampleModelWithLoss()
 8model.train()
 9
10optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
11
12momentum_loss = None
13
14for batch, target in training_data:
15    # Zero gradients
16    optimizer.zero_grad()
17
18    # Run model.
19    _, loss = model(batch, target)
20
21    # Back propagate the gradients.
22    loss.backward()
23
24    # Update the weights.
25    optimizer.step()
26
27    if momentum_loss is None:
28        momentum_loss = loss
29    else:
30        momentum_loss = momentum_loss * 0.95 + loss * 0.05
31
32    if momentum_loss < 0.1:
33        optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001)
Listing 1.2 Equivalent code using PopTorch to train on the IPU
 1# Set up the PyTorch DataLoader to load that much data at each iteration
 2opts = poptorch.Options()
 3opts.deviceIterations(10)
 4training_data = poptorch.DataLoader(options=opts,
 5                                    dataset=ExampleDataset(shape=[1],
 6                                                           length=20000),
 7                                    batch_size=10,
 8                                    shuffle=True,
 9                                    drop_last=True)
10
11model = ExampleModelWithLoss()
12model.train()
13
14optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
15
16# Wrap the model in a PopTorch training wrapper
17poptorch_model = poptorch.trainingModel(model,
18                                        options=opts,
19                                        optimizer=optimizer)
20
21momentum_loss = None
22
23for batch, target in training_data:
24    # Performs forward pass, loss function evaluation,
25    # backward pass and weight update in one go on the device.
26    _, loss = poptorch_model(batch, target)
27
28    if momentum_loss is None:
29        momentum_loss = loss
30    else:
31        momentum_loss = momentum_loss * 0.95 + loss * 0.05
32
33    # Optimizer can be updated via setOptimizer.
34    if momentum_loss < 0.1:
35        poptorch_model.setOptimizer(
36            torch.optim.AdamW(model.parameters(), lr=0.0001))

  • Under the hood, PopTorch uses PyTorch’s torch.jit.trace API. That means it inherits the constraints of that API. These include:

    • Inputs must be PyTorch tensors or tuples containing PyTorch tensors.

    • None can be used as a default value for a parameter but cannot be explicitly passed as an input value.

    • torch.jit.trace cannot handle control flow or shape variations within the model. That is, the inputs passed at run-time cannot vary the control flow of the model or the shapes/sizes of results. If you attempt this, the graph will be frozen to whichever control flow path was traced as a result of the first inputs given to the wrapped model.

    • Not all PyTorch operations have been implemented by the backend yet.

  • PopTorch follows PopART’s data batching semantics. By default, this means you will just pass in data of the normal batch size. However, there are a number of options provided in PopTorch which will enable more efficient data loading. See Efficient data batching for more information.