update dags

dags/train.py

@@ -0,0 +1,35 @@
+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()