diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index cafcef4e..e3cb92f9 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -4,9 +4,6 @@ on: push: branches: - main - pull_request: - branches: - - main jobs: build: runs-on: ubuntu-latest @@ -29,10 +26,10 @@ jobs: run: pip wheel --no-deps -w dist . release: needs: build + if: github.repository == 'EleutherAI/bergson' && github.event_name == 'push' && github.ref == 'refs/heads/main' && !contains(github.event.head_commit.message, 'chore(release):') permissions: contents: write id-token: write - if: github.event_name == 'push' && github.ref == 'refs/heads/main' && !contains(github.event.head_commit.message, 'chore(release):') runs-on: ubuntu-latest concurrency: release steps: diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml new file mode 100644 index 00000000..9f944241 --- /dev/null +++ b/.github/workflows/unit_tests.yml @@ -0,0 +1,62 @@ +name: Unit Tests +env: + UV_SYSTEM_PYTHON: 1 +on: + push: + branches: + - 'main' + pull_request: + branches: + - 'main' + workflow_dispatch: + +jobs: + linter: + name: Linters + runs-on: ubuntu-latest + timeout-minutes: 5 + + steps: + - uses: actions/checkout@v5 + - name: "Set up Python" + uses: actions/setup-python@v6 + with: + python-version: "3.10" + - name: Checkout + uses: astral-sh/setup-uv@v6 + with: + enable-cache: true + - name: Install dependencies + run: uv pip install -e ".[dev]" --torch-backend=auto + - name: Pre-Commit + uses: pre-commit/action@v3.0.1 + - name: Type Checking + uses: jakebailey/pyright-action@v2 + - name: Cleanup + run: uv cache prune --ci + + testcpu: + name: CPU Tests + runs-on: ubuntu-latest + strategy: + fail-fast: true + matrix: + python-version: [ "3.10", "3.11", "3.12" ] + + steps: + - uses: actions/checkout@v5 + + - name: Checkout + uses: astral-sh/setup-uv@v6 + with: + python-version: ${{ matrix.python-version }} + enable-cache: true + + - name: Install dependencies + run: uv sync --extra dev + + - name: Run tests + run: uv run pytest tests --showlocals -s -vv -n=auto + + - name: Cleanup + run: uv cache prune --ci diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b58386c7..2f1d88d8 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,18 +1,16 @@ # See https://pre-commit.com for more information # See https://pre-commit.com/hooks.html for more hooks repos: -- repo: https://github.com/pre-commit/pre-commit-hooks - rev: v5.0.0 + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v6.0.0 hooks: - - id: trailing-whitespace - - id: end-of-file-fixer - exclude: \.txt$ -- repo: https://github.com/psf/black - rev: 25.1.0 + - id: trailing-whitespace + - id: end-of-file-fixer + exclude: \.txt$ + - id: no-commit-to-branch + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.14.5 hooks: - - id: black -- repo: https://github.com/astral-sh/ruff-pre-commit - rev: 'v0.11.9' - hooks: - - id: ruff - args: [--fix, --exit-non-zero-on-fix] + - id: ruff-check + args: [ --fix ] + - id: ruff-format diff --git a/bergson/__main__.py b/bergson/__main__.py index 7b4c55cc..789e9eab 100644 --- a/bergson/__main__.py +++ b/bergson/__main__.py @@ -57,8 +57,7 @@ def execute(self): if self.index_cfg.projection_dim != 0: print( - "Warning: projection_dim is not 0. " - "Compressed gradients will be scored." + "Warning: projection_dim is not 0. Compressed gradients will be scored." ) score_dataset(self.index_cfg, self.score_cfg) diff --git a/bergson/gradients.py b/bergson/gradients.py index 992a335f..27533640 100644 --- a/bergson/gradients.py +++ b/bergson/gradients.py @@ -155,9 +155,9 @@ def to_adafactor(self) -> AdafactorNormalizer: and the factored second moments. """ # We assume avg_sq is a square matrix of shape [O, I] - assert ( - self.avg_sq.ndim == 2 - ), f"Expected 2D tensor for avg_sq, got {self.avg_sq.ndim}D" + assert self.avg_sq.ndim == 2, ( + f"Expected 2D tensor for avg_sq, got {self.avg_sq.ndim}D" + ) # Compute row and column means return AdafactorNormalizer( @@ -563,7 +563,6 @@ def _process_grad(self, module: nn.Module, _, grad_out): # If we are using AdamNormalizer, or including bias gradients # we need to materialize the full gradient and then project if isinstance(norm, AdamNormalizer) or include_bias: - P = G.mT @ I # [N, O, S] @ [N, S, I] → [N, O, I] if include_bias: # Append the bias gradient to the input diff --git a/bergson/huggingface.py b/bergson/huggingface.py index 8dbb9693..7c76a502 100644 --- a/bergson/huggingface.py +++ b/bergson/huggingface.py @@ -243,9 +243,9 @@ def on_step_end( # Record training order if enabled if self.order is not None: - assert ( - self.batch_indices is not None - ), "Batch indices are not available for training order tracking" + assert self.batch_indices is not None, ( + "Batch indices are not available for training order tracking" + ) epoch = int(state.epoch or 0) global_step = state.global_step diff --git a/bergson/query/faiss_index.py b/bergson/query/faiss_index.py index 58a1b65d..073c5c9b 100644 --- a/bergson/query/faiss_index.py +++ b/bergson/query/faiss_index.py @@ -242,9 +242,9 @@ def create_index( shard_sizes[-1] += remainder # Verify all gradients will be consumed - assert ( - sum(shard_sizes) == total_grads - ), f"Shard sizes {shard_sizes} don't sum to total_grads {total_grads}" + assert sum(shard_sizes) == total_grads, ( + f"Shard sizes {shard_sizes} don't sum to total_grads {total_grads}" + ) dl = gradients_loader(gradients_path) buffer: list[NDArray] = [] diff --git a/pyproject.toml b/pyproject.toml index c0f9ecf1..9b02a510 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,10 +21,14 @@ dependencies = [ version = "0.2.0" [project.optional-dependencies] dev = [ + "matplotlib", "pre-commit", "pytest", + "pytest-xdist", "pyright", + "setuptools", "trl", + "wandb", # Generate documentation "furo", "myst-parser", @@ -51,10 +55,12 @@ reportPrivateImportUsage = false include = ["bergson*"] [tool.ruff] -lint.ignore = ["E741"] # Ambiguous variable name +lint.ignore = ["E741", # Ambiguous variable name + "E501", # line-too-long (formatter takes care of it) +] # Enable pycodestyle (`E`), Pyflakes (`F`), and isort (`I`) codes # See https://beta.ruff.rs/docs/rules/ for more possible rules -lint.select = ["E", "F", "I"] +lint.extend-select = ["E", "F", "I"] # Same as Black. line-length = 88 diff --git a/tests/test_build.py b/tests/test_build.py index 6319499e..393c1c75 100644 --- a/tests/test_build.py +++ b/tests/test_build.py @@ -80,9 +80,9 @@ def test_split_attention_build(tmp_path: Path, model, dataset): attention_cfgs=attention_cfgs, ) - assert any( - Path(cfg.partial_run_path).iterdir() - ), "Expected artifacts in the temp run_path" + assert any(Path(cfg.partial_run_path).iterdir()), ( + "Expected artifacts in the temp run_path" + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @@ -107,9 +107,9 @@ def test_conv1d_build(tmp_path: Path, dataset): cfg=cfg, ) - assert any( - Path(cfg.partial_run_path).iterdir() - ), "Expected artifacts in the run path" + assert any(Path(cfg.partial_run_path).iterdir()), ( + "Expected artifacts in the run path" + ) index = load_gradients(cfg.partial_run_path) diff --git a/tests/test_reduce.py b/tests/test_reduce.py index 6718a298..76ec0d60 100644 --- a/tests/test_reduce.py +++ b/tests/test_reduce.py @@ -31,6 +31,6 @@ def test_reduce_e2e(tmp_path: Path): text=True, ) - assert ( - "error" not in result.stderr.lower() - ), f"Error found in stderr: {result.stderr}" + assert "error" not in result.stderr.lower(), ( + f"Error found in stderr: {result.stderr}" + ) diff --git a/tests/test_score.py b/tests/test_score.py index 294d6361..ee727304 100644 --- a/tests/test_score.py +++ b/tests/test_score.py @@ -64,9 +64,9 @@ def test_large_gradients_query(tmp_path: Path, dataset): ) assert result.returncode == 0 - assert ( - "error" not in result.stderr.lower() - ), f"Error found in stderr: {result.stderr}" + assert "error" not in result.stderr.lower(), ( + f"Error found in stderr: {result.stderr}" + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")