modify pre step (overlap step) can get from arsgs

torch-ext/optimizer/muon.py

@@ -516,6 +516,11 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
+        overlap_step : How many all2all gather, compute operations are launched in advance
+                       before the corresponding all2all scatter steps begin.
+                       A higher overlap_step increases memory usage but can improve
+                       performance by overlapping communication.
+                       Parallel muon only. 
     """
 
     def __init__(self,
@@ -534,7 +539,8 @@ class Muon(torch.optim.Optimizer):
                      "k_indices": [],
                      "head_dim": 128,
                      "threshold": 100
-                 }):
+                 },
+                 overlap_step=5):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -564,6 +570,7 @@ class Muon(torch.optim.Optimizer):
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
+        self.overlap_step = overlap_step
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -821,16 +828,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        PRE_STEP = 5
-        for i in range(0, PRE_STEP):
+        overlap_step = self.overlap_step
+        for i in range(0, overlap_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_computes(i + overlap_step * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)