37 lines
701 B
Python
37 lines
701 B
Python
import os
|
|
import torch
|
|
import torch.distributed as dist
|
|
import time
|
|
|
|
|
|
def main():
|
|
dist.init_process_group("nccl")
|
|
|
|
local_rank = int(os.environ["LOCAL_RANK"])
|
|
torch.cuda.set_device(local_rank)
|
|
|
|
rank = dist.get_rank()
|
|
world_size = dist.get_world_size()
|
|
|
|
print(f"[rank {rank}/{world_size}] started")
|
|
|
|
model = torch.nn.Linear(16, 16).cuda()
|
|
|
|
for step in range(50):
|
|
x = torch.randn(32, 16).cuda()
|
|
loss = model(x).sum()
|
|
|
|
loss.backward()
|
|
|
|
if rank == 0 and step % 10 == 0:
|
|
print(f"step={step}, loss={loss.item()}")
|
|
|
|
time.sleep(0.2)
|
|
|
|
dist.barrier()
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|