--- title: XLAStrategy description: >- Strategy for training multiple TPU devices using the torch_xla.distributed.xla_multiprocessing.spawn() method. canonical: >- https://lightning.ai/docs/fabric/stable/full-api-reference/api/generated/lightning.fabric.strategies.XLAStrategy --- # XLAStrategy {#xlastrategy} ### `class lightning.fabric.strategies.XLAStrategy(accelerator=None, parallel_devices=None, checkpoint_io=None, precision=None, sync_module_states=True)` {#lightning.fabric.strategies.XLAStrategy} Bases: [`ParallelStrategy`](/docs/fabric/2.6.6/full-api-reference/api/generated/lightning.fabric.strategies.ParallelStrategy.md#lightning.fabric.strategies.ParallelStrategy "lightning.fabric.strategies.parallel.ParallelStrategy") Strategy for training multiple TPU devices using the [`torch_xla.distributed.xla_multiprocessing.spawn()`](https://docs.pytorch.org/xla/release/2.0/index.html#torch_xla.distributed.xla_multiprocessing.spawn "(in PyTorch/XLA vmaster (2.0 ))") method. #### `_configure_launcher()` {#lightning.fabric.strategies.XLAStrategy._configure_launcher} Attach the launcher based on Strategy. Return type: [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)") #### `all_gather(tensor, group=None, sync_grads=False)` {#lightning.fabric.strategies.XLAStrategy.all_gather} Function to gather a tensor from several distributed processes. Parameters: * **tensor** ([`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)")) – tensor to all-gather. * **group** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – unused. * **sync\_grads** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – flag that allows users to synchronize gradients for the all-gather operation. Return type: [`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)") Returns: A tensor of shape (world\_size, …) #### `all_reduce(output, group=None, reduce_op=None)` {#lightning.fabric.strategies.XLAStrategy.all_reduce} Reduces the given tensor (e.g. across GPUs/processes). Parameters: * **tensor** – the tensor to sync and reduce * **group** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – the process group to reduce * **reduce\_op** (`Union`\[[`ReduceOp`](https://docs.pytorch.org/docs/stable/distributed.html#torch.distributed.ReduceOp "(in PyTorch v2.15)"), [`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – the reduction operation. Defaults to ‘mean’. Can also be a string ‘sum’ or ReduceOp. Return type: [`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)") #### `barrier(name=None, *args, **kwargs)` {#lightning.fabric.strategies.XLAStrategy.barrier} Synchronizes all processes which blocks processes until the whole group enters this function. Parameters: **name** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")\]) – an optional name to pass into barrier. Return type: [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)") #### `broadcast(obj, src=0)` {#lightning.fabric.strategies.XLAStrategy.broadcast} Broadcasts an object to all processes. Parameters: * **obj** ([`TypeVar`](https://docs.python.org/3/library/typing.html#typing.TypeVar "(in Python v3.15)")(`TBroadcast`)) – the object to broadcast * **src** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – source rank Return type: [`TypeVar`](https://docs.python.org/3/library/typing.html#typing.TypeVar "(in Python v3.15)")(`TBroadcast`) #### `module_to_device(module)` {#lightning.fabric.strategies.XLAStrategy.module_to_device} Moves the model to the correct device. Return type: [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)") #### `process_dataloader(dataloader)` {#lightning.fabric.strategies.XLAStrategy.process_dataloader} Wraps the dataloader if necessary. Parameters: **dataloader** ([`DataLoader`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "(in PyTorch v2.15)")) – iterable. Ideally of type: [`torch.utils.data.DataLoader`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "(in PyTorch v2.15)") Return type: `MpDeviceLoader` #### `save_checkpoint(path, state, storage_options=None, filter=None)` {#lightning.fabric.strategies.XLAStrategy.save_checkpoint} Save model, optimizer, and other state as a checkpoint file. Parameters: * **path** (`Union`\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Path`](https://docs.python.org/3/library/pathlib.html#pathlib.Path "(in Python v3.15)")\]) – A path to where the file(s) should be saved * **state** ([`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), `Union`\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`Optimizer`](https://docs.pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]\]) – 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`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – Additional options for the `CheckpointIO` plugin * **filter** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Callable`](https://docs.python.org/3/library/typing.html#typing.Callable "(in Python v3.15)")\[\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\], [`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")\]\]\]) – An optional dictionary of the same format as `state` mapping keys to callables that return a boolean indicating whether the given parameter should be saved (`True`) or filtered out (`False`). Return type: [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)") #### `setup_environment()` {#lightning.fabric.strategies.XLAStrategy.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`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)") #### `setup_module(module)` {#lightning.fabric.strategies.XLAStrategy.setup_module} Performs setup for the model, e.g., by wrapping it by another class. Return type: [`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)") #### `property root_device: device` {#lightning.fabric.strategies.XLAStrategy.root_device} Returns the root device. --- **Navigation** - Parent: [Generated](/docs/fabric/2.6.6/full-api-reference/api/generated.md)