DeepSpeedStrategy

class lightning.pytorch.strategies.DeepSpeedStrategy(accelerator=None, zero_optimization=True, stage=2, remote_device='cpu', 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='auto', 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_plugin=None, process_group_backend=None)

Bases: lightning.pytorch.strategies.ddp.DDPStrategy

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 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) – Enable ZeRO optimization. This is compatible with either precision=”16-mixed” or precision=”bf16-mixed”.

• stage (int) – 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 (str) – Device to instantiate the model on initially (cpu or nvme).

• offload_optimizer (bool) – Enable offloading optimizer memory and computation to CPU or NVMe based on offload_optimizer_device.

• offload_parameters (bool) – When using ZeRO Stage 3, Enable offloading parameter memory and computation to CPU or NVMe based on offload_params_device.

• offload_params_device (str) – When offloading parameters choose the device to offload to, cpu or nvme.

• offload_optimizer_device (str) – When offloading optimizer state choose the device to offload to, cpu or nvme.

• params_buffer_count (int) – Number of buffers in buffer pool for parameter offloading when offload_params_device is nvme.

• params_buffer_size (int) – Size of buffers in buffer pool for parameter offloading when offload_params_device is nvme.

• max_in_cpu (int) – Number of parameter elements to maintain in CPU memory when offloading to NVMe is enabled.

• nvme_path (str) – Filesystem path for NVMe device for optimizer/parameter state offloading.

• optimizer_buffer_count (int) – Number of buffers in buffer pool for optimizer state offloading when offload_optimizer_device is set to 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) – When using NVMe Offloading, the I/O block size in bytes.

• queue_depth (int) – When using NVMe Offloading, the I/O queue depth.

• single_submit (bool) – When using NVMe Offloading, submit requests to storage device as multiple individual requests, as opposed to one block of requests.

• overlap_events (bool) – When using NVMe Offloading, submit requests to storage device in an overlapped fashion without waiting for completion of earlier requests.

• thread_count (int) – When using NVMe Offloading, Intra-request parallelism for each read/write submitted by a user thread.

• pin_memory (bool) – 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) – 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) – Copies gradients to a continuous buffer as they are produced. Avoids memory fragmentation during backwards. Useful when training large models.

• overlap_comm (bool) – Overlap the reduction (synchronization) of gradients with the backwards computation. This is a speed optimization when training across multiple GPUs/machines.

• allgather_partitions (bool) – All gather updated parameters at the end of training step, instead of using a series of broadcast collectives.

• reduce_scatter (bool) – Use reduce/scatter instead of allreduce to average gradients.

• allgather_bucket_size (int) – 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) – 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) – Allow untested optimizers to be used with ZeRO. Currently only Adam is a DeepSpeed supported optimizer when using ZeRO.

• logging_batch_size_per_gpu (Union[str, int]) – Config used in DeepSpeed to calculate verbose timing for logging on a per sample per second basis (only displayed if logging=logging.INFO). If set to “auto”, the plugin tries to infer this from the train DataLoader’s BatchSampler, else defaults to 1. To obtain accurate logs when using datasets that do not support batch samplers, set this to the actual per gpu batch size (trainer.batch_size).

• config (Union[str, Path, Dict[str, Any], None]) – Pass in a deepspeed formatted config dict, or path to a deepspeed config: https://www.deepspeed.ai/docs/config-json. All defaults will be ignored if a config is passed in.

• logging_level (int) – Set logging level for deepspeed.

• loss_scale (float) – Loss scaling value for FP16 training. 0.0 results in dynamic loss scaling, otherwise static.

• initial_scale_power (int) – Power of the initial dynamic loss scale value. Loss scale is computed by 2^initial_scale_power.

• loss_scale_window (int) – Window in which to raise/lower the dynamic FP16 loss scaling value.

• hysteresis (int) – FP16 Delay shift in Dynamic Loss scaling.

• min_loss_scale (int) – The minimum FP16 dynamic loss scaling value.

• partition_activations (bool) – 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.

• cpu_checkpointing (bool) – Offloads partitioned activations to CPU if partition_activations is enabled.

• contiguous_memory_optimization (bool) – Copies partitioned activations so that they are contiguous in memory. Not supported by all models.

• synchronize_checkpoint_boundary (bool) – Insert torch.cuda.synchronize() at each checkpoint boundary.

• load_full_weights (bool) – 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.

batch_to_device(batch, device=None, dataloader_idx=0)

Moves the batch to the correct device.

The returned batch is of the same type as the input batch, just having all tensors on the correct device.

Parameters

• batch (Any) – The batch of samples to move to the correct device

• device (Optional[device]) – The target device

• dataloader_idx (int) – The index of the dataloader to which the batch belongs.

Return type

Any

model_sharded_context()

Provide hook to create modules in a distributed aware context. This is useful for when we’d like to shard the model instantly, which is useful for extremely large models which can save memory and initialization time.

Returns: Model parallel context.

Return type

Generator[None, None, None]

predict_step(*args, **kwargs)

The actual predict step.

See predict_step() for more details

Return type

Union[Tensor, Dict[str, Any]]

save_checkpoint(checkpoint, filepath, storage_options=None)

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

Parameters

• checkpoint (Dict) – The checkpoint state dictionary

• filepath (Union[str, Path]) – write-target file’s path

• storage_options (Optional[Any]) – not used for DeepSpeedStrategy as CheckpointIO is not used

Raises

TypeError – If storage_options arg is passed in

Return type

None

setup(trainer)

Setup plugins for the trainer fit and creates optimizers.

Parameters

trainer (Trainer) – the trainer instance

Return type

None

setup_optimizers(trainer)

Creates optimizers and schedulers.

Parameters

trainer (Trainer) – the Trainer, these optimizers should be connected to

Return type

None

test_step(*args, **kwargs)

The actual test step.

See test_step() for more details

Return type

Union[Tensor, Dict[str, Any], None]

validation_step(*args, **kwargs)

The actual validation step.

See validation_step() for more details

Return type

Union[Tensor, Dict[str, Any], None]

property handles_gradient_accumulation: bool

Whether the plugin handles gradient accumulation internally.

Return type

bool

property lightning_restore_optimizer: bool

Override to disable Lightning restoring optimizers/schedulers.

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

Return type

bool

property restore_checkpoint_after_setup: bool

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.

Return type

bool

Returns

If True, restore checkpoint after strategy setup.