---
title: DeepSpeedStrategy
description: >-
  Provides capabilities to run training using the DeepSpeed library, with
  training optimizations for large

  billion parameter models. For more information: https://pytorch-

  lightning.readthedocs.io/en/st
canonical: >-
  https://lightning.ai/docs/fabric/stable/full-api-reference/api/generated/lightning.fabric.strategies.DeepSpeedStrategy
---
# DeepSpeedStrategy {#deepspeedstrategy}

### `class lightning.fabric.strategies.DeepSpeedStrategy(accelerator=None, zero_optimization=True, stage=2, remote_device=None, offload_optimizer=False, offload_parameters=False, offload_params_device='cpu', nvme_path='/local_nvme', params_buffer_count=5, params_buffer_size=100000000, max_in_cpu=1000000000, offload_optimizer_device='cpu', optimizer_buffer_count=4, block_size=1048576, queue_depth=8, single_submit=False, overlap_events=True, thread_count=1, pin_memory=False, sub_group_size=1000000000000, contiguous_gradients=True, overlap_comm=True, allgather_partitions=True, reduce_scatter=True, allgather_bucket_size=200000000, reduce_bucket_size=200000000, zero_allow_untested_optimizer=True, logging_batch_size_per_gpu=None, config=None, logging_level=30, parallel_devices=None, cluster_environment=None, loss_scale=0, initial_scale_power=16, loss_scale_window=1000, hysteresis=2, min_loss_scale=1, partition_activations=False, cpu_checkpointing=False, contiguous_memory_optimization=False, synchronize_checkpoint_boundary=False, load_full_weights=False, precision=None, process_group_backend=None, timeout=datetime.timedelta(seconds=1800), exclude_frozen_parameters=False)` {#lightning.fabric.strategies.DeepSpeedStrategy}

