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) – 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.

abstract property root_device: torch.device

Returns the root device.