Anonymous-123/ImageNet-Editing

1"""2Helpers for distributed training.3"""4 5import io6import os7import socket8 9import blobfile as bf10from mpi4py import MPI11import torch as th12import torch.distributed as dist13 14# Change this to reflect your cluster layout.15# The GPU for a given rank is (rank % GPUS_PER_NODE).16GPUS_PER_NODE = 817 18SETUP_RETRY_COUNT = 319 20 21def setup_dist():22    """23    Setup a distributed process group.24    """25    if dist.is_initialized():26        return27    os.environ["CUDA_VISIBLE_DEVICES"] = f"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}"28 29    comm = MPI.COMM_WORLD30    backend = "gloo" if not th.cuda.is_available() else "nccl"31 32    if backend == "gloo":33        hostname = "localhost"34    else:35        hostname = socket.gethostbyname(socket.getfqdn())36    os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0)37    os.environ["RANK"] = str(comm.rank)38    os.environ["WORLD_SIZE"] = str(comm.size)39 40    port = comm.bcast(_find_free_port(), root=0)41    os.environ["MASTER_PORT"] = str(port)42    dist.init_process_group(backend=backend, init_method="env://")43 44 45def dev():46    """47    Get the device to use for torch.distributed.48    """49    if th.cuda.is_available():50        return th.device(f"cuda")51    return th.device("cpu")52 53 54def load_state_dict(path, **kwargs):55    """56    Load a PyTorch file without redundant fetches across MPI ranks.57    """58    chunk_size = 2 ** 30  # MPI has a relatively small size limit59    if MPI.COMM_WORLD.Get_rank() == 0:60        with bf.BlobFile(path, "rb") as f:61            data = f.read()62        num_chunks = len(data) // chunk_size63        if len(data) % chunk_size:64            num_chunks += 165        MPI.COMM_WORLD.bcast(num_chunks)66        for i in range(0, len(data), chunk_size):67            MPI.COMM_WORLD.bcast(data[i : i + chunk_size])68    else:69        num_chunks = MPI.COMM_WORLD.bcast(None)70        data = bytes()71        for _ in range(num_chunks):72            data += MPI.COMM_WORLD.bcast(None)73 74    return th.load(io.BytesIO(data), **kwargs)75 76 77def sync_params(params):78    """79    Synchronize a sequence of Tensors across ranks from rank 0.80    """81    for p in params:82        with th.no_grad():83            dist.broadcast(p, 0)84 85 86def _find_free_port():87    try:88        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)89        s.bind(("", 0))90        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)91        return s.getsockname()[1]92    finally:93        s.close()94