Strategy

class lightning_fabric.strategies.Strategy(accelerator=None, checkpoint_io=None, precision=None)

Bases: abc.ABC

Base class for all strategies that change the behaviour of the training, validation and test- loop.

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

Perform an all_gather on all processes.

Parameters

• tensor (Tensor) – the tensor to all_gather

• group (Optional[Any]) – the process group to gather results from

• sync_grads (bool) – flag that allows users to synchronize gradients for all_gather op

Return type

Tensor

abstract all_reduce(tensor, group=None, reduce_op='mean')

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

Parameters

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

• group (Optional[Any]) – the process group to reduce

• reduce_op (Union[ReduceOp, str, None]) – the reduction operation. Defaults to ‘mean’. Can also be a string ‘sum’ or ReduceOp.

Return type

Union[Tensor, Any]

backward(tensor, module, *args, **kwargs)

Forwards backward-calls to the precision plugin.

Return type

None

abstract 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)

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

Return type

Any

abstract 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)

get_module_state_dict(module)

Returns model state.

Return type

Dict[str, Union[Any, Tensor]]

get_optimizer_state(optimizer)

Returns state of an optimizer.

Allows for syncing/collating optimizer state from processes in custom plugins.

Return type

Dict[str, Tensor]

abstract module_to_device(module)

Moves the model to the correct device.

Return type

None

optimizer_step(optimizer, **kwargs)

Performs the actual optimizer step.

Parameters

• optimizer (Optimizable) – the optimizer performing the step

• **kwargs – Any extra arguments to optimizer.step

Return type

Any

process_dataloader(dataloader)

Wraps the dataloader if necessary.

Parameters

dataloader (DataLoader) – iterable. Ideally of type: torch.utils.data.DataLoader

Return type

DataLoader

reduce_boolean_decision(decision, all=True)

Reduce a boolean decision across all processes.

Return type

bool

remove_checkpoint(filepath)

Remove checkpoint filepath from the filesystem.

Parameters

filepath (Union[str, Path]) – Path to checkpoint

Return type

None

save_checkpoint(checkpoint, filepath, storage_options=None)

Save model/training states as a checkpoint file through state-dump and file-write.

Parameters

• checkpoint (Dict[str, Any]) – dict containing model and trainer state

• filepath (Union[str, Path]) – write-target file’s path

• storage_options (Optional[Any]) – parameter for how to save to storage, passed to CheckpointIO plugin

Return type

None

setup_environment()

Setup any processes or distributed connections.

This must be called by the framework at the beginning of every process, before any distributed communication takes place.

Return type

None

setup_module(module)

Performs setup for the model, e.g., by wrapping it by another class.

Return type

Module

setup_module_and_optimizers(module, optimizers)

Set up a model and multiple optimizers together.

The returned objects are expected to be in the same order they were passed in. The default implementation will call setup_module() and setup_optimizer() on the inputs.

Return type

Tuple[Module, List[Optimizer]]

setup_optimizer(optimizer)

Performs setup for the optimizer, e.g., by wrapping it by another class.

Return type

Optimizer

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

abstract 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

abstract property root_device: torch.device

Returns the root device.

Return type

device