update dags

dags/train.py

@@ -1,33 +1,34 @@
+import os
 import torch
 import torch.distributed as dist
-import torch.nn as nn
-import torch.optim as optim
+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()
-    torch.cuda.set_device(0)
+    world_size = dist.get_world_size()
 
-    model = nn.Linear(10, 10).cuda()
-    ddp = torch.nn.parallel.DistributedDataParallel(model, device_ids=[0])
+    print(f"[rank {rank}/{world_size}] started")
 
-    opt = optim.SGD(ddp.parameters(), lr=0.01)
-    loss_fn = nn.MSELoss()
+    model = torch.nn.Linear(16, 16).cuda()
 
-    for i in range(5):
-        x = torch.randn(32, 10).cuda()
-        y = torch.randn(32, 10).cuda()
+    for step in range(50):
+        x = torch.randn(32, 16).cuda()
+        loss = model(x).sum()
 
-        opt.zero_grad()
-        out = ddp(x)
-        loss = loss_fn(out, y)
         loss.backward()
-        opt.step()
 
-        print(f"rank {rank} step {i} loss {loss.item()}")
+        if rank == 0 and step % 10 == 0:
+            print(f"step={step}, loss={loss.item()}")
 
+        time.sleep(0.2)
+
+    dist.barrier()
     dist.destroy_process_group()