"""torchrun --standalone --nproc_per_node=2 examples/12_ddp_minimal.py"""
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP


def main():
    dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
    rank = dist.get_rank()
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
    if device.type == "cuda":
        torch.cuda.set_device(device)
    torch.manual_seed(1234)
    model = DDP(torch.nn.Linear(4, 1).to(device), device_ids=[local_rank] if device.type == "cuda" else None)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    x = torch.full((2, 4), float(rank + 1), device=device)
    y = torch.zeros((2, 1), device=device)
    loss = torch.nn.functional.mse_loss(model(x), y)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    if rank == 0:
        print("world_size=", dist.get_world_size(), "loss=", float(loss))
    dist.destroy_process_group()


if __name__ == "__main__":
    main()