class_docstr: To use ZeroRedundancyOptimizer, we need to shard the optimizer states when
  saving/loading checkpoints.
class_docstr_tok: ["To","use","ZeroRedundancyOptimizer",",","we","need","to","shard","the","optimizer","states","when","saving\/loading","checkpoints","."]
code: def optimizer_state(self, optimizer: Optimizer) -> Optional[dict]:
      if isinstance(optimizer, LightningOptimizer):
          optimizer = optimizer._optimizer
      if isinstance(optimizer, ZeroRedundancyOptimizer):
          return get_zero_optimizer_state_dict_local(
              optimizer, self.global_rank
          )
      else:
          return optimizer.state_dict()
  
file_path: flash-attention/training/src/utils/ddp_zero1.py
imports: from typing import Any, Callable, Dict, List, Optional, Union
  from pathlib import Path
  import torch
  from torch.optim.optimizer import Optimizer
  from torch.distributed.optim import ZeroRedundancyOptimizer
  from pytorch_lightning.strategies.ddp import DDPStrategy
  from pytorch_lightning.core.optimizer import LightningOptimizer
label_desc: Use image node_id 1 for calling the DDPStrategyZero1 obj's underlying member method code with example usage: obj.optimizer_state(optimizer) and returns: get_zero_optimizer_state_dict_local, optimizer