Fabric

class lightning_fabric.fabric.Fabric(accelerator=None, strategy=None, devices=None, num_nodes=1, precision=32, plugins=None, callbacks=None, loggers=None)

Bases: object

Fabric accelerates your PyTorch training or inference code with minimal changes required.

• Automatic placement of models and data onto the device.

• Automatic support for mixed and double precision (smaller memory footprint).

• Seamless switching between hardware (CPU, GPU, TPU) and distributed training strategies (data-parallel training, sharded training, etc.).

• Automated spawning of processes, no launch utilities required.

• Multi-node support.

Parameters

• accelerator (Union[str, Accelerator, None]) – The hardware to run on. Possible choices are: "cpu", "cuda", "mps", "gpu", "tpu", "auto".

• strategy (Union[str, Strategy, None]) – Strategy for how to run across multiple devices. Possible choices are: "dp", "ddp", "ddp_spawn", "deepspeed", "fsdp".

• devices (Union[int, str, List[int], None]) – Number of devices to train on (int), which GPUs to train on (list or str), or "auto". The value applies per node.

• num_nodes (int) – Number of GPU nodes for distributed training.

• precision (Union[Literal[64, 32, 16], Literal[‘64’, ‘32’, ‘16’, ‘bf16’]]) – Double precision (64), full precision (32), half precision (16), or bfloat16 precision ("bf16").

• plugins (Union[Precision, ClusterEnvironment, CheckpointIO, str, List[Union[Precision, ClusterEnvironment, CheckpointIO, str]], None]) – One or several custom plugins

• callbacks (Union[List[Any], Any, None]) – A single callback or a list of callbacks. A callback can contain any arbitrary methods that can be invoked through call() by the user.

• loggers (Union[Logger, List[Logger], None]) – A single logger or a list of loggers. See log() for more information.

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

Gather tensors or collections of tensors from multiple processes.

Parameters

• data (Union[Tensor, Dict, List, Tuple]) – int, float, tensor of shape (batch, …), or a (possibly nested) collection thereof.

• group (Optional[Any]) – the process group to gather results from. Defaults to all processes (world)

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

Return type

Union[Tensor, Dict, List, Tuple]

Returns

A tensor of shape (world_size, batch, …), or if the input was a collection the output will also be a collection with tensors of this shape.

autocast()

A context manager to automatically convert operations for the chosen precision.

Use this only if the forward method of your model does not cover all operations you wish to run with the chosen precision setting.

Return type

Generator[None, None, None]

backward(tensor, *args, model=None, **kwargs)

Replaces loss.backward() in your training loop. Handles precision and automatically for you.

Parameters

• tensor (Tensor) – The tensor (loss) to back-propagate gradients from.

• *args – Optional positional arguments passed to the underlying backward function.

• model (Optional[_FabricModule]) – Optional model instance for plugins that require the model for backward().

• **kwargs – Optional named keyword arguments passed to the underlying backward function.

Note

When using strategy="deepspeed" and multiple models were set up, it is required to pass in the model as argument here.

Return type

None

barrier(name=None)

Wait for all processes to enter this call. Use this to synchronize all parallel processes, but only if necessary, otherwise the overhead of synchronization will cause your program to slow down.

Example:

if self.global_rank == 0:
    # let process 0 download the dataset
    dataset.download_files()

# let all processes wait before reading the dataset
self.barrier()

# now all processes can read the files and start training

Return type

None

call(hook_name, *args, **kwargs)

Trigger the callback methods with the given name and arguments.

Not all objects registered via Fabric(callbacks=...) must implement a method with the given name. The ones that have a matching method name will get called.

Parameters

• hook_name (str) – The name of the callback method.

• *args – Optional positional arguments that get passed down to the callback method.

• **kwargs – Optional keyword arguments that get passed down to the callback method.

Example:

class MyCallback:
    def on_train_epoch_end(self, results):
        ...

fabric = Fabric(callbacks=[MyCallback()])
fabric.call("on_train_epoch_end", results={...})

Return type

None

load(filepath)

Load a checkpoint from a file.

How and which processes load gets determined by the strategy

Parameters

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

Return type

Any

log(name, value, step=None)

Log a scalar to all loggers that were added to Fabric.

Parameters

• name (str) – The name of the metric to log.

• value (Any) – The metric value to collect. If the value is a torch.Tensor, it gets detached from the graph automatically.

• step (Optional[int]) – Optional step number. Most Logger implementations auto-increment the step value by one with every log call. You can specify your own value here.

Return type

None

log_dict(metrics, step=None)

Log multiple scalars at once to all loggers that were added to Fabric.

Parameters

• metrics (Mapping[str, Any]) – A dictionary where the key is the name of the metric and the value the scalar to be logged. Any torch.Tensor in the dictionary get detached from the graph automatically.

• step (Optional[int]) – Optional step number. Most Logger implementations auto-increment this value by one with every log call. You can specify your own value here.

Return type

None

no_backward_sync(module, enabled=True)

Skip gradient synchronization during backward to avoid redundant communication overhead.

