class_docstr: Handles partitioning and gathering of parameters.
docstr: for each partitioned parameter, kick off an async allgather and store
  the work handle for the in flight parameters.
docstr_tok: ["for","each","partitioned","parameter",",","kick","off","an","async","allgather","and","store","the","work","handle","for","the","in","flight","parameters","."]
code: def __all_gather_params(self, params: Set[Parameter]) -> None:
      """for each partitioned parameter, kick off an async allgather and store
      the work handle for the in flight parameters."""
      partitioned_params = []
      for param in params:
          if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
              partitioned_params.append(param)
              self.__n_available_params += param.ds_numel
  
      if partitioned_params:
          with torch.cuda.stream(self.__allgather_stream):
              handle = partitioned_params[0].all_gather_coalesced(
                  partitioned_params
              )
  
          for param in partitioned_params:
              assert (
                  param.ds_status == ZeroParamStatus.INFLIGHT
              ), param.ds_summary()
              self.__inflight_param_registry[param] = handle
  
          # Release swap buffers for persisted params on nvme since they will never be partitioned or evicted from GPU
          swap_persisted_params = [
              p
              for p in partitioned_params
              if p.ds_persist
              and p.ds_tensor.final_location == OffloadDeviceEnum.nvme
          ]
          if swap_persisted_params:
              swap_persisted_params[
                  0
              ].nvme_swapper.remove_partition_and_release_buffers(
                  swap_persisted_params
              )
  
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 17 for calling the PartitionedParameterCoordinator obj's underlying member method code with example usage: obj.__all_gather_params(params) without return types