Bases: [`DDPStrategy`](/docs/fabric/2.6.6/full-api-reference/api/generated/lightning.fabric.strategies.DDPStrategy.md#lightning.fabric.strategies.DDPStrategy "lightning.fabric.strategies.ddp.DDPStrategy"), `_Sharded`

Provides capabilities to run training using the DeepSpeed library, with training optimizations for large billion parameter models. For more information: https://pytorch- lightning.readthedocs.io/en/stable/advanced/model\_parallel.html#deepspeed.

Warning

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

Defaults have been set to enable ZeRO-Offload and some have been taken from the link below. These defaults have been set generally, but may require tuning for optimum performance based on your model size. For more information: https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training.

Parameters:

*   **zero\_optimization** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Enable ZeRO optimization. This is compatible with either `precision="16-mixed"` or `precision="bf16-mixed"`.
    
*   **stage** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Different stages of the ZeRO Optimizer. 0 is disabled, 1 is optimizer state partitioning, 2 is optimizer+gradient state partitioning, 3 is optimizer+gradient\_parameter partitioning using the infinity engine.
    
*   **remote\_device** ([`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)")\]) – Device to instantiate the model on initially (`cpu` or `nvme`). Defaults to GPU.
    
*   **offload\_optimizer** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Enable offloading optimizer memory and computation to CPU or NVMe based on `offload_optimizer_device`.
    
*   **offload\_parameters** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – When using ZeRO Stage 3, Enable offloading parameter memory and computation to CPU or NVMe based on `offload_params_device`.
    
*   **offload\_params\_device** ([`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")) – When offloading parameters choose the device to offload to, `cpu` or `nvme`.
    
*   **offload\_optimizer\_device** ([`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")) – When offloading optimizer state choose the device to offload to, `cpu` or `nvme`.
    
*   **params\_buffer\_count** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Number of buffers in buffer pool for parameter offloading when `offload_params_device` is `nvme`.
    
*   **params\_buffer\_size** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Size of buffers in buffer pool for parameter offloading when `offload_params_device` is `nvme`.
    
*   **max\_in\_cpu** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Number of parameter elements to maintain in CPU memory when offloading to NVMe is enabled.
    
*   **nvme\_path** ([`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")) – Filesystem path for NVMe device for optimizer/parameter state offloading.
    
*   **optimizer\_buffer\_count** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Number of buffers in buffer pool for optimizer state offloading when `offload_optimizer_device` is set to `nvme`. This should be at least the number of states maintained per parameter by the optimizer. For example, Adam optimizer has 4 states (parameter, gradient, momentum, and variance).
    
*   **block\_size** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – When using NVMe Offloading, the I/O block size in bytes.
    
*   **queue\_depth** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – When using NVMe Offloading, the I/O queue depth.
    
*   **single\_submit** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – When using NVMe Offloading, submit requests to storage device as multiple individual requests, as opposed to one block of requests.
    
*   **overlap\_events** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – When using NVMe Offloading, submit requests to storage device in an overlapped fashion without waiting for completion of earlier requests.
    
*   **thread\_count** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – When using NVMe Offloading, Intra-request parallelism for each read/write submitted by a user thread.
    
*   **pin\_memory** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – When using ZeRO stage 3, pin optimizer state memory on CPU. This could boost throughput at the cost of extra memory overhead.
    
*   **sub\_group\_size** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – When using ZeRO stage 3, defines the number of parameters within a sub group to offload at a time. Smaller numbers require more communication, but improve memory efficiency.
    
*   **contiguous\_gradients** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Copies gradients to a continuous buffer as they are produced. Avoids memory fragmentation during backwards. Useful when training large models.
    
*   **overlap\_comm** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Overlap the reduction (synchronization) of gradients with the backwards computation. This is a speed optimization when training across multiple GPUs/machines.
    
*   **allgather\_partitions** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – All gather updated parameters at the end of training step, instead of using a series of broadcast collectives.
    
*   **reduce\_scatter** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Use reduce/scatter instead of allreduce to average gradients.
    
*   **allgather\_bucket\_size** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Number of elements to allgather at once. Used to limit the memory required for larger model sizes, with a tradeoff with speed.
    
*   **reduce\_bucket\_size** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Number of elements to reduce at once. Used to limit the memory required for larger model sizes, with a tradeoff with speed.
    
*   **zero\_allow\_untested\_optimizer** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Allow untested optimizers to be used with ZeRO. Currently only Adam is a DeepSpeed supported optimizer when using ZeRO.
    
*   **logging\_batch\_size\_per\_gpu** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")\]) – Config used in DeepSpeed to calculate verbose timing for logging on a per sample per second basis (only displayed if logging=logging.INFO). To obtain accurate logs when using datasets that do not support batch samplers, set this to the actual per gpu batch size.
    
*   **config** (`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)"), [`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)")\], [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – Pass in a deepspeed formatted config dict, or path to a deepspeed config: [https://www.deepspeed.ai/docs/config-json](https://www.deepspeed.ai/docs/config-json). All defaults will be ignored if a config is passed in.
    
*   **logging\_level** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Set logging level for deepspeed.
    
*   **loss\_scale** ([`float`](https://docs.python.org/3/builtins/functions.html#float "(in Python v3.15)")) – Loss scaling value for FP16 training. 0.0 results in dynamic loss scaling, otherwise static.
    
*   **initial\_scale\_power** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Power of the initial dynamic loss scale value. Loss scale is computed by `2^initial_scale_power`.
    
*   **loss\_scale\_window** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Window in which to raise/lower the dynamic FP16 loss scaling value.
    
*   **hysteresis** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – FP16 Delay shift in Dynamic Loss scaling.
    
*   **min\_loss\_scale** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – The minimum FP16 dynamic loss scaling value.
    
*   **partition\_activations** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Enables partition activation when used with ZeRO stage 3 and model parallelism. Still requires you to wrap your forward functions in deepspeed.checkpointing.checkpoint. See [deepspeed tutorial](https://www.deepspeed.ai/tutorials/megatron/#deepspeed-activation-checkpoints-optional).
    
*   **cpu\_checkpointing** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Offloads partitioned activations to CPU if `partition_activations` is enabled.
    
*   **contiguous\_memory\_optimization** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Copies partitioned activations so that they are contiguous in memory. Not supported by all models.
    
*   **synchronize\_checkpoint\_boundary** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Insert [`torch.cuda.synchronize()`](https://docs.pytorch.org/docs/stable/generated/torch.cuda.synchronize.html#torch.cuda.synchronize "(in PyTorch v2.15)") at each checkpoint boundary.
    
*   **load\_full\_weights** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – True when loading a single checkpoint file containing the model state dict when using ZeRO Stage 3. This differs from the DeepSpeed checkpoint which contains shards per worker.
    
*   **exclude\_frozen\_parameters** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Exclude frozen parameters when saving checkpoints.
    

#### `_initialize_engine(model, optimizer=None, scheduler=None)` {#lightning.fabric.strategies.DeepSpeedStrategy._initialize_engine}

Initialize one model and one optimizer with an optional learning rate scheduler.

This calls `deepspeed.initialize` internally.

Return type:

[`tuple`](https://docs.python.org/3/builtins/stdtypes.html#tuple "(in Python v3.15)")\[DeepSpeedEngine, [`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)")\]

#### `_restore_zero_state(module, ckpt)` {#lightning.fabric.strategies.DeepSpeedStrategy._restore_zero_state}

Overrides the normal load\_state\_dict behaviour in PyTorch to ensure we gather parameters that may be sharded across processes before loading the state dictionary when using ZeRO stage 3. This is then automatically synced across processes.

Parameters:

**ckpt** ([`Mapping`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Mapping "(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)")\]) – The ckpt file.

Return type:

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

#### `clip_gradients_norm(module, optimizer, max_norm, norm_type=2.0, error_if_nonfinite=True)` {#lightning.fabric.strategies.DeepSpeedStrategy.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.DeepSpeedStrategy.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.DeepSpeedStrategy.load_checkpoint}

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

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 is located
    
*   **state** (`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)"), [`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)")\]\], [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – A dictionary of objects whose state will be restored in-place from the checkpoint path. This should contain exactly one model, and the model must already be set up by DeepSpeed.
    
*   **strict** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Whether to enforce that the keys in state match the keys in the checkpoint.
    

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

Returns:

Dictionary with the state inside DeepSpeed’s engine

Raises:

*   [**ValueError**](https://docs.python.org/3/builtins/exceptions.html#ValueError "(in Python v3.15)") – If no state is provided, when no `deepspeed.DeepSpeedEngine` objects were found in the state, or when multiple `deepspeed.DeepSpeedEngine` objects were found.
    
*   [**RuntimeError**](https://docs.python.org/3/builtins/exceptions.html#RuntimeError "(in Python v3.15)") – If DeepSpeed was unable to load the checkpoint due to missing files or because the checkpoint is not in the expected DeepSpeed format.
    

#### `module_init_context(empty_init=None)` {#lightning.fabric.strategies.DeepSpeedStrategy.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.DeepSpeedStrategy.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)")

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

Save model, optimizer, and other state in a checkpoint directory.

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 files 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)")\]) – Unused by this strategy, since it doesn’t use a `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)")\]\]\]) – Unsupported.
    

Raises:

*   [**TypeError**](https://docs.python.org/3/builtins/exceptions.html#TypeError "(in Python v3.15)") – If the unused `storage_options` gets passed.
    
*   [**ValueError**](https://docs.python.org/3/builtins/exceptions.html#ValueError "(in Python v3.15)") – When no `deepspeed.DeepSpeedEngine` objects were found in the state, or when multiple `deepspeed.DeepSpeedEngine` objects were found.
    

Return type:

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

#### `setup_environment()` {#lightning.fabric.strategies.DeepSpeedStrategy.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.DeepSpeedStrategy.setup_module}

Set up a module for inference (no optimizers).

For training, see [`setup_module_and_optimizers()`](#lightning.fabric.strategies.DeepSpeedStrategy.setup_module_and_optimizers "lightning.fabric.strategies.DeepSpeedStrategy.setup_module_and_optimizers").

Return type:

`DeepSpeedEngine`

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

Set up a model and multiple optimizers together, along with an optional learning rate scheduler. Currently, only a single optimizer is supported.

Return type:

[`tuple`](https://docs.python.org/3/builtins/stdtypes.html#tuple "(in Python v3.15)")\[DeepSpeedEngine, [`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)")\], [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]

Returns:

The model wrapped into a `deepspeed.DeepSpeedEngine`, a list with a single deepspeed optimizer, and an optional learning rate scheduler.

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

Optimizers can only be set up jointly with the model in this strategy.

Please use [`setup_module_and_optimizers()`](#lightning.fabric.strategies.DeepSpeedStrategy.setup_module_and_optimizers "lightning.fabric.strategies.DeepSpeedStrategy.setup_module_and_optimizers") to set up both module and optimizer together.

Return type:

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

#### `property distributed_sampler_kwargs: dict[str, int]` {#lightning.fabric.strategies.DeepSpeedStrategy.distributed_sampler_kwargs}

Arguments for the `DistributedSampler`.

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

---

**Navigation**

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