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_cap_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_cap_mb","(","int",",","Optional",")",":","bucket","size","for","communicating",".","Buckets","are","sub-divided","based","on","world_size",".","Values","<","=","0","disable","bucketing","."]
code: def _get_bucket(
      self, tensor: Tensor, group: "ProcessGroup"
  ) -> Bucket:
      # TODO (Min): the `group` used here in the key is the object hash, not the content
      #     hash. That means if FSDP instances are initialized with different process groups,
      #     even when the group members are in fact the same, we end up creating different
      #     buckets here.
      key = (tensor.dtype, tensor.device, group)
      if key not in self.buckets:
          # buckets are divided into world_size pieces, bucket.data shaped (world_size, shard_size)
          world_size = group.size()
          shard_size = self._get_shard_size(
              tensor.element_size(), world_size
          )
          data = tensor.new_zeros((world_size, shard_size))
          self.buckets[key] = Bucket(data, group)
      self.buckets[key].setup()
      return self.buckets[key]
  
file_path: fairscale/fairscale/internal/reduce_scatter_bucketer.py
imports: import functools
  import os
  from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple
  import torch
  from torch import Tensor
  import torch.distributed
label_desc: Use image node_id 6 for calling the ReduceScatterBucketer obj's underlying member method code with example usage: obj._get_bucket(tensor, group) and returns: self