Add GitHub workflow for single GPU tests

.github/workflows/gpu-tests.yaml

@@ -0,0 +1,26 @@
+name: gpu-tests
+
+on: [push, pull_request]
+
+jobs:
+  gpu-tests:
+    name: "Python 3.10"
+    runs-on: 4-core-ubuntu-gpu-t4
+    steps:
+        - uses: actions/checkout@v4
+        - name: Set up and update uv.
+          run: |
+            curl -LsSf https://astral.sh/uv/install.sh | sh
+            source $HOME/.cargo/env
+            uv self update
+        - name: Install Python.
+          run: uv python install 3.10
+        - name: Create venv and install the package.
+          run: |
+            uv venv && source .venv/bin/activate
+            uv pip install setuptools
+            uv pip install -e ".[dev]"
+        - name: Run tests.
+          run: |
+            source .venv/bin/activate
+            make test-gpu

CONTRIBUTING.md

@@ -11,7 +11,7 @@ We actively welcome your pull requests for existing optimizers.
 1. Fork the repo and create your branch from `main`. Install the package inside of your Python environment with `pip install -e ".[dev]"`.
 2. If you've added code that should be tested, add tests.
 3. If you've changed APIs, update the documentation.
-4. Ensure the test suite passes. To run the subset of the tests that can be run on CPU use `make test`; to run the subset of tests that require four GPUs use `make test-gpu`.
+4. Ensure the test suite passes. To run the subset of the tests that can be run on CPU use `make test`; to run the tests for a single GPU use `make test-gpu` and to run the subset of tests that require 2-4 GPUs use `make test-multi-gpu`.
 5. Make sure your code lints. You can use `make lint` and `make format` to automatically lint and format the code where possible.
 6. If you haven't already, complete the Contributor License Agreement ("CLA").
 

README.md

@@ -3,6 +3,7 @@
 [![Python
 3.10 | 3.11 | 3.12](https://img.shields.io/badge/python-3.10_|_3.11_|_3.12-blue.svg)](https://www.python.org/downloads/)
 ![tests](https://github.com/facebookresearch/optimizers/actions/workflows/tests.yaml/badge.svg)
+![gpu-tests](https://github.com/facebookresearch/optimizers/actions/workflows/gpu-tests.yaml/badge.svg)
 ![lint-ruff](https://github.com/facebookresearch/optimizers/actions/workflows/lint-ruff.yaml/badge.svg)
 ![format-ruff](https://github.com/facebookresearch/optimizers/actions/workflows/format-ruff.yaml/badge.svg)
 ![format-usort](https://github.com/facebookresearch/optimizers/actions/workflows/format-usort.yaml/badge.svg)

makefile

@@ -14,6 +14,10 @@ test:
 	@torchrun --standalone --nnodes=1 --nproc_per_node=2 -m unittest distributed_shampoo/utils/gpu_tests/shampoo_ddp_distributor_test.py
 
 test-gpu:
+	@python3 -m unittest discover -s distributed_shampoo/gpu_tests/ -p "*_test.py"
+	@python3 -m unittest distributed_shampoo/utils/gpu_tests/shampoo_dist_utils_test.py
+
+test-multi-gpu:
 	@torchrun --standalone --nnodes=1 --nproc_per_node=2 -m unittest discover -s distributed_shampoo/gpu_tests/ -p "*_test.py"
 	@torchrun --standalone --nnodes=1 --nproc_per_node=2 -m unittest distributed_shampoo/utils/gpu_tests/shampoo_dist_utils_test.py
 	@torchrun --standalone --nnodes=1 --nproc_per_node=2 -m unittest distributed_shampoo/utils/gpu_tests/shampoo_ddp_distributor_test.py