---
title: XLAFSDPStrategy
description: >-
  Strategy for training multiple XLA devices using the

  torch_xla.distributed.xla_fully_sharded_data_parallel.XlaFullyShardedDataParallel()
  method.
canonical: >-
  https://lightning.ai/docs/fabric/stable/full-api-reference/api/generated/lightning.fabric.strategies.XLAFSDPStrategy
---
# XLAFSDPStrategy {#xlafsdpstrategy}

### `class lightning.fabric.strategies.XLAFSDPStrategy(accelerator=None, parallel_devices=None, checkpoint_io=None, precision=None, auto_wrap_policy=None, activation_checkpointing_policy=None, state_dict_type='sharded', sequential_save=False, **kwargs)` {#lightning.fabric.strategies.XLAFSDPStrategy}

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"), `_Sharded`

Strategy for training multiple XLA devices using the `torch_xla.distributed.xla_fully_sharded_data_parallel.XlaFullyShardedDataParallel()` method.

Warning

This is an [experimental](/docs/pytorch/stable/versioning#experimental-api "(in PyTorch Lightning v2.7.0.dev)") feature.

For more information check out [https://github.com/pytorch/xla/blob/v2.5.0/docs/fsdp.md](https://github.com/pytorch/xla/blob/v2.5.0/docs/fsdp.md)

Parameters:

*   **auto\_wrap\_policy** (`Union`\[[`set`](https://docs.python.org/3/builtins/stdtypes.html#set "(in Python v3.15)")\[[`type`](https://docs.python.org/3/builtins/functions.html#type "(in Python v3.15)")\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)")\]\], [`Callable`](https://docs.python.org/3/library/typing.html#typing.Callable "(in Python v3.15)")\[\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)"), [`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")\], [`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")\], [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – Same as `auto_wrap_policy` parameter in `torch_xla.distributed.fsdp.XlaFullyShardedDataParallel`. For convenience, this also accepts a set of the layer classes to wrap.
    
*   **activation\_checkpointing\_policy** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`set`](https://docs.python.org/3/builtins/stdtypes.html#set "(in Python v3.15)")\[[`type`](https://docs.python.org/3/builtins/functions.html#type "(in Python v3.15)")\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)")\]\]\]) – Used when selecting the modules for which you want to enable activation checkpointing. Enabling this can free up a significant amount of memory at the cost of speed since activations in these layers need to be recomputed during backpropagation. This accepts a set of the layer classes to wrap.
    
*   **state\_dict\_type** ([`Literal`](https://docs.python.org/3/library/typing.html#typing.Literal "(in Python v3.15)")\[`'full'`, `'sharded'`\]) –
    
    The format in which the state of the model and optimizers gets saved into the checkpoint.
    
    *   `"full"`: The full weights and optimizer states get assembled on rank 0 and saved to a single file.
        
    *   `"sharded"`: Each rank saves its shard of weights and optimizer states to a file. The checkpoint is a folder with files for each shard in the host. Note that TPU VM multihost does not have a shared filesystem.
        
    
*   **sequential\_save** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – With this enabled, individual ranks consecutively save their state dictionary shards, reducing peak system RAM usage, although it elongates the saving process.
    
*   **\*\*kwargs** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – See available parameters in `torch_xla.distributed.fsdp.XlaFullyShardedDataParallel`.
    

#### `_configure_launcher()` {#lightning.fabric.strategies.XLAFSDPStrategy._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.XLAFSDPStrategy.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.XLAFSDPStrategy.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.XLAFSDPStrategy.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.XLAFSDPStrategy.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`)

#### `clip_gradients_norm(module, optimizer, max_norm, norm_type=2.0, error_if_nonfinite=True)` {#lightning.fabric.strategies.XLAFSDPStrategy.clip_gradients_norm}

Clip gradients by norm.

Return type:

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

#### `clip_gradients_value(module, optimizer, clip_val)` {#lightning.fabric.strategies.XLAFSDPStrategy.clip_gradients_value}

Clip gradients by value.

Return type:

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

#### `load_checkpoint(path, state=None, strict=True, weights_only=None)` {#lightning.fabric.strategies.XLAFSDPStrategy.load_checkpoint}

Given a folder, load the contents from a checkpoint and restore the state of the given objects.

The strategy currently only supports saving and loading sharded checkpoints which are stored in form of a directory of multiple files rather than a single file.

Return type:

[`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)")\]

#### `module_init_context(empty_init=None)` {#lightning.fabric.strategies.XLAFSDPStrategy.module_init_context}

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`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")\]) – 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`](https://docs.python.org/3/library/contextlib.html#contextlib.AbstractContextManager "(in Python v3.15)")

#### `module_sharded_context()` {#lightning.fabric.strategies.XLAFSDPStrategy.module_sharded_context}

A context manager that goes over the instantiation of an [`torch.nn.Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)") and handles sharding of parameters on creation.

By sharding layers directly on instantiation, one can reduce peak memory usage and initialization time.

Return type:

[`AbstractContextManager`](https://docs.python.org/3/library/contextlib.html#contextlib.AbstractContextManager "(in Python v3.15)")

#### `module_to_device(module)` {#lightning.fabric.strategies.XLAFSDPStrategy.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)")

#### `optimizer_step(optimizer, **kwargs)` {#lightning.fabric.strategies.XLAFSDPStrategy.optimizer_step}

Overrides default tpu optimizer\_step since FSDP should not call torch\_xla.core.xla\_model.optimizer\_step. Performs the actual optimizer step.

Parameters:

*   **optimizer** (`Optimizable`) – the optimizer performing the step
    
*   **\*\*kwargs** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – Any extra arguments to `optimizer.step`
    

Return type:

[`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")

#### `process_dataloader(dataloader)` {#lightning.fabric.strategies.XLAFSDPStrategy.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.XLAFSDPStrategy.save_checkpoint}

Save model, optimizer, and other state in the provided checkpoint directory.

If the user specifies sharded checkpointing, the directory will contain one file per process, with model- and optimizer shards stored per file. If the user specifies full checkpointing, the directory will contain a consolidated checkpoint combining all of the sharded checkpoints.

Return type:

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

#### `setup_environment()` {#lightning.fabric.strategies.XLAFSDPStrategy.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.XLAFSDPStrategy.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)")

#### `setup_module_and_optimizers(module, optimizers, scheduler=None)` {#lightning.fabric.strategies.XLAFSDPStrategy.setup_module_and_optimizers}

Returns NotImplementedError since for XLAFSDP optimizer setup must happen after module setup.

Return type:

[`tuple`](https://docs.python.org/3/builtins/stdtypes.html#tuple "(in Python v3.15)")\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`list`](https://docs.python.org/3/builtins/stdtypes.html#list "(in Python v3.15)")\[[`Optimizer`](https://docs.pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.15)")\], [`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[`_LRScheduler`\]\]

#### `setup_optimizer(optimizer)` {#lightning.fabric.strategies.XLAFSDPStrategy.setup_optimizer}

Set up an optimizer for a model wrapped with XLAFSDP.

This setup method doesn’t modify the optimizer or wrap the optimizer. The only thing it currently does is verify that the optimizer was created after the model was wrapped with [`setup_module()`](#lightning.fabric.strategies.XLAFSDPStrategy.setup_module "lightning.fabric.strategies.XLAFSDPStrategy.setup_module") with a reference to the flattened parameters.

Return type:

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

#### `property root_device: device` {#lightning.fabric.strategies.XLAFSDPStrategy.root_device}

Returns the root device.

---

**Navigation**

- Parent: [Generated](/docs/fabric/2.6.6/full-api-reference/api/generated.md)