From 3ea0b9c25839d035de595c195f9fd63f460abef5 Mon Sep 17 00:00:00 2001 From: Baber Abbasi Date: Tue, 11 Nov 2025 17:31:52 +0000 Subject: [PATCH 1/6] add req --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index b45c4ed9..0953c36f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,11 +12,13 @@ license = {text = "MIT License"} dependencies = [ "accelerate", # For device_map in from_pretrained "datasets", + "matplotlib>=3.10.7", "natsort", "peft>=0.17.0", "simple-parsing", "torch", "transformers", + "wandb>=0.22.3", ] version = "0.1.1" [project.optional-dependencies] From e1bb1038dff20da8c6d9f29e714609d4d8fc9d41 Mon Sep 17 00:00:00 2001 From: Baber Date: Fri, 14 Nov 2025 16:44:30 +0500 Subject: [PATCH 2/6] ci: add unit tests workflow and restrict build to upstream repo --- .github/workflows/build.yml | 2 +- .github/workflows/unit_tests.yml | 54 ++++++++++++++++++++++++++++++++ pyproject.toml | 5 +-- 3 files changed, 58 insertions(+), 3 deletions(-) create mode 100644 .github/workflows/unit_tests.yml diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index cafcef4e..db5738b7 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -29,10 +29,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..b2840553 --- /dev/null +++ b/.github/workflows/unit_tests.yml @@ -0,0 +1,54 @@ +name: Unit Tests + +on: + push: + branches: + - 'main' + pull_request: + branches: + - 'main' + workflow_dispatch: + +jobs: + linter: + name: Linters + runs-on: ubuntu-latest + timeout-minutes: 5 + + steps: + - name: Pre-commit Checks + uses: actions/checkout@v5 + - name: Set up Python 3.10 + uses: actions/setup-python@v6 + with: + python-version: '3.10' + cache: pip + cache-dependency-path: pyproject.toml + - name: Pre-Commit + uses: pre-commit/action@v3.0.1 + + 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 --locked --dev + + - name: Run tests + run: uv run pytest --showlocals -s -vv -n=auto tests + + - name: Cleanup + run: uv cache prune --ci diff --git a/pyproject.toml b/pyproject.toml index d804354c..cc387197 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,21 +12,22 @@ license = {text = "MIT License"} dependencies = [ "accelerate", # For device_map in from_pretrained "datasets", - "matplotlib>=3.10.7", "natsort", "peft>=0.17.0", "simple-parsing", "torch", "transformers", - "wandb>=0.22.3", ] version = "0.2.0" [project.optional-dependencies] dev = [ + "matplotlib", "pre-commit", "pytest", + "pytest-xdist", "pyright", "trl", + "wandb" ] example = [ "trl", From 50d4d0d6ee06ae35f67e38a693937db1208495a6 Mon Sep 17 00:00:00 2001 From: Baber Date: Fri, 14 Nov 2025 21:02:20 +0500 Subject: [PATCH 3/6] nit --- .github/workflows/build.yml | 3 --- .github/workflows/unit_tests.yml | 9 ++++++--- pyproject.toml | 1 + 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index db5738b7..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 diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index b2840553..8db2d1a0 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -1,5 +1,6 @@ name: Unit Tests - +env: + UV_SYSTEM_PYTHON: 1 on: push: branches: @@ -26,6 +27,8 @@ jobs: cache-dependency-path: pyproject.toml - name: Pre-Commit uses: pre-commit/action@v3.0.1 + - name: Type Checking + uses: jakebailey/pyright-action@v1 testcpu: name: CPU Tests @@ -45,10 +48,10 @@ jobs: enable-cache: true - name: Install dependencies - run: uv sync --locked --dev + run: uv sync --extra dev - name: Run tests - run: uv run pytest --showlocals -s -vv -n=auto tests + run: uv run pytest tests --showlocals -s -vv -n=auto - name: Cleanup run: uv cache prune --ci diff --git a/pyproject.toml b/pyproject.toml index cc387197..bffe0d51 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ dev = [ "pytest", "pytest-xdist", "pyright", + "setuptools", "trl", "wandb" ] From 78fa64390f6b93ec6f8bfa27864c75524fb97f72 Mon Sep 17 00:00:00 2001 From: Baber Date: Mon, 17 Nov 2025 18:58:13 +0500 Subject: [PATCH 4/6] nit --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 5e5c4a47..101f1076 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,7 @@ dev = [ "pyright", "setuptools", "trl", - "wandb" + "wandb", # Generate documentation "furo", "myst-parser", From 45841c295ab288be98f31ff6a2715943e2d14fc3 Mon Sep 17 00:00:00 2001 From: Baber Date: Tue, 18 Nov 2025 04:40:10 +0500 Subject: [PATCH 5/6] use ruff-format --- .pre-commit-config.yaml | 24 +++++++++++------------- pyproject.toml | 6 ++++-- 2 files changed, 15 insertions(+), 15 deletions(-) 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/pyproject.toml b/pyproject.toml index 101f1076..9b02a510 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,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 From b2c1017673f429ca8e5c03158f8ab1f8a402a3ec Mon Sep 17 00:00:00 2001 From: Baber Date: Tue, 18 Nov 2025 04:59:28 +0500 Subject: [PATCH 6/6] fix typing --- .github/workflows/unit_tests.yml | 19 ++++++++++++------- bergson/__main__.py | 3 +-- bergson/gradients.py | 7 +++---- bergson/huggingface.py | 6 +++--- bergson/query/faiss_index.py | 6 +++--- tests/test_build.py | 12 ++++++------ tests/test_reduce.py | 6 +++--- tests/test_score.py | 6 +++--- 8 files changed, 34 insertions(+), 31 deletions(-) diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index 8db2d1a0..9f944241 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -17,18 +17,23 @@ jobs: timeout-minutes: 5 steps: - - name: Pre-commit Checks - uses: actions/checkout@v5 - - name: Set up Python 3.10 + - uses: actions/checkout@v5 + - name: "Set up Python" uses: actions/setup-python@v6 with: - python-version: '3.10' - cache: pip - cache-dependency-path: pyproject.toml + 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@v1 + uses: jakebailey/pyright-action@v2 + - name: Cleanup + run: uv cache prune --ci testcpu: name: CPU Tests 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/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")