repository·sdk-release-3.4·Indexed 41 days ago
https://github.com/graphcore/poptorchPyTorch extensions for training and evaluating models on Graphcore IPUs. Includes PopTorch Geometric for Graph Neural Network (GNN) workloads and a Dispatch Tracing system with JIT and MLIR backends for model tracing in inference mode.
PopTorch Geometric is a set of extensions for PyTorch Geometric designed to enable Graph Neural Network (GNN) models to be trained, evaluated, and used on Graphcore IPUs.
:warning: This project is under active development. All APIs should be considered volatile.
PopTorch Geometric is currently an experimental release. Users should be aware of the following:
PopTorch Geometric is a set of extensions for PyTorch Geometric designed to enable Graph Neural Network (GNN) models to be trained, evaluated, and used on Graphcore IPU hardware.
It is built to require minimal changes to existing PyTorch Geometric models, though it introduces some differences from native PyTorch Geometric execution to optimize performance for IPU hardware. PopTorch Geometric relies on the core functionality provided by PopTorch.
PopTorch uses a custom implementation of torch::jit::trace called Dispatch Tracing. This allows tracing autograd functions and sidestepping standard PyTorch JIT constraints. There are two available backends:
Key Constraint: Models can only be traced in inference mode. For training, PopART can optionally apply its own autograd to convert the traced inference graph into a training graph.
PopTorch uses a hierarchy of Blocks $\rightarrow$ Stages $\rightarrow$ Phases to manage execution:
user_id or stage_id.poptorch.AutoStage to automatically assign stages.AutoStage.SameAsIpu is true, stage_id matches the ipu_id.AutoStage.AutoIncrement is true, each block gets a unique, incrementing stage_id. This is useful in PipelinedExecution to avoid scheduling conflicts where the model flow contradicts the stage_id order.SerialPhasedExecution and ParallelPhasedExecution. A phase can contain one or more stages.In phased execution, weights and activations are transferred to/from Streaming Memory between phases. This allows running larger models that exceed IPU memory by trading off performance for capacity.
When training with half or float16 values, use loss scaling to prevent gradient underflow. PopTorch scales the loss by the loss_scaling parameter before gradient calculation and multiplies gradients by the inverse scale before the optimizer update. This process does not affect training hyper-parameters.
Options:
loss_scaling parameter directly on the optimizer.poptorch.options._TrainingOptions.setAutomaticLossScaling to enable a global automatic scaling factor.Note: If automatic loss scaling causes convergence issues or compilation errors, switch to manual scaling.
When slicing a tensor in a PopTorch model, you must satisfy one of two conditions for the operation to be supported on IPU hardware:
tensor[x:x+5]).Note on behavior: If the end value of a slice exceeds the length of the tensor, PopTorch behaves differently than PyTorch. While PyTorch returns a smaller tensor, PopTorch allows the slice to wrap around to the start of the relevant dimension.
PopTorch supports PyTorch buffers to make tensors persistent across runs without making them model parameters.
Important Requirements:
+= or methods ending in _ like torch.Tensor.copy_) to modify the buffer. Standard assignments (e.g., model.i += 1) will fail because the PyTorch dispatcher captures the value during graph construction and freezes it as a constant.poptorch.inferenceModel, buffers modified by the model are not automatically copied to the host. You must call poptorch.PoplarExecutor.copyWeightsToHost before reading the updated buffer values.broadcastBuffers option to False using poptorch.Options().broadcastBuffers(False).# Correct way to use a persistent buffer
class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.register_buffer('i', torch.zeros(1))
def forward(self, x):
self.i += 1 # Must be an in-place operation
return xPopTorch optimizes performance by embedding constant attributes directly in the program.
For poptorch.optim optimizers:
poptorch.optim.VariableAttributes.markAsConstant or poptorch.optim.VariableAttributes.markAsVariable before model compilation.For native torch.optim optimizers:
poptorch.optim instead.# Constant and variable attributes for PopTorch optimizers
optimizer = poptorch.optim.Adam(model.parameters(), lr=0.01)
optimizer.markAsConstant('lr') # This would make lr a constantIn PopTorch, the optimizer state is encapsulated within the poptorch_model.
Important considerations:
poptorch_model.setOptimizer(updated_optimizer) to synchronize the state.poptorch.optim. These are wrapper classes that offer performance/memory advantages and allow for parameters like loss scaling and velocity scaling.# If you update the optimizer externally, sync it with the model:
poptorch_model.setOptimizer(new_optimizer)When using poptorch.trainingModel, the training workflow changes significantly from native PyTorch due to PopTorch's custom automatic differentiation (autograd) engine:
torch.nn.Module must have at least one loss value built into its forward pass. The model should return the loss value(s) as part of its output.loss.backward(). PopTorch automatically detects the loss value and backpropagates gradients from it.torch.optim.AdamW) as an argument to poptorch.trainingModel. PopTorch manages the optimizer state and applies the algorithm automatically; you do not call optimizer.step() in your training loop.poptorch.DataLoader class for efficient data batching on the IPU.PopTorch provides several abstractions to define how computation is distributed across IPUs:
poptorch.PipelinedExecution: For pipeline parallelism.poptorch.ShardedExecution: For sharding parameters/data.poptorch.SerialPhasedExecution and poptorch.ParallelPhasedExecution: For controlling execution phases.poptorch.Block and poptorch.BeginBlock: Define computation blocks.poptorch.Stage and poptorch.AutoStage: Define stages in a pipeline.poptorch.Phase: Defines specific phases of execution.poptorch.replicaGrouping(tensor, comm_group_type, shards, variable_retrieval_mode) on a weight tensor (after wrapping with trainingModel or inferenceModel) to configure how replicas receive different values of that tensor.