maybe_init_distributed() — pytorch Function Reference

Architecture documentation for the maybe_init_distributed() function in common.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking called by 1

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ maybe_init_distributed() — pytorch Function Reference

Dependency Diagram

graph TD
  b723bea5_3f53_dc94_5332_34181acc6752["maybe_init_distributed()"]
  27245bba_7bce_3446_86ed_29630ca8fe8a["process_entry()"]
  27245bba_7bce_3446_86ed_29630ca8fe8a -->|calls| b723bea5_3f53_dc94_5332_34181acc6752
  style b723bea5_3f53_dc94_5332_34181acc6752 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/common.py lines 1743–1755

def maybe_init_distributed(should_init_distributed, rank, world_size, port="6789"):
    try:
        if should_init_distributed:
            torch.cuda.set_device(rank)
            os.environ["MASTER_ADDR"] = "localhost"
            os.environ["MASTER_PORT"] = port
            torch.distributed.init_process_group(
                "nccl", rank=rank, world_size=world_size
            )
        yield
    finally:
        if should_init_distributed:
            torch.distributed.destroy_process_group()