class_docstr: Helper for bucketing multiple reduce-scatter operations on small tensors
  into larger reduce-scatter ops to improve communication efficiency.
  
  Usage::
  
      bucketer = ReduceScatterBucketer()
      bucketer.reduce_scatter_async(
          small_tensors, callback_fn=lambda result: print("small")
      )
      bucketer.reduce_scatter_async(
          big_tensors, callback_fn=lambda result: print("big")
      )
      bucketer.reduce_scatter_async(
          more_small_tensors, callback_fn=lambda result: print("small2")
      )
      bucketer.flush()  # callbacks only guaranteed to be called after flush()
      # Example output (note that it is out of order, due to bucketing):
      # big
      # small
      # small2
  
  Args:
      bucket_size_mb (int, Optional): bucket size for communicating. Buckets
          are sub-divided based on world_size. Values <= 0 disable bucketing.
class_docstr_tok: ["Helper","for","bucketing","multiple","reduce-scatter","operations","on","small","tensors","into","larger","reduce-scatter","ops","to","improve","communication","efficiency",".","Usage",":",":","bucketer","=","ReduceScatterBucketer","(",")","bucketer.reduce_scatter_async","(","small_tensors",",","callback_fn=lambda","result",":","print","(","``","small","''",")",")","bucketer.reduce_scatter_async","(","big_tensors",",","callback_fn=lambda","result",":","print","(","``","big","''",")",")","bucketer.reduce_scatter_async","(","more_small_tensors",",","callback_fn=lambda","result",":","print","(","``","small2","''",")",")","bucketer.flush","(",")","#","callbacks","only","guaranteed","to","be","called","after","flush","(",")","#","Example","output","(","note","that","it","is","out","of","order",",","due","to","bucketing",")",":","#","big","#","small","#","small2","Args",":","bucket_size_mb","(","int",",","Optional",")",":","bucket","size","for","communicating",".","Buckets","are","sub-divided","based","on","world_size",".","Values","<","=","0","disable","bucketing","."]
code: def __init__(self, bucket_size_mb: int = 25):
      self.bucket_size_mb = bucket_size_mb
      self.buckets: Dict[
          Tuple[torch.dtype, torch.device, ProcessGroup], Bucket
      ] = {}
  
file_path: ColossalAI/colossalai/legacy/zero/sharded_model/reduce_scatter.py
imports: import functools
  import os
  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 1 to create a new ReduceScatterBucketer object  with example: obj = ReduceScatterBucketer(bucket_size_mb)