code: def maybe_init_dist() -> Optional[int]:
      try:
          # provided by torchrun
          rank = _get_rank()
          world_size = _get_world_size()
  
          if world_size < 2:
              # too few gpus to parallelize, tp is no-op
              return None
      except KeyError:
          # not run via torchrun, no-op
          return None
  
      torch.cuda.set_device(rank)
      dist.init_process_group(
          backend="nccl", rank=rank, world_size=world_size
      )
      return rank
  
file_path: gpt-fast/tp.py
imports: import os
  from typing import List, Optional
  import torch
  import torch.distributed
  from torch import nn
  from torch.distributed import _functional_collectives
  from model import Attention, FeedForward, Transformer
  from quantize import WeightOnlyInt4Linear
label_desc: Use image node_id 5 for calling a global function with example usage: maybe_init_dist() and returns: rank, None, None