Setup Github CI for tests

Author: MohsenDehghankar

.github/workflows/tests.yml

@@ -0,0 +1,39 @@
+name: Tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  tests:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+
+      - name: Install system build tools
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y build-essential
+
+      - name: Install Python dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements.txt
+
+      - name: Rebuild CPU kernel libraries
+        run: |
+          make -C kernels/bit_1/cpu clean all
+          make -C kernels/bit_1_58/cpu clean all
+
+      - name: Run test suite
+        run: pytest -q tests

pyproject.toml

@@ -0,0 +1,26 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "rsr-core"
+version = "0.1.0"
+description = "Core RSR kernels, multipliers, integrations, and benchmarks."
+readme = "README.md"
+requires-python = ">=3.10"
+dynamic = ["dependencies"]
+
+[project.urls]
+Repository = "https://github.com/UIC-InDeXLab/RSR-Core"
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.dynamic]
+dependencies = { file = ["requirements.txt"] }
+
+[tool.setuptools.packages.find]
+include = ["benchmarking*", "integrations*", "multiplier*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

requirements.txt

@@ -1 +1,7 @@
-torch
+matplotlib
+numpy
+pytest
+safetensors
+torch
+tqdm
+transformers

tests/test_model_infer.py

@@ -21,6 +21,7 @@
     preprocess_layer_cuda,
     save_preprocessed,
 )
+from multiplier.bit_1_58.cpu.rsr_runtime import select_cpu_tensor_keys
 
 
 # ---------------------------------------------------------------------------
@@ -316,7 +317,7 @@ def test_roundtrip(self):
 
             with safe_open(out / "rsr_weights.safetensors", framework="pt") as handle:
                 tensors = {}
-                for key in ["perms", "group_ends", "pos_masks", "neg_masks", "block_meta"]:
+                for key in select_cpu_tensor_keys(n_cols, k):
                     tensors[key] = handle.get_tensor(f"test_layer.{key}")
 
             layer = RSRLinear("test_layer", layer_meta["test_layer"], tensors)

tests/test_model_prep.py

@@ -20,6 +20,7 @@
     save_preprocessed,
     unpack_ternary_weights,
 )
+from multiplier.bit_1_58.cpu.rsr_runtime import select_cpu_tensor_keys
 
 
 # ---------------------------------------------------------------------------
@@ -207,15 +208,17 @@ def test_excludes_packed_weight_only(self):
 # ---------------------------------------------------------------------------
 
 class TestPreprocessLayerCpu:
-    EXPECTED_KEYS = {"perms", "group_ends", "pos_masks", "neg_masks", "block_meta"}
+    @staticmethod
+    def _expected_keys(n_cols, k):
+        return set(select_cpu_tensor_keys(n_cols, k))
 
     @pytest.mark.parametrize(
         "n_rows,n_cols", [(16, 32), (32, 16), (64, 64), (24, 48)]
     )
     def test_output_keys(self, n_rows, n_cols):
         w = torch.randint(-1, 2, (n_rows, n_cols), dtype=torch.int8)
         result = preprocess_layer_cpu(w, k=8)
-        assert set(result.keys()) == self.EXPECTED_KEYS
+        assert set(result.keys()) == self._expected_keys(n_cols, 8)
 
     def test_all_tensors_int32(self):
         w = torch.randint(-1, 2, (32, 32), dtype=torch.int8)
@@ -234,12 +237,12 @@ def test_identity_matrix(self):
         n = 32
         w = torch.eye(n, dtype=torch.int8)
         result = preprocess_layer_cpu(w, k=8)
-        assert set(result.keys()) == self.EXPECTED_KEYS
+        assert set(result.keys()) == self._expected_keys(n, 8)
 
     def test_all_zeros(self):
         w = torch.zeros(16, 32, dtype=torch.int8)
         result = preprocess_layer_cpu(w, k=8)
-        assert set(result.keys()) == self.EXPECTED_KEYS
+        assert set(result.keys()) == self._expected_keys(32, 8)
 
 
 # ---------------------------------------------------------------------------
@@ -291,17 +294,16 @@ def test_no_column_padding(self):
 
 class TestSavePreprocessed:
     def _make_dummy(self):
+        layer_meta = {"layer1": {"n_rows": 32, "n_cols": 64, "k": 8}}
         rsr_data = {
-            "layer1": {
-                "perms": torch.zeros(10, dtype=torch.int32),
-                "group_ends": torch.zeros(5, dtype=torch.int32),
-                "pos_masks": torch.zeros(5, dtype=torch.int32),
-                "neg_masks": torch.zeros(5, dtype=torch.int32),
-                "block_meta": torch.zeros(4, dtype=torch.int32),
-            }
+            "layer1": {"perms": torch.zeros(10, dtype=torch.int32),
+                       "group_ends": torch.zeros(5, dtype=torch.int32),
+                       "block_meta": torch.zeros(4, dtype=torch.int32)}
         }
+        for key in select_cpu_tensor_keys(64, 8):
+            if key not in rsr_data["layer1"]:
+                rsr_data["layer1"][key] = torch.zeros(5, dtype=torch.int32)
         non_quantized = {"embed.weight": torch.randn(100, 64)}
-        layer_meta = {"layer1": {"n_rows": 32, "n_cols": 64, "k": 8}}
         return rsr_data, non_quantized, layer_meta
 
     def test_creates_output_files(self):
@@ -345,7 +347,7 @@ def test_safetensors_roundtrip(self):
                 k=8, version="adaptive", model_name="test",
             )
             loaded = load_file(str(out / "rsr_weights.safetensors"))
-            for key in ["perms", "group_ends", "pos_masks", "neg_masks", "block_meta"]:
+            for key in rsr_data["layer1"]:
                 full_key = f"layer1.{key}"
                 assert full_key in loaded
                 torch.testing.assert_close(loaded[full_key], rsr_data["layer1"][key])
@@ -403,7 +405,7 @@ def test_unpack_preprocess_save_roundtrip(self):
         torch.testing.assert_close(unpacked, weights)
 
         result = preprocess_layer_cpu(unpacked, k=k)
-        assert len(result) == 5
+        assert len(result) == len(select_cpu_tensor_keys(n_cols, k))
 
         with tempfile.TemporaryDirectory() as tmpdir:
             out = Path(tmpdir) / "model"