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

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

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

Enables user-defined parallelism applied to a model.

Warning

This is an [experimental](/docs/pytorch/stable/versioning#experimental-api "(in PyTorch Lightning v2.7.0.dev)") 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. Requires PyTorch 2.4 or newer.

Parameters:

*   **parallelize\_fn** ([`Callable`](https://docs.python.org/3/library/typing.html#typing.Callable "(in Python v3.15)")\[\[[`TypeVar`](https://docs.python.org/3/library/typing.html#typing.TypeVar "(in Python v3.15)")(`TModel`, bound= [`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)")), [`DeviceMesh`](https://docs.pytorch.org/docs/stable/distributed.html#torch.distributed.device_mesh.DeviceMesh "(in PyTorch v2.15)")\], [`TypeVar`](https://docs.python.org/3/library/typing.html#typing.TypeVar "(in Python v3.15)")(`TModel`, bound= [`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"))\]) – A function that applies parallelisms to a module. The strategy will provide the model and device mesh as input.
    
*   **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.
    

#### `_configure_launcher()` {#lightning.fabric.strategies.ModelParallelStrategy._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_reduce(tensor, group=None, reduce_op='mean')` {#lightning.fabric.strategies.ModelParallelStrategy.all_reduce}

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

Parameters:

*   **tensor** ([`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.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)")

#### `barrier(*args, **kwargs)` {#lightning.fabric.strategies.ModelParallelStrategy.barrier}

Synchronizes all processes which blocks processes until the whole group enters this function.

Parameters:

**name** – 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.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`)

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

Load the contents from a checkpoint and restore the state of the given objects.

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

#### `save_checkpoint(path, state, storage_options=None, filter=None)` {#lightning.fabric.strategies.ModelParallelStrategy.save_checkpoint}

Save model, optimizer, and other state to a checkpoint on disk.

If distributed checkpointing is enabled (default), the checkpoint gets saved as a directory containing one file per process, with model- and optimizer shards stored per file. Additionally, it creates a metadata file meta.pt with the rest of the user’s state (only saved from rank 0). If distributed checkpointing is disabled (`save_distributed_checkpoint=False`), the checkpoint will be written to a single file containing the weights, optimizer state and other metadata.

Return type:

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

#### `setup_environment()` {#lightning.fabric.strategies.ModelParallelStrategy.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.ModelParallelStrategy.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 distributed_sampler_kwargs: dict[str, Any]` {#lightning.fabric.strategies.ModelParallelStrategy.distributed_sampler_kwargs}

Arguments for the `DistributedSampler`.

If this method is not defined, or it returns `None`, then the `DistributedSampler` will not be used.

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

Returns the root device.

---

**Navigation**

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