import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim


def main():
    dist.init_process_group("nccl")

    rank = dist.get_rank()
    torch.cuda.set_device(0)

    model = nn.Linear(10, 10).cuda()
    ddp = torch.nn.parallel.DistributedDataParallel(model, device_ids=[0])

    opt = optim.SGD(ddp.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()

    for i in range(5):
        x = torch.randn(32, 10).cuda()
        y = torch.randn(32, 10).cuda()

        opt.zero_grad()
        out = ddp(x)
        loss = loss_fn(out, y)
        loss.backward()
        opt.step()

        print(f"rank {rank} step {i} loss {loss.item()}")

    dist.destroy_process_group()


if __name__ == "__main__":
    main()