docstr: 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 :ref:`experimental <versioning:Experimental API>` 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`.
  
  Arguments:
  
      zero_optimization: Enable ZeRO optimization. This is compatible with either `precision="16-mixed"` or
          `precision="bf16-mixed"`.
  
      stage: 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: Device to instantiate the model on initially (``cpu`` or ``nvme``). Defaults to GPU.
  
      offload_optimizer: Enable offloading optimizer memory and computation to CPU or NVMe
          based on ``offload_optimizer_device``.
  
      offload_parameters: When using ZeRO Stage 3, Enable offloading parameter memory and computation
          to CPU or NVMe based on ``offload_params_device``.
  
      offload_params_device: When offloading parameters choose the device to offload to, ``cpu`` or ``nvme``.
  
      offload_optimizer_device: When offloading optimizer state choose the device to offload to,
          ``cpu`` or ``nvme``.
  
      params_buffer_count: Number of buffers in buffer pool for
          parameter offloading when ``offload_params_device`` is ``nvme``.
  
      params_buffer_size: Size of buffers in buffer pool for parameter offloading
          when ``offload_params_device`` is ``nvme``.
  
      max_in_cpu: Number of parameter elements to maintain in CPU memory when offloading to NVMe is enabled.
  
      nvme_path: Filesystem path for NVMe device for optimizer/parameter state offloading.
  
      optimizer_buffer_count: 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: When using NVMe Offloading, the I/O block size in bytes.
  
      queue_depth: When using NVMe Offloading, the I/O queue depth.
  
      single_submit: When using NVMe Offloading,
          submit requests to storage device as multiple individual requests,
          as opposed to one block of requests.
  
      overlap_events: When using NVMe Offloading,
          submit requests to storage device in an overlapped fashion
          without waiting for completion of earlier requests.
  
      thread_count: When using NVMe Offloading,
          Intra-request parallelism for each read/write submitted by a user thread.
  
      pin_memory: 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: 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: Copies gradients to a continuous buffer as they are produced.
          Avoids memory fragmentation during backwards. Useful when training large models.
  
      overlap_comm: Overlap the reduction (synchronization) of gradients with the backwards computation.
          This is a speed optimization when training across multiple GPUs/machines.
  
      allgather_partitions: All gather updated parameters at the end of training step,
          instead of using a series of broadcast collectives.
  
      reduce_scatter: Use reduce/scatter instead of allreduce to average gradients.
  
      allgather_bucket_size: 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: 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: Allow untested optimizers to be used with ZeRO. Currently only Adam is a
          DeepSpeed supported optimizer when using ZeRO.
  
      logging_batch_size_per_gpu: 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 strategy 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: 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: Set logging level for deepspeed.
  
      loss_scale: Loss scaling value for FP16 training.
          0.0 results in dynamic loss scaling, otherwise static.
  
      initial_scale_power: Power of the initial dynamic loss scale value. Loss scale is computed
          by ``2^initial_scale_power``.
  
      loss_scale_window: Window in which to raise/lower the dynamic FP16 loss scaling value.
  
      hysteresis: FP16 Delay shift in Dynamic Loss scaling.
  
      min_loss_scale: The minimum FP16 dynamic loss scaling value.
  
      partition_activations: 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: Offloads partitioned activations to CPU if ``partition_activations`` is enabled.
  
      contiguous_memory_optimization: Copies partitioned activations so that they are contiguous in memory.
          Not supported by all models.
  
      synchronize_checkpoint_boundary: Insert :func:`torch.cuda.synchronize` at each checkpoint boundary.
  
      load_full_weights: 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.
docstr_tok: ["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",":","ref",":","`","experimental","<","versioning",":","Experimental","API",">","`","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","`",".","Arguments",":","zero_optimization",":","Enable","ZeRO","optimization",".","This","is","compatible","with","either","`","precision=","''","16-mixed","''","`","or","`","precision=","''","bf16-mixed","''","`",".","stage",":","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",":","Device","to","instantiate","the","model","on","initially","(","``","cpu","``","or","``","nvme","``",")",".","Defaults","to","GPU",".","offload_optimizer",":","Enable","offloading","optimizer","memory","and","computation","to","CPU","or","NVMe","based","on","``","offload_optimizer_device","``",".","offload_parameters",":","When","using","ZeRO","Stage","3",",","Enable","offloading","parameter","memory","and","computation","to","CPU","or","NVMe","based","on","``","offload_params_device","``",".","offload_params_device",":","When","offloading","parameters","choose","the","device","to","offload","to",",","``","cpu","``","or","``","nvme","``",".","offload_optimizer_device",":","When","offloading","optimizer","state","choose","the","device","to","offload","to",",","``","cpu","``","or","``","nvme","``",".","params_buffer_count",":","Number","of","buffers","in","buffer","pool","for","parameter","offloading","when","``","offload_params_device","``","is","``","nvme","``",".","params_buffer_size",":","Size","of","buffers","in","buffer","pool","for","parameter","offloading","when","``","offload_params_device","``","is","``","nvme","``",".","max_in_cpu",":","Number","of","parameter","elements","to","maintain","in","CPU","memory","when","offloading","to","NVMe","is","enabled",".","nvme_path",":","Filesystem","path","for","NVMe","device","for","optimizer\/parameter","state","offloading",".","optimizer_buffer_count",":","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",":","When","using","NVMe","Offloading",",","the","I\/O","block","size","in","bytes",".","queue_depth",":","When","using","NVMe","Offloading",",","the","I\/O","queue","depth",".","single_submit",":","When","using","NVMe","Offloading",",","submit","requests","to","storage","device","as","multiple","individual","requests",",","as","opposed","to","one","block","of","requests",".","overlap_events",":","When","using","NVMe","Offloading",",","submit","requests","to","storage","device","in","an","overlapped","fashion","without","waiting","for","completion","of","earlier","requests",".","thread_count",":","When","using","NVMe","Offloading",",","Intra-request","parallelism","for","each","read\/write","submitted","by","a","user","thread",".","pin_memory",":","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",":","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",":","Copies","gradients","to","a","continuous","buffer","as","they","are","produced",".","Avoids","memory","fragmentation","during","backwards",".","Useful","when","training","large","models",".","overlap_comm",":","Overlap","the","reduction","(","synchronization",")","of","gradients","with","the","backwards","computation",".","This","is","a","speed","optimization","when","training","across","multiple","GPUs\/machines",".","allgather_partitions",":","All","gather","updated","parameters","at","the","end","of","training","step",",","instead","of","using","a","series","of","broadcast","collectives",".","reduce_scatter",":","Use","reduce\/scatter","instead","of","allreduce","to","average","gradients",".","allgather_bucket_size",":","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",":","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",":","Allow","untested","optimizers","to","be","used","with","ZeRO",".","Currently","only","Adam","is","a","DeepSpeed","supported","optimizer","when","using","ZeRO",".","logging_batch_size_per_gpu",":","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","strategy","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",":","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",":","Set","logging","level","for","deepspeed",".","loss_scale",":","Loss","scaling","value","for","FP16","training",".","0.0","results","in","dynamic","loss","scaling",",","otherwise","static",".","initial_scale_power",":","Power","of","the","initial","dynamic","loss","scale","value",".","Loss","scale","is","computed","by","``","2^initial_scale_power","``",".","loss_scale_window",":","Window","in","which","to","raise\/lower","the","dynamic","FP16","loss","scaling","value",".","hysteresis",":","FP16","Delay","shift","in","Dynamic","Loss","scaling",".","min_loss_scale",":","The","minimum","FP16","dynamic","loss","scaling","value",".","partition_activations",":","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",":","Offloads","partitioned","activations","to","CPU","if","``","partition_activations","``","is","enabled",".","contiguous_memory_optimization",":","Copies","partitioned","activations","so","that","they","are","contiguous","in","memory",".","Not","supported","by","all","models",".","synchronize_checkpoint_boundary",":","Insert",":","func",":","`","torch.cuda.synchronize","`","at","each","checkpoint","boundary",".","load_full_weights",":","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","."]
code: def __init__(
      self,
      accelerator: Optional["pl.accelerators.Accelerator"] = None,
      zero_optimization: bool = True,
      stage: int = 2,
      remote_device: Optional[str] = None,
      offload_optimizer: bool = False,
      offload_parameters: bool = False,
      offload_params_device: str = "cpu",
      nvme_path: str = "/local_nvme",
      params_buffer_count: int = 5,
      params_buffer_size: int = 100_000_000,
      max_in_cpu: int = 1_000_000_000,
      offload_optimizer_device: str = "cpu",
      optimizer_buffer_count: int = 4,
      block_size: int = 1048576,
      queue_depth: int = 8,
      single_submit: bool = False,
      overlap_events: bool = True,
      thread_count: int = 1,
      pin_memory: bool = False,
      sub_group_size: int = 1_000_000_000_000,
      contiguous_gradients: bool = True,
      overlap_comm: bool = True,
      allgather_partitions: bool = True,
      reduce_scatter: bool = True,
      allgather_bucket_size: int = 200_000_000,
      reduce_bucket_size: int = 200_000_000,
      zero_allow_untested_optimizer: bool = True,
      logging_batch_size_per_gpu: Union[str, int] = "auto",
      config: Optional[Union[_PATH, Dict[str, Any]]] = None,
      logging_level: int = logging.WARN,
      parallel_devices: Optional[List[torch.device]] = None,
      cluster_environment: Optional[ClusterEnvironment] = None,
      loss_scale: float = 0,
      initial_scale_power: int = 16,
      loss_scale_window: int = 1000,
      hysteresis: int = 2,
      min_loss_scale: int = 1,
      partition_activations: bool = False,
      cpu_checkpointing: bool = False,
      contiguous_memory_optimization: bool = False,
      synchronize_checkpoint_boundary: bool = False,
      load_full_weights: bool = False,
      precision_plugin: Optional[Precision] = None,
      process_group_backend: Optional[str] = None,
  ) -> None:
      """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 :ref:`experimental <versioning:Experimental API>` 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`.
  
      Arguments:
  
          zero_optimization: Enable ZeRO optimization. This is compatible with either `precision="16-mixed"` or
              `precision="bf16-mixed"`.
  
          stage: 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: Device to instantiate the model on initially (``cpu`` or ``nvme``). Defaults to GPU.
  
          offload_optimizer: Enable offloading optimizer memory and computation to CPU or NVMe
              based on ``offload_optimizer_device``.
  
          offload_parameters: When using ZeRO Stage 3, Enable offloading parameter memory and computation
              to CPU or NVMe based on ``offload_params_device``.
  
          offload_params_device: When offloading parameters choose the device to offload to, ``cpu`` or ``nvme``.
  
          offload_optimizer_device: When offloading optimizer state choose the device to offload to,
              ``cpu`` or ``nvme``.
  
          params_buffer_count: Number of buffers in buffer pool for
              parameter offloading when ``offload_params_device`` is ``nvme``.
  
          params_buffer_size: Size of buffers in buffer pool for parameter offloading
              when ``offload_params_device`` is ``nvme``.
  
          max_in_cpu: Number of parameter elements to maintain in CPU memory when offloading to NVMe is enabled.
  
          nvme_path: Filesystem path for NVMe device for optimizer/parameter state offloading.
  
          optimizer_buffer_count: 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: When using NVMe Offloading, the I/O block size in bytes.
  
          queue_depth: When using NVMe Offloading, the I/O queue depth.
  
          single_submit: When using NVMe Offloading,
              submit requests to storage device as multiple individual requests,
              as opposed to one block of requests.
  
          overlap_events: When using NVMe Offloading,
              submit requests to storage device in an overlapped fashion
              without waiting for completion of earlier requests.
  
          thread_count: When using NVMe Offloading,
              Intra-request parallelism for each read/write submitted by a user thread.
  
          pin_memory: 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: 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: Copies gradients to a continuous buffer as they are produced.
              Avoids memory fragmentation during backwards. Useful when training large models.
  
          overlap_comm: Overlap the reduction (synchronization) of gradients with the backwards computation.
              This is a speed optimization when training across multiple GPUs/machines.
  
          allgather_partitions: All gather updated parameters at the end of training step,
              instead of using a series of broadcast collectives.
  
          reduce_scatter: Use reduce/scatter instead of allreduce to average gradients.
  
          allgather_bucket_size: 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: 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: Allow untested optimizers to be used with ZeRO. Currently only Adam is a
              DeepSpeed supported optimizer when using ZeRO.
  
          logging_batch_size_per_gpu: 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 strategy 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: 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: Set logging level for deepspeed.
  
          loss_scale: Loss scaling value for FP16 training.
              0.0 results in dynamic loss scaling, otherwise static.
  
          initial_scale_power: Power of the initial dynamic loss scale value. Loss scale is computed
              by ``2^initial_scale_power``.
  
          loss_scale_window: Window in which to raise/lower the dynamic FP16 loss scaling value.
  
          hysteresis: FP16 Delay shift in Dynamic Loss scaling.
  
          min_loss_scale: The minimum FP16 dynamic loss scaling value.
  
          partition_activations: 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: Offloads partitioned activations to CPU if ``partition_activations`` is enabled.
  
          contiguous_memory_optimization: Copies partitioned activations so that they are contiguous in memory.
              Not supported by all models.
  
          synchronize_checkpoint_boundary: Insert :func:`torch.cuda.synchronize` at each checkpoint boundary.
  
          load_full_weights: 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.
  
      """
      if not _DEEPSPEED_AVAILABLE:
          raise MisconfigurationException(
              "To use the `DeepSpeedStrategy`, you must have DeepSpeed installed."
              " Install it by running `pip install -U deepspeed`."
          )
  
      super().__init__(
          accelerator=accelerator,
          parallel_devices=parallel_devices,
          cluster_environment=cluster_environment,
          precision_plugin=precision_plugin,
          process_group_backend=process_group_backend,
      )
  
      self.config = self._load_config(config)
      if self.config is None:
          # User has not overridden config, set defaults
          self.config = self._create_default_config(
              zero_optimization,
              zero_allow_untested_optimizer,
              logging_batch_size_per_gpu,
              offload_optimizer=offload_optimizer,
              offload_parameters=offload_parameters,
              nvme_path=nvme_path,
              offload_params_device=offload_params_device,
              params_buffer_count=params_buffer_count,
              params_buffer_size=params_buffer_size,
              max_in_cpu=max_in_cpu,
              pin_memory=pin_memory,
              offload_optimizer_device=offload_optimizer_device,
              optimizer_buffer_count=optimizer_buffer_count,
              block_size=block_size,
              queue_depth=queue_depth,
              single_submit=single_submit,
              overlap_events=overlap_events,
              thread_count=thread_count,
              partition_activations=partition_activations,
              cpu_checkpointing=cpu_checkpointing,
              contiguous_memory_optimization=contiguous_memory_optimization,
              synchronize_checkpoint_boundary=synchronize_checkpoint_boundary,
              stage=stage,
              contiguous_gradients=contiguous_gradients,
              overlap_comm=overlap_comm,
              allgather_partitions=allgather_partitions,
              reduce_scatter=reduce_scatter,
              allgather_bucket_size=allgather_bucket_size,
              reduce_bucket_size=reduce_bucket_size,
              sub_group_size=sub_group_size,
          )
      import deepspeed
  
      self._config_initialized = False
      deepspeed.utils.logging.logger.setLevel(logging_level)
  
      self.remote_device = remote_device
      self.load_full_weights = load_full_weights
  
      # default FP16 parameters.
      self.loss_scale = loss_scale
      self.initial_scale_power = initial_scale_power
      self.loss_scale_window = loss_scale_window
      self.hysteresis = hysteresis
      self.min_loss_scale = min_loss_scale
  
