code: def all_reduce_async(
      self,
      tensor: Tensor,
      group: ProcessGroup,
      callback_fn: Optional[Callable] = None,
  ) -> None:
      bucket_size = self._get_bucket_size(tensor.element_size())
  
      if tensor.numel() >= bucket_size:
          dist.all_reduce(tensor, group=group)
          if callback_fn is not None:
              callback_fn(tensor)
          return
  
      bucket = self._get_bucket(tensor, group)
      if tensor.numel() > bucket.avail_size:
          # not enough space remaining in bucket, flush it now
          bucket.flush()
      bucket.append(tensor, callback_fn)
  
file_path: ColossalAI/colossalai/legacy/nn/parallel/reducer.py
imports: import functools
  from typing import Callable, Dict, List, Optional, Tuple
  import torch
  import torch.distributed
  from torch import Tensor
  from torch.distributed import ProcessGroup
label_desc: Use image node_id 2 for calling the Reducer obj's underlying member method code with example usage: obj.all_reduce_async(tensor, group, callback_fn) without return types