INNER CODE UNIT · Python

init_processes

huggingface/large_language_model_training_playbook · throughput/all_reduce_bench.py:60

def init_processes(local_rank, fn, backend='nccl'):
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend)
    fn(local_rank)


if __name__ == "__main__":
    rank = int(os.environ["LOCAL_RANK"])
    printflock("local_rank: %d" % rank)
    init_processes(local_rank=rank, fn=run)

View source record →