imports: import argparse
  import json
  import logging
  import os
  import platform
  from collections import OrderedDict
  from contextlib import contextmanager
  from pathlib import Path
  from typing import TYPE_CHECKING, Any, Dict, Generator, List, Mapping, Optional, Tuple, Union
  import torch
  from torch.nn import Module
  from torch.optim import Optimizer
  from typing_extensions import override
  import lightning.pytorch
  from lightning.fabric.plugins import ClusterEnvironment
  from lightning.fabric.strategies import _StrategyRegistry
  from lightning.fabric.strategies.deepspeed import _DEEPSPEED_AVAILABLE, _format_precision_config, _validate_checkpoint_directory, _validate_device_index_selection
  from lightning.fabric.utilities.optimizer import _optimizers_to_device
  from lightning.fabric.utilities.seed import reset_seed
  from lightning.fabric.utilities.types import _PATH, LRScheduler, ReduceLROnPlateau
  from lightning.pytorch.accelerators.cuda import CUDAAccelerator
  from lightning.pytorch.core.optimizer import _init_optimizers_and_lr_schedulers
  from lightning.pytorch.plugins.precision import Precision
  from lightning.pytorch.strategies.ddp import DDPStrategy
  from lightning.pytorch.trainer.states import TrainerFn
  from lightning.pytorch.utilities import GradClipAlgorithmType
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from lightning.pytorch.utilities.model_helpers import is_overridden
  from lightning.pytorch.utilities.rank_zero import WarningCache, rank_zero_info, rank_zero_warn
  from lightning.pytorch.utilities.types import LRSchedulerConfig
label_desc: Use image node_id 1 to create a new DeepSpeedStrategy object from inherited base classes: DDPStrategy with example: obj = DeepSpeedStrategy(accelerator, zero_optimization, stage, remote_device, offload_optimizer, offload_parameters, offload_params_device, nvme_path, params_buffer_count, params_buffer_size, max_in_cpu, offload_optimizer_device, optimizer_buffer_count, block_size, queue_depth, single_submit, overlap_events, thread_count, pin_memory, sub_group_size, contiguous_gradients, overlap_comm, allgather_partitions, reduce_scatter, allgather_bucket_size, reduce_bucket_size, zero_allow_untested_optimizer, logging_batch_size_per_gpu, config, logging_level, parallel_devices, cluster_environment, loss_scale, initial_scale_power, loss_scale_window, hysteresis, min_loss_scale, partition_activations, cpu_checkpointing, contiguous_memory_optimization, synchronize_checkpoint_boundary, load_full_weights, precision_plugin, process_group_backend)