docstr: Load a checkpoint.
  
  For DeepSpeed, we need every worker to independently load back the model weights, as the checkpoints themselves
  may be sharded (when using DeepSpeed Zero3).
  
  https://deepspeed.readthedocs.io/en/latest/model-checkpointing.html#loading-training-checkpoints
docstr_tok: ["Load","a","checkpoint",".","For","DeepSpeed",",","we","need","every","worker","to","independently","load","back","the","model","weights",",","as","the","checkpoints","themselves","may","be","sharded","(","when","using","DeepSpeed","Zero3",")",".","https",":","\/\/deepspeed.readthedocs.io\/en\/latest\/model-checkpointing.html","#","loading-training-checkpoints"]
code: def load(
      self, save_path: str, device: Optional[torch.device] = None
  ) -> bool:
      """Load a checkpoint.
  
      For DeepSpeed, we need every worker to independently load back the model weights, as the checkpoints themselves
      may be sharded (when using DeepSpeed Zero3).
  
      https://deepspeed.readthedocs.io/en/latest/model-checkpointing.html#loading-training-checkpoints
      """
      # NOTE(geoffrey): `load_module_strict=False` because this code path is frequently used to load models trained
      # using adapter-based fine-tuning, where the checkpoints only contain the adapter weights, and not the full
      # model weights. This may lead to silent, unexpected behavior for resuming full model fine-tuning,
      # where all the model weights *must* be loaded in.
      # TODO(geoffrey): Add a boolean arg to function to control load_module_strict behavior.
      _, client_state = self.model.load_checkpoint(
          save_path,
          load_lr_scheduler_states=False,
          load_module_strict=False,
      )
      self.global_step = self._get_global_step(client_state, save_path)
      if (
          self.scheduler is not None
          and "scheduler_state" in client_state
      ):
          self.scheduler.load_state_dict(
              client_state["scheduler_state"]
          )
      return True
  
imports: import logging
  import os
  import warnings
  from typing import Any, Dict, List, Mapping, Optional, Tuple, TYPE_CHECKING, Union
  import deepspeed
  import deepspeed.comm
  import torch
  from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
  from packaging import version
  from torch import nn
  from torch.optim.optimizer import Optimizer
  from ludwig.constants import MIN_POSSIBLE_BATCH_SIZE
  from ludwig.distributed.ddp import DDPStrategy
  from ludwig.modules.optimization_modules import get_optimizer_class_and_kwargs
  from ludwig.utils.checkpoint_utils import Checkpoint
  from ludwig.utils.model_utils import extract_tensors, replace_tensors
label_desc: Use image node_id 2 for calling the DeepSpeedCheckpoint obj's underlying member method code with example usage: obj.load(save_path, device) and returns: True