applied lint

.pre-commit-config.yaml

@@ -15,6 +15,7 @@ repos:
   rev: v1.34.0
   hooks:
   - id: typos
+    exclude: '.gitattributes'
 - repo: https://github.com/PyCQA/isort
   rev: 6.0.1
   hooks:

optimizer/dummy.cu

@@ -3,4 +3,4 @@ namespace {
 __global__ void dummy() {
   // This kernel does nothing but serves as a placeholder
 }
-}
+} // namespace

torch-ext/optimizer/muon.py

@@ -59,7 +59,9 @@ def _gather(p, state, rank, comm_stream, none_grad):
 
     if rank == state.worker_rank:
         num_ranks = dist.get_world_size(group=state.process_group)
-        gather_list = [torch.empty_like(g.to_local()) for _ in range(num_ranks)]
+        gather_list = [
+            torch.empty_like(g.to_local()) for _ in range(num_ranks)
+        ]
     else:
         gather_list = None
 
@@ -73,8 +75,7 @@ def _gather(p, state, rank, comm_stream, none_grad):
         if rank == state.worker_rank:
             if state.gathered_grad is not None:
                 raise RuntimeError(
-                    "Gather event already exists, which should not happen."
-                )
+                    "Gather event already exists, which should not happen.")
             state.gathered_grad = torch.cat(gather_list, dim=0)
             state.gather_event = torch.cuda.Event()
             state.gather_event.record()
@@ -240,9 +241,10 @@ class Muon(torch.optim.Optimizer):
         """
         Get the shard mesh for a parameter p on the given rank.
         """
-        assert isinstance(p, DTensor), "Parallel Muon only supports DTensor parameters."
+        assert isinstance(
+            p, DTensor), "Parallel Muon only supports DTensor parameters."
 
-        if p.placements == (Shard(dim=0),):
+        if p.placements == (Shard(dim=0), ):
             # Case for FSDP
             return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
         elif p.placements == (Replicate(), Shard(dim=0)):
@@ -269,11 +271,12 @@ class Muon(torch.optim.Optimizer):
             total_flops += flops
 
         if self.debug:
-            print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs", flush=True)
+            print(f"Total TFLOPs for Muon: {total_flops / 1e12:.2f} TFLOPs",
+                  flush=True)
 
-        ordered_params = sorted(
-            params, key=lambda p: param_to_flops[id(p)], reverse=True
-        )
+        ordered_params = sorted(params,
+                                key=lambda p: param_to_flops[id(p)],
+                                reverse=True)
 
         round_robin = 0
         mesh = None
@@ -369,28 +372,29 @@ class Muon(torch.optim.Optimizer):
             p.grad = g
 
         param_to_state, ordered_params = self.init_state_and_assign_params(
-            params, group
-        )
+            params, group)
 
         def enqueue_gathers(start_idx, chunk_size):
-            for p in ordered_params[start_idx : start_idx + chunk_size]:
+            for p in ordered_params[start_idx:start_idx + chunk_size]:
                 state = param_to_state[id(p)]
-                _gather(p, state, self.rank, self.comm_stream, group["none_grad"])
+                _gather(p, state, self.rank, self.comm_stream,
+                        group["none_grad"])
 
         def enqueue_computes(start_idx, chunk_size):
-            for p in ordered_params[start_idx : start_idx + chunk_size]:
+            for p in ordered_params[start_idx:start_idx + chunk_size]:
                 state = param_to_state[id(p)]
-                _compute_u(state, group["ns_steps"], self.rank, self.compute_stream)
+                _compute_u(state, group["ns_steps"], self.rank,
+                           self.compute_stream)
 
         def enqueue_scatters(start_idx, chunk_size):
-            for p in ordered_params[start_idx : start_idx + chunk_size]:
+            for p in ordered_params[start_idx:start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(
-                    p, state, lr, adjusted_lr, weight_decay, self.rank, self.comm_stream
-                )
+                _scatter(p, state, lr, adjusted_lr, weight_decay, self.rank,
+                         self.comm_stream)
 
-        chunk_size = dist.get_world_size(param_to_state[id(params[0])].process_group)
+        chunk_size = dist.get_world_size(param_to_state[id(
+            params[0])].process_group)
 
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
@@ -436,15 +440,16 @@ class Muon(torch.optim.Optimizer):
                     continue
                 if isinstance(p.data, DTensor):
                     if all(
-                        isinstance(placement, Replicate) for placement in p.placements
-                    ):
+                            isinstance(placement, Replicate)
+                            for placement in p.placements):
                         param_tensors.append(p)
                     else:
                         param_dtensors.append(p)
                 elif isinstance(p.data, torch.Tensor):
                     param_tensors.append(p)
                 else:
-                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+                    raise TypeError(
+                        f"Unsupported parameter type: {type(p.data)}")
 
             if self.debug:
                 print(
@@ -479,7 +484,9 @@ class Muon(torch.optim.Optimizer):
             #       AdamW backup       #
             ############################
 
-            params = [p for p in group["params"] if not self.state[p]["use_muon"]]
+            params = [
+                p for p in group["params"] if not self.state[p]["use_muon"]
+            ]
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]