class_docstr: Handles partitioning and gathering of parameters.
docstr: indicate that we have completed one fwd+bwd for the model
docstr_tok: ["indicate","that","we","have","completed","one","fwd+bwd","for","the","model"]
code: def reset_step(self) -> None:
      """indicate that we have completed one fwd+bwd for the model"""
      if self.__inflight_param_registry:
          raise RuntimeError(
              f"still have inflight params "
              f"{[p.ds_summary for p in self.__inflight_param_registry.keys()]}"
          )
  
      if not self.is_complete_trace():  # not self.trace_complete:
          # Make sure that recorded submodule orders are identical across ranks
          assert_ints_same_as_other_ranks(
              [m.id for m in self.__submodule_order]
          )
  
          if self.is_record_trace():
              # Successfully recorded a trace
              self.construct_parameter_trace_from_module_trace()
              # Make sure that recorded parameter orders are identical across ranks
              assert_ints_same_as_other_ranks(
                  [p.param.ds_id for p in self.__param_order]
              )
              assert_ints_same_as_other_ranks(
                  [p.step_id_last_used_at for p in self.__param_order]
              )
  
              self.__submodule_order = tuple(
                  self.__submodule_order
              )  # freeze
              self.__param_order = tuple(self.__param_order)  # freeze
              self.__trace_mode = ZeRoTraceMode.COMPLETE
              print_rank_0(
                  f"completed record trace: {[m.id for m in self.__submodule_order]}",
                  force=False,
              )
          else:
              # Enable trace recording for next forward/backward pass
              self.__trace_mode = ZeRoTraceMode.RECORD
  
      self.__param_queue = collections.deque(
          self.__param_order
      )  # reset fetch queue
      self.__most_recent_step_id_param_fetched_for = (
          collections.defaultdict(lambda: int(-1e10))
      )
      self.__step_id_module_fetched_for = collections.defaultdict(
          lambda: collections.deque()
      )
      self.__step_id = 0
      self.__n_available_params = 0
  
imports: from dataclasses import dataclass
  import collections
  from collections import UserDict
  from typing import Deque, Set
  from torch.cuda import Event, Stream
  from deepspeed import comm
  from deepspeed.utils.logging import logger
  from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum
  from deepspeed.runtime.zero.partition_parameters import *
  from deepspeed.runtime.swap_tensor.partitioned_param_swapper import PartitionedParamStatus
  from deepspeed.utils.debug import debug_module2name_id, debug_param2name_id
label_desc: Use image node_id 11 for calling the PartitionedParameterCoordinator obj's underlying member method code with example usage: obj.reset_step() without return types