20 lines
468 B
Python
20 lines
468 B
Python
import os
|
|
import torch
|
|
import torch.distributed as dist
|
|
import datetime
|
|
|
|
def main():
|
|
dist.init_process_group("nccl",timeout=datetime.timedelta(seconds=10))
|
|
local_rank = int(os.environ["LOCAL_RANK"])
|
|
torch.cuda.set_device(local_rank)
|
|
|
|
print(f"Rank {dist.get_rank()} initialized on GPU {local_rank}")
|
|
|
|
x = torch.ones(10).cuda()
|
|
dist.all_reduce(x)
|
|
print(f"Rank {dist.get_rank()} result: {x[0].item()}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|