docstr: 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.
  
  Args:
      ckpt: The ckpt file.
docstr_tok: ["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",".","Args",":","ckpt",":","The","ckpt","file","."]
code: def _restore_zero_state(self, ckpt: Mapping[str, Any]) -> None:
      """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.
  
      Args:
          ckpt: The ckpt file.
  
      """
      import deepspeed
  
      assert self.lightning_module is not None
  
      def load(module: torch.nn.Module, prefix: str = "") -> None:
          missing_keys: List[str] = []
          unexpected_keys: List[str] = []
          error_msgs: List[str] = []
          state_dict = ckpt["state_dict"]
  
          # copy state_dict so _load_from_state_dict can modify it
          metadata = getattr(state_dict, "_metadata", None)
          state_dict = state_dict.copy()
          if metadata is not None:
              state_dict._metadata = metadata
  
          local_metadata = (
              {} if metadata is None else metadata.get(prefix[:-1], {})
          )
          # because zero3 puts placeholders in model params, this context
          # manager gathers (unpartitions) the params of the current layer, then loads from
          # the state dict and then re-partitions them again
          with deepspeed.zero.GatheredParameters(
              list(module.parameters(recurse=False)), modifier_rank=0
          ):
              if self.is_global_zero:
                  module._load_from_state_dict(
                      state_dict=state_dict,
                      prefix=prefix,
                      local_metadata=local_metadata,
                      strict=True,
                      missing_keys=missing_keys,
                      unexpected_keys=unexpected_keys,
                      error_msgs=error_msgs,
                  )
  
          for name, child in module._modules.items():
              if child is not None:
                  load(child, prefix + name + ".")
  
      load(self.lightning_module, prefix="")
  
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 27 for calling the DeepSpeedStrategy obj's underlying member method code with example usage: obj._restore_zero_state(ckpt) without return types