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()