Use this context manager when performing gradient accumulation to speed up training with multiple devices.

Example:

# Accumulate gradient 8 batches at a time
with self.no_backward_sync(model, enabled=(batch_idx % 8 != 0)):
    output = model(input)
    loss = ...
    self.backward(loss)
    ...

For those strategies that don’t support it, a warning is emitted. For single-device strategies, it is a no-op. Both the model’s .forward() and the self.backward() call need to run under this context.

Parameters

• module (_FabricModule) – The module for which to control the gradient synchronization.

• enabled (bool) – Whether the context manager is enabled or not. True means skip the sync, False means do not skip.

Return type

Generator

print(*args, **kwargs)

Print something only on the first process.

Arguments passed to this method are forwarded to the Python built-in print() function.

Return type

None

run(*args, **kwargs)

All the code inside this run method gets accelerated by Fabric.

You can pass arbitrary arguments to this function when overriding it.

Return type

Any

save(content, filepath)

Save checkpoint contents to a file.

How and which processes save gets determined by the strategy. For example, the ddp strategy saves checkpoints only on process 0.

Parameters

• content (Dict[str, Any]) – A dictionary with contents, i.e., the state dict of your model

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

Return type

None

static seed_everything(seed=None, workers=None)

Helper function to seed everything without explicitly importing Lightning.

See pytorch_lightning.seed_everything() for more details.

Return type

int

setup(module, *optimizers, move_to_device=True)

Set up a model and its optimizers for accelerated training.

Parameters

• module (Module) – A torch.nn.Module to set up

• *optimizers – The optimizer(s) to set up (no optimizers is also possible)

• move_to_device (bool) – If set True (default), moves the model to the correct device. Set this to False and alternatively use to_device() manually.

Return type

Any

Returns

The tuple containing wrapped module and the optimizers, in the same order they were passed in.

setup_dataloaders(*dataloaders, replace_sampler=True, move_to_device=True)

Set up one or multiple dataloaders for accelerated training. If you need different settings for each dataloader, call this method individually for each one.

Parameters

• *dataloaders – A single dataloader or a sequence of dataloaders.

• replace_sampler (bool) – If set True (default), automatically wraps or replaces the sampler on the dataloader(s) for distributed training. If you have a custom sampler defined, set this to this argument to False.

• move_to_device (bool) – If set True (default), moves the data returned by the dataloader(s) automatically to the correct device. Set this to False and alternatively use to_device() manually on the returned data.

Return type

Union[DataLoader, List[DataLoader]]

Returns

The wrapped dataloaders, in the same order they were passed in.

setup_module(module, move_to_device=True)

Set up a model for accelerated training or inference.

This is the same as calling .setup(model) with no optimizers. It is useful for inference or for certain strategies like FSDP that require setting up the module before the optimizer can be created and set up. See also setup_optimizers().

Parameters

• module (Module) – A torch.nn.Module to set up

• move_to_device (bool) – If set True (default), moves the model to the correct device. Set this to False and alternatively use to_device() manually.

Return type

_FabricModule

Returns

The wrapped model.

setup_optimizers(*optimizers)

Set up one or more optimizers for accelerated training.

Some strategies do not allow setting up model and optimizer independently. For them, you should call .setup(model, optimizer, ...) instead to jointly set them up.

Parameters

*optimizers – One or more optmizers to set up.

Return type

Union[_FabricOptimizer, Tuple[_FabricOptimizer, ...]]

Returns

The wrapped optimizer(s).

sharded_model()

Shard the parameters of the model instantly when instantiating the layers.

Use this context manager with strategies that support sharding the model parameters to save peak memory usage.

Example:

with self.sharded_model():
    model = MyModel()

The context manager is strategy-agnostic and for the ones that don’t do sharding, it is a no-op.

Return type

Generator

to_device(obj: torch.nn.modules.module.Module) → torch.nn.modules.module.Module

to_device(obj: torch.Tensor) → torch.Tensor

to_device(obj: Any) → Any

Move a torch.nn.Module or a collection of tensors to the current device, if it is not already on that device.

Parameters

obj (Union[Module, Tensor, Any]) – An object to move to the device. Can be an instance of torch.nn.Module, a tensor, or a (nested) collection of tensors (e.g., a dictionary).

Return type

Union[Module, Tensor, Any]

Returns

A reference to the object that was moved to the new device.

property device: torch.device

The current device this process runs on.

Use this to create tensors directly on the device if needed.

Return type

device

property global_rank: int

The global index of the current process across all devices and nodes.

Return type

int

property is_global_zero: bool

Whether this rank is rank zero.

Return type

bool

property local_rank: int

The index of the current process among the processes running on the local node.

Return type

int

property logger: lightning_fabric.loggers.logger.Logger

Returns the first logger in the list passed to Fabric, which is considered the main logger.

Return type

Logger

property loggers: List[lightning_fabric.loggers.logger.Logger]

Returns all loggers passed to Fabric.

Return type

List[Logger]

property node_rank: int

The index of the current node.

Return type

int

property world_size: int

The total number of processes running across all devices and nodes.

Return type

int