DeepSpeedPlugin(zero_optimization=True, stage=2, cpu_offload=False, cpu_offload_params=False, cpu_offload_use_pin_memory=False, contiguous_gradients=True, overlap_comm=True, allgather_partitions=True, reduce_scatter=True, allgather_bucket_size=200000000.0, reduce_bucket_size=200000000.0, zero_allow_untested_optimizer=True, logging_batch_size_per_gpu='auto', config=None, logging_level=30, num_nodes=None, 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, save_full_weights=True)¶
Provides capabilities to run training using the DeepSpeed library, with training optimizations for large billion parameter models. For more information: https://www.deepspeed.ai/.
DeepSpeedPluginis in beta and subject to change.
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.
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).
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. (Default:
bool) – Enables partition activation when used with ZeRO stage 3. Still requires you to wrap your forward functions in deepspeed.checkpointing.checkpoint. See deepspeed tutorial
bool) – Gathers weights across all processes before saving to disk when using ZeRO Stage 3. This allows a single weight file to contain the entire model, rather than individual sharded weight files. Disable to save sharded states individually. (Default: True)
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.
Hook to do something before the training/evaluation/prediction starts.
restore_model_state_from_ckpt_path(ckpt_path, map_location=<function DeepSpeedPlugin.<lambda>>)¶
This function is used to load and restore the model state.
checkpoint: Return loaded checkpoint bool: Wether to load optimizer / lr_schedulers states from checkpoint
Save model/training states as a checkpoint file through state-dump and file-write.
Provide a hook to count optimizer step calls.
Returns: New optimizer step calls
- Return type
Returns the pure LightningModule without potential wrappers