code: def get_process_group_cached(
      name: ProcessGroupName = ProcessGroupName.default,
      ranks: Optional[Sequence[int]] = None,
  ) -> "ProcessGroup":
      """
      Singleton PyTorch distributed group cache. Inspired by the code from fairseq.
  
      Just like torch.distributed.new_group, this method needs to be called on all ranks
      at the same time when a new group is created. This is true for all ranks irrespective
      of their group membership status.
  
      For FSDP, it is important to use the same group between outer and inner FSDP instances,
      otherwise, inner FSDP instances will not share the gradient reduction bucket buffer with
      the root instance. This will result in increased GPU memory utilization.
  
      Each separate process group also uses separate NCCL library instances, which will have
      a significant effect on GPU memory use if too many process groups are created and used.
      Setting NCCL_BUFFSIZE=102400 env variable is a useful technique to check if the NCCL
      memory is causing GPU OOM. Note, the NCCL buffers are not allocated
      through the PyTorch caching allocator, therefore, you may see GPU OOM even when
      torch.cuda.reserved_memory() is still way below the total amount of GPU memory.
  
      Extra process groups can also reduce training speed (observed on VISSL models).
  
      Args:
          name ProcessGroupName:
              There are two process groups when reduce_scatter overlap is enabled. The "default" process group is the
              default process group. The other group is "reduce_scatter" group.
              Default: ProcessGroupName.default
          ranks (Optional[List[int]]):
              Ranks requested in the target group. None for all ranks.
              Default: None
  
      Returns:
          (ProcessGroup):
              Return the requested process group. Throws RuntimeError if torch.distributed module is not yet initialized.
      """
      if not dist.is_initialized():
          # Likely caused by initiating a dummy pg for unit test, skip checking.
          if (
              name == ProcessGroupName.reduce_scatter
              and "pytest" in sys.modules
          ):
              return None
          else:
              raise RuntimeError(
                  "torch.distributed is not yet initialized but process group is requested."
              )
  
      # Init the cache if needed.
      if not hasattr(get_process_group_cached, "_global_group_cache"):
          get_process_group_cached._global_group_cache = {}  # type: ignore
          # Populate with default process group.
          cache = get_process_group_cached._global_group_cache  # type: ignore
  
          default_pg = dist.new_group(ranks=ranks)
          cache[None] = default_pg
          cache[(ProcessGroupName.default, None)] = default_pg
          cache[
              (
                  ProcessGroupName.default,
                  frozenset(list(range(dist.get_world_size()))),
              )
          ] = default_pg
  
      # Lookup and fill the cache if needed.
      cache = get_process_group_cached._global_group_cache  # type: ignore
      if ranks is not None:
          # take care of ordering and duplicates in the ranks list. use tuple so that ranks
          # can be used as a cache index.
          ranks = tuple(sorted(list(set(ranks))))
      if (name, ranks) not in cache:
          cache[(name, ranks)] = dist.new_group(ranks=ranks)
  
      return cache[(name, ranks)]
  
file_path: fairscale/fairscale/internal/parallel.py
imports: from enum import Enum
  import sys
  from typing import TYPE_CHECKING, List, Optional, Sequence
  import torch
  import torch.distributed
  import torch.nn.functional
label_desc: Use image node_id 4 for calling a global function with example usage: get_process_group_cached(name, ranks) and returns: cache, None