class_docstr: Optimizer used for ZeRO-1 and ZeRO-2.
docstr: Returns dictionaries containing a whole state of the module one by one. The max size of dictionary shard is specified by ``max_shard_size``.
     Only include the 'state' in state_dict.
  
  Args:
      max_shard_size (int, optional): max size of state shard (in MB). Defaults to 1024.
  
  Yields:
      Iterator[OrderedDict]: A generator of state dict shard
docstr_tok: ["Returns","dictionaries","containing","a","whole","state","of","the","module","one","by","one",".","The","max","size","of","dictionary","shard","is","specified","by","``","max_shard_size","``",".","Only","include","the","'state","'","in","state_dict",".","Args",":","max_shard_size","(","int",",","optional",")",":","max","size","of","state","shard","(","in","MB",")",".","Defaults","to","1024",".","Yields",":","Iterator","[","OrderedDict","]",":","A","generator","of","state","dict","shard"]
code: def state_dict_shard(
      self, max_shard_size: int = 1024
  ) -> Iterator[Tuple[Dict, int]]:
      """Returns dictionaries containing a whole state of the module one by one. The max size of dictionary shard is specified by ``max_shard_size``.
         Only include the 'state' in state_dict.
  
      Args:
          max_shard_size (int, optional): max size of state shard (in MB). Defaults to 1024.
  
      Yields:
          Iterator[OrderedDict]: A generator of state dict shard
      """
      ret_block = dict()
      ret_block_size = 0
  
      device = get_current_device()
      local_states = self.optim.state_dict()["state"]
      for param_idx, states in local_states.items():
          current_block_size = 0
          current_block = copy.deepcopy(states)
  
          # find the working param of current param_id
          for (
              group_id,
              pg,
          ) in self._master_param_groups_of_current_rank.items():
              if (group_id + 1) * len(pg) < param_idx:
                  continue
              master_param = pg[param_idx - (group_id) * len(pg)]
              working_param = self._param_store.master_to_working_param[
                  id(master_param)
              ]
  
          for k, v in states.items():
              if isinstance(v, torch.Tensor) and k != "step":
                  if (
                      self.moe_extra_dp_pg is not None
                      and is_moe_tensor(v)
                  ):
                      state_tensor = [
                          torch.zeros(
                              v.shape, device=device, dtype=v.dtype
                          )
                          for _ in range(self.moe_extra_dp_pg_size)
                      ]
                      dist.all_gather(
                          state_tensor,
                          v.to(device),
                          group=self.moe_extra_dp_pg,
                      )
                  else:
                      state_tensor = [
                          torch.zeros(
                              v.shape, device=device, dtype=v.dtype
                          )
                          for _ in range(self._world_size)
                      ]
                      dist.all_gather(
                          state_tensor, v.to(device), group=self.dp_pg
                      )
                  state_tensor = (
                      torch.stack(state_tensor)
                      .view(-1)[: working_param.numel()]
                      .reshape_as(working_param)
                      .cpu()
                  )
                  current_block_size += state_tensor.numel()
                  current_block[k] = state_tensor
  
          if (
              ret_block_size + current_block_size > max_shard_size
              and len(ret_block) > 0
          ):
              yield ret_block, ret_block_size
              ret_block = dict()
              ret_block_size = 0
  
          ret_block[param_idx] = current_block
          ret_block_size += current_block_size
  
      yield ret_block, ret_block_size
  
imports: import copy
  from contextlib import contextmanager
  from functools import partial
  from typing import Dict, Iterator, List, Optional, Tuple
  import torch
  import torch.distributed
  import torch.nn
  from torch import Tensor, inf
  from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
  from torch.distributed import ProcessGroup
  from torch.optim import Optimizer
  import colossalai.utils.device
  from colossalai.amp.naive_amp.mixed_precision_mixin import BF16MixedPrecisionMixin, FP16MixedPrecisionMixin, MixedPrecisionMixin
  from colossalai.interface import OptimizerWrapper
  from colossalai.logging import get_dist_logger
  from colossalai.tensor.moe_tensor.api import is_moe_tensor
  from colossalai.utils.device import IS_NPU_AVAILABLE, get_current_device
  from ._utils import calculate_global_norm_from_list, flatten, has_inf_or_nan, release_param_grad, sync_tensor
  from .bookkeeping import BucketStore, GradientStore, ParameterStore
label_desc: Use image node_id 25 for calling the LowLevelZeroOptimizer obj's underlying member method code with example usage: obj.state_dict_shard(max_shard_size) without return types