---
title: ModelParallelStrategy
description: Enables user-defined parallelism applied to a model.
canonical: >-
  https://lightning.ai/docs/pytorch/stable/optional-api/api/lightning.pytorch.strategies.ModelParallelStrategy
---
# ModelParallelStrategy {#modelparallelstrategy}

### `class lightning.pytorch.strategies.ModelParallelStrategy(data_parallel_size='auto', tensor_parallel_size='auto', save_distributed_checkpoint=True, process_group_backend=None, timeout=datetime.timedelta(seconds=1800))` {#lightning.pytorch.strategies.ModelParallelStrategy}

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

Enables user-defined parallelism applied to a model.

Warning

This is an [experimental](/docs/pytorch/2.6.6/reference/versioning.md#experimental-api) feature.

Currently supports up to 2D parallelism. Specifically, it supports the combination of Fully Sharded Data-Parallel 2 (FSDP2) with Tensor Parallelism (DTensor). These PyTorch APIs are currently still experimental in PyTorch (see [https://pytorch.org/docs/stable/distributed.tensor.parallel.html](https://pytorch.org/docs/stable/distributed.tensor.parallel.html)). Requires PyTorch 2.4 or newer.

Parameters:

*   **data\_parallel\_size** (`Union`\[[`Literal`](https://docs.python.org/3/library/typing.html#typing.Literal "(in Python v3.15)")\[`'auto'`\], [`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")\]) – The number of devices within a data-parallel group. Defaults to `"auto"`, which sets this size to the number of nodes in the cluster.
    
*   **tensor\_parallel\_size** (`Union`\[[`Literal`](https://docs.python.org/3/library/typing.html#typing.Literal "(in Python v3.15)")\[`'auto'`\], [`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")\]) – The number of devices within a tensor-parallel group. Defaults to `"auto"`, which sets this size to the number of GPUs in a single node.
    
*   **save\_distributed\_checkpoint** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – If `True`, each rank saves its shard of weights and optimizer states to a file. The checkpoint is a folder with as many files as the world size. If `False`, the full weights and optimizer states get assembled on rank 0 and saved to a single file.
    

#### `barrier(name=None)` {#lightning.pytorch.strategies.ModelParallelStrategy.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.ModelParallelStrategy.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`)

#### `lightning_module_state_dict()` {#lightning.pytorch.strategies.ModelParallelStrategy.lightning_module_state_dict}

Collects the state dict of the model.

Only returns a non-empty state dict on rank 0 if `save_distributed_checkpoint=False`.

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

#### `model_to_device()` {#lightning.pytorch.strategies.ModelParallelStrategy.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)")

#### `optimizer_state(optimizer)` {#lightning.pytorch.strategies.ModelParallelStrategy.optimizer_state}

Collects the state of the given optimizer.

Only returns a non-empty state dict on rank 0 if `save_distributed_checkpoint=False`.

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

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

Reduces the given tensor (e.g. across GPUs/processes).

Parameters:

*   **tensor** (`Union`\[[`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – 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)")

#### `save_checkpoint(checkpoint, filepath, storage_options=None)` {#lightning.pytorch.strategies.ModelParallelStrategy.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.ModelParallelStrategy.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.ModelParallelStrategy.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)")

#### `setup_optimizers(trainer)` {#lightning.pytorch.strategies.ModelParallelStrategy.setup_optimizers}

Creates optimizers and schedulers.

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, these optimizers should be connected to

Return type:

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

#### `teardown()` {#lightning.pytorch.strategies.ModelParallelStrategy.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)")

#### `tensor_init_context(empty_init=None)` {#lightning.pytorch.strategies.ModelParallelStrategy.tensor_init_context}

Controls how tensors get created (device, dtype).

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:

[`Generator`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Generator "(in Python v3.15)")\[[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)"), [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)"), [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]

#### `property lightning_restore_optimizer: bool` {#lightning.pytorch.strategies.ModelParallelStrategy.lightning_restore_optimizer}

Override to disable Lightning restoring optimizers/schedulers.

This is useful for strategies which manage restoring optimizers/schedulers.

#### `property restore_checkpoint_after_setup: bool` {#lightning.pytorch.strategies.ModelParallelStrategy.restore_checkpoint_after_setup}

Override to delay restoring from checkpoint till after the setup phase has completed. This is useful when the strategy requires all the setup hooks to run before loading checkpoint.

Returns:

If `True`, restore checkpoint after strategy setup.

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

Return the root device.

---

**Navigation**

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