class_docstr: To use Apex's DistributedFusedAdam, we need to shard the optimizer states when
  saving/loading checkpoints.
docstr: Save model/training states as a checkpoint file through state-dump and file-write.
  Args:
      checkpoint: dict containing model and trainer state
      filepath: write-target file's path
      storage_options: parameter for how to save to storage, passed to ``CheckpointIO`` plugin
docstr_tok: ["Save","model\/training","states","as","a","checkpoint","file","through","state-dump","and","file-write",".","Args",":","checkpoint",":","dict","containing","model","and","trainer","state","filepath",":","write-target","file","'s","path","storage_options",":","parameter","for","how","to","save","to","storage",",","passed","to","``","CheckpointIO","``","plugin"]
code: def save_checkpoint(
      self,
      checkpoint: Dict[str, Any],
      filepath: _PATH,
      storage_options: Optional[Any] = None,
  ) -> None:
      """Save model/training states as a checkpoint file through state-dump and file-write.
      Args:
          checkpoint: dict containing model and trainer state
          filepath: write-target file's path
          storage_options: parameter for how to save to storage, passed to ``CheckpointIO`` plugin
      """
      filepath = Path(filepath)
      filepath.mkdir(parents=True, exist_ok=True)
      local_optimizer_states = checkpoint.pop("optimizer_states")
      if self.is_global_zero:
          self.checkpoint_io.save_checkpoint(
              checkpoint,
              filepath / "model_states.pt",
              storage_options=storage_options,
          )
      self.checkpoint_io.save_checkpoint(
          local_optimizer_states,
          filepath / f"{self.global_rank:03d}_optim_states.pt",
          storage_options=storage_options,
      )
  
imports: from typing import Any, Callable, Dict, List, Optional, Union
  from pathlib import Path
  import types
  import torch
  from torch.optim.optimizer import Optimizer
  from torch.optim import LBFGS
  from apex.contrib.optimizers.distributed_fused_adam import DistributedFusedAdam
  from pytorch_lightning.strategies.ddp import DDPStrategy
  from pytorch_lightning.plugins.precision import PrecisionPlugin, NativeMixedPrecisionPlugin
  from pytorch_lightning.core.optimizer import LightningOptimizer
  from pytorch_lightning.utilities.exceptions import MisconfigurationException
label_desc: Use image node_id 5 for calling the DDPStrategyZero2 obj's underlying member method code with example usage: obj.save_checkpoint(checkpoint, filepath, storage_options) without return types