class_docstr: Handles partitioning and gathering of parameters.
code: def trace_prologue(self, sub_module: Module) -> None:
      if self.is_complete_trace():
          # sub_module must match expectation else invalidate trace cache
          if sub_module != self.__submodule_order[self.__step_id]:
              expected_module_id = self.__submodule_order[
                  self.__step_id
              ].id
              debug_rank0(
                  f"Invalidate trace cache @ step {self.__step_id}: "
                  f"expected module {expected_module_id}, but got module {sub_module.id}"
              )
              self._invalidate_trace()
  
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 7 for calling the PartitionedParameterCoordinator obj's underlying member method code with example usage: obj.trace_prologue(sub_module) without return types