---
title: XLAStrategy
description: >-
  Strategy for training multiple TPU devices using the
  torch_xla.distributed.xla_multiprocessing.spawn()

  method.
canonical: >-
  https://lightning.ai/docs/pytorch/stable/optional-api/api/lightning.pytorch.strategies.XLAStrategy
---
# XLAStrategy {#xlastrategy}

### `class lightning.pytorch.strategies.XLAStrategy(accelerator=None, parallel_devices=None, checkpoint_io=None, precision_plugin=None, debug=False, sync_module_states=True, **_)` {#lightning.pytorch.strategies.XLAStrategy}

Bases: [`DDPStrategy`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.strategies.DDPStrategy.md#lightning.pytorch.strategies.DDPStrategy "lightning.pytorch.strategies.ddp.DDPStrategy")

Strategy for training multiple TPU devices using the `torch_xla.distributed.xla_multiprocessing.spawn()` method.

#### `all_gather(tensor, group=None, sync_grads=False)` {#lightning.pytorch.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, …)

#### `barrier(name=None, *args, **kwargs)` {#lightning.pytorch.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.pytorch.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`)

#### `model_to_device()` {#lightning.pytorch.strategies.XLAStrategy.model_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.pytorch.strategies.XLAStrategy.process_dataloader}

Wraps the dataloader if necessary.

Parameters:

**dataloader** ([`object`](https://docs.python.org/3/builtins/functions.html#object "(in Python v3.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`

#### `reduce(output, group=None, reduce_op='mean')` {#lightning.pytorch.strategies.XLAStrategy.reduce}

Reduces a tensor from several distributed processes to one aggregated tensor.

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 gather results from. Defaults to all processes (world)
    
*   **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’/’avg’. Can also be a string ‘sum’ to calculate the sum during reduction.
    

Return type:

[`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)")

Returns:

reduced value, except when the input was not a tensor the output remains is unchanged

#### `remove_checkpoint(filepath)` {#lightning.pytorch.strategies.XLAStrategy.remove_checkpoint}

Remove checkpoint filepath from the filesystem.

Parameters:

**filepath** (`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)")\]) – Path to checkpoint

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `save_checkpoint(checkpoint, filepath, storage_options=None)` {#lightning.pytorch.strategies.XLAStrategy.save_checkpoint}

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

Parameters:

*   **checkpoint** ([`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)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – dict containing model and trainer state
    
*   **filepath** (`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)")\]) – write-target file’s path
    
*   **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)")\]) – parameter for how to save to storage, passed to `CheckpointIO` plugin
    

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `setup(trainer)` {#lightning.pytorch.strategies.XLAStrategy.setup}

Sets up the accelerator, plugins and initializes the optimizers (if needed).

Parameters:

**trainer** ([`Trainer`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md#lightning.pytorch.trainer.trainer.Trainer "lightning.pytorch.trainer.trainer.Trainer")) – the trainer instance

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `setup_environment()` {#lightning.pytorch.strategies.XLAStrategy.setup_environment}

Setup any processes or distributed connections.

This is called before the LightningModule/DataModule setup hook which allows the user to access the accelerator environment before setup is complete.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `teardown()` {#lightning.pytorch.strategies.XLAStrategy.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`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `property root_device: device` {#lightning.pytorch.strategies.XLAStrategy.root_device}

Return the root device.

---

**Navigation**

- Parent: [Api](/docs/pytorch/2.6.6/optional-api/api.md)