IPUStrategy

class lightning.pytorch.strategies.IPUStrategy(accelerator=None, device_iterations=1, autoreport=False, autoreport_dir=None, parallel_devices=None, cluster_environment=None, checkpoint_io=None, precision_plugin=None, training_opts=None, inference_opts=None)

Bases: ParallelStrategy

Plugin for training on IPU devices.

Warning

This is an experimental feature.

Parameters

• device_iterations – Number of iterations to run on device at once before returning to host. This can be used as an optimization to speed up training. https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/batching.html

• autoreport – Enable auto-reporting for IPUs using PopVision https://docs.graphcore.ai/projects/graphcore-popvision-user-guide/en/latest/graph/graph.html

• autoreport_dir – Optional directory to store autoReport output.

• training_opts – Optional poptorch.Options to override the default created options for training.

• inference_opts – Optional poptorch.Options to override the default created options for validation/testing and predicting.

all_gather(tensor, group=None, sync_grads=False)

Perform a all_gather on all processes.

Return type

Tensor

barrier(name=None)

Synchronizes all processes which blocks processes until the whole group enters this function.

Parameters

name (Optional[str]) – an optional name to pass into barrier.

Return type

None

batch_to_device(batch, device=None, dataloader_idx=0)

Moves the batch to the correct device.

The returned batch is of the same type as the input batch, just having all tensors on the correct device.

Parameters

• batch (Any) – The batch of samples to move to the correct device

• device (Optional[device]) – The target device

• dataloader_idx (int) – The index of the dataloader to which the batch belongs.

Return type

Any

broadcast(obj, src=0)

Broadcasts an object to all processes.

Parameters

• obj (TypeVar(TBroadcast)) – the object to broadcast

• src (int) – source rank

Return type

TypeVar(TBroadcast)

model_to_device()

Moves the model to the correct device.

Return type

None

on_predict_end()

Called when predict ends.

Return type

None

on_predict_start()

Called when predict begins.

Return type

None

on_test_end()

Called when test end.

Return type

None

on_test_start()

Called when test begins.

Return type

None

on_train_batch_start(batch, batch_idx)

Called in the training loop before anything happens for that batch.

Return type

None

on_train_end()

Called when train ends.

Return type

None

on_train_start()

Called when train begins.

Return type

None

on_validation_end()

Called when validation ends.

Return type

None

on_validation_start()

Called when validation begins.

Return type

None

predict_step(*args, **kwargs)

The actual predict step.

See predict_step() for more details

Return type

Union[Tensor, Dict[str, Any]]

reduce(tensor, *args, **kwargs)

Reduces the given tensor (e.g. across GPUs/processes).

Parameters

• tensor (Union[Tensor, Any]) – the tensor to sync and reduce

• group – the process group to reduce

• reduce_op – the reduction operation. Defaults to ‘mean’. Can also be a string ‘sum’ or ReduceOp.

Return type

Union[Tensor, Any]

setup(trainer)

Setup plugins for the trainer fit and creates optimizers.

Parameters

trainer (Trainer) – the trainer instance

Return type

None

setup_optimizers(trainer)

Creates optimizers and schedulers.

Parameters

trainer (Trainer) – the Trainer, these optimizers should be connected to

Return type

None

teardown()

This method is called to teardown the training process.

It is the right place to release memory and free other resources.

Return type

None

test_step(*args, **kwargs)

The actual test step.

See test_step() for more details

Return type

Union[Tensor, Dict[str, Any], None]

training_step(*args, **kwargs)

The actual training step.

See training_step() for more details

Return type

Union[Tensor, Dict[str, Any]]

validation_step(*args, **kwargs)

The actual validation step.

See validation_step() for more details

Return type

Union[Tensor, Dict[str, Any], None]

property is_global_zero: bool

Whether the current process is the rank zero process not only on the local node, but for all nodes.

Return type

bool

property root_device: device

Return the root device.

Return type

device