Merge pull request #1 from MotifTechnologies/pre-commit_test_and_apply_lint

.github/workflows/pre-commit.yml

@@ -0,0 +1,30 @@
+name: pre-commit
+
+on:
+  pull_request:
+  push:
+    branches: [ main, master ]
+
+jobs:
+  run-pre-commit:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      pull-requests: read
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Cache pre-commit
+        uses: actions/cache@v4
+        with:
+          path: ~/.cache/pre-commit
+          key: pre-commit-${{ runner.os }}-${{ hashFiles('.pre-commit-config.yaml') }}
+          restore-keys: |
+            pre-commit-${{ runner.os }}-
+
+      - name: Run pre-commit
+        uses: pre-commit/[email protected]

.pre-commit-config.yaml

@@ -0,0 +1,37 @@
+default_install_hook_types:
+  - pre-commit
+  - commit-msg
+default_stages:
+  - pre-commit # Run locally
+  - manual # Run in CI
+exclude: '(build|result)/.*'
+repos:
+- repo: https://github.com/google/yapf
+  rev: v0.43.0
+  hooks:
+  - id: yapf
+    args: [--in-place, --verbose]
+- repo: https://github.com/crate-ci/typos
+  rev: v1.34.0
+  hooks:
+  - id: typos
+    exclude: '.gitattributes'
+- repo: https://github.com/PyCQA/isort
+  rev: 6.0.1
+  hooks:
+  - id: isort
+- repo: https://github.com/pre-commit/mirrors-clang-format
+  rev: v20.1.3
+  hooks:
+  - id: clang-format
+    types_or: [c++, cuda]
+    args: [--style=file, --verbose]
+- repo: https://github.com/jackdewinter/pymarkdown
+  rev: v0.9.29
+  hooks:
+  - id: pymarkdown
+    args: [fix]
+- repo: https://github.com/rhysd/actionlint
+  rev: v1.7.7
+  hooks:
+  - id: actionlint

README.md

@@ -10,7 +10,7 @@ Optimizer is a python package that provides:
 - PyTorch implementation of recent optimizer algorithms
 - with support for parallelism techniques for efficient large-scale training.
 
-### Currently implemented
+## Currently implemented
 - [Parallel Muon with FSDP2](./docs/muon/parallel_muon.pdf)
 
 ## Usage
@@ -31,4 +31,49 @@ optim = optimizer.Muon(
     momentum=0.9,
     weight_decay=1e-4,
 )
-```
+```
+
+## Pre-commit Hooks
+
+This project uses [pre-commit](https://pre-commit.com/) to automatically check and format code before commits.
+
+### Setup
+
+1. Install pre-commit:
+
+   ```bash
+   pip install pre-commit
+   ```
+
+2. Install the git hooks:
+  
+```bash
+   pre-commit install
+   ```
+
+Once installed, the configured hooks will run automatically on each commit.
+
+### Included Hooks
+
+The following tools are run via pre-commit:
+
+- **[yapf](https://github.com/google/yapf)** – Python code formatter  
+- **[typos](https://github.com/crate-ci/typos)** – Spell checker for common typos  
+- **[isort](https://github.com/PyCQA/isort)** – Organizes and sorts Python imports  
+- **[clang-format](https://clang.llvm.org/docs/ClangFormat.html)** – Formats C++/CUDA code (`--style=file`)  
+- **[pymarkdown](https://github.com/jackdewinter/pymarkdown)** – Lints and auto-fixes Markdown files  
+- **[actionlint](https://github.com/rhysd/actionlint)** – Validates GitHub Actions workflows  
+
+### Usage
+
+- Run all checks on the entire codebase:
+
+   ```bash
+   pre-commit run --all-files
+   ```
+
+- Run a specific hook (example: isort):
+  
+ ```bash
+   pre-commit run isort --all-files
+   ```

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"]