chore: add .gitignore

.gitignore

@@ -0,0 +1,22 @@
+__pycache__
+.idea
+.DS_Store
+*.egg-info
+build
+outputs
+dist/*
+.vscode
+
+# data
+data
+out
+wandb
+
+torchtitan/datasets/**/*.model
+torchtitan/experiments/flux/assets/*
+
+# temp files
+*.log
+error.json
+_remote_module_non_scriptable.py
+.git_disabled/

torch-ext/optimizer/muon.py

@@ -83,6 +83,7 @@ def _gather(p, state, rank, comm_stream, none_grad):
             state.gathered_grad = None
             state.gather_event = None
         if none_grad:
+            p.grad.record_stream(comm_stream)
             p.grad = None
 
 
@@ -98,6 +99,7 @@ def _compute_u(state, steps, rank, compute_stream):
             state.compute_event = torch.cuda.Event()
             state.compute_event.record()
             # Clear the gathered gradient to free memory
+            state.gathered_grad.record_stream(compute_stream)
             state.gathered_grad = None
         else:
             state.computed_u = None
@@ -106,7 +108,6 @@ def _compute_u(state, steps, rank, compute_stream):
 
 @torch.no_grad()
 def _scatter(p, state, lr, adjusted_lr, weight_decay, rank, comm_stream):
-    u = state.computed_u
 
     with torch.cuda.stream(comm_stream):
         if rank == state.worker_rank:
@@ -114,27 +115,31 @@ def _scatter(p, state, lr, adjusted_lr, weight_decay, rank, comm_stream):
             if state.compute_event is None:
                 raise RuntimeError("Compute event must be set before scatter.")
             comm_stream.wait_event(state.compute_event)
+            u = state.computed_u
             scatter_list = list(torch.split(u, p.size(0) // num_ranks, dim=0))
+            scatter_list = [s.contiguous() for s in scatter_list]
         else:
             scatter_list = None
 
-        u = torch.empty_like(p.to_local())
+        u_received = torch.empty_like(p.to_local())
         torch.distributed.scatter(
-            u,
+            u_received,
             scatter_list=scatter_list,
             src=state.worker_rank,
             group=state.process_group,
         )
         if rank == state.worker_rank:
             # Clear u to free memory
+            state.computed_u.record_stream(comm_stream)
             state.computed_u = None
-        u = DTensor.from_local(
-            u,
+
+        u_dtensor = DTensor.from_local(
+            u_received,
             placements=p.placements,
             device_mesh=p.device_mesh,
         )
         p.data.mul_(1 - lr * weight_decay)
-        p.data.add_(u, alpha=-adjusted_lr)
+        p.data.add_(u_dtensor, alpha=-adjusted_lr)
 
 
 def default_is_muon(x, name):