Strategy

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

Bases: ABC

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

_configure_launcher()

Attach the launcher based on Strategy.

Return type:

None

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)

clip_gradients_norm(module, optimizer, max_norm, norm_type=2.0, error_if_nonfinite=True)

Clip gradients by norm.

Return type:

Tensor

clip_gradients_value(module, optimizer, clip_val)

Clip gradients by value.

Return type:

None

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]

load_checkpoint(path, state=None, strict=True)

Load the contents from a checkpoint and restore the state of the given objects.

Parameters:

• path (Union[str, Path]) – A path to where the file is located

• state (Union[Module, Optimizer, dict[str, Union[Module, Optimizer, Any]], None]) –

  Can be one of:

  • A dictionary of objects whose state will be restored in-place from the checkpoint path.

  • None or the empty dict: The loaded checkpoint will be returned in full.

  • A Module instance, if the checkpoint file contains a raw module state dict.

  • A Optimizer instance, if the checkpoint file contains a raw optimizer state.

• strict (bool) – Whether to enforce that the keys in state match the keys in the checkpoint.

Return type:

dict[str, Any]

Returns:

The remaining items that were not restored into the given state dictionary. If no state dictionary is given, the full checkpoint will be returned.

load_module_state_dict(module, state_dict, strict=True)

Loads the given state into the model.

Return type:

None

module_init_context(empty_init=None)

A context manager wrapping the model instantiation.

Here, the strategy can control how the parameters of the model get created (device, dtype) and or apply other patches to the model.

Parameters:

empty_init (Optional[bool]) – Whether to initialize the model with empty weights (uninitialized memory). If None, the strategy will decide. Some strategies may not support all options.

Return type:

AbstractContextManager

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

save_checkpoint(path, state, storage_options=None, filter=None)

Save model, optimizer, and other state as a checkpoint file.

Parameters:

• path (Union[str, Path]) – A path to where the file(s) should be saved

• state (dict[str, Union[Module, Optimizer, Any]]) – A dictionary with contents to be saved. If the dict contains modules or optimizers, their state-dict will be retrieved and converted automatically.

• storage_options (Optional[Any]) – Additional options for the CheckpointIO plugin

• filter (Optional[dict[str, Callable[[str, Any], bool]]]) – An optional dictionary containing filter callables that return a boolean indicating whether the given item should be saved (True) or filtered out (False). Each filter key should match a state key, where its filter will be applied to the state_dict generated.

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, scheduler=None)

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], Optional[_LRScheduler]]

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

tensor_init_context()

Controls how tensors get created (device, dtype).

Return type:

AbstractContextManager

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: device

Returns the root device.