Adds store serialization tests for GPTQ (#21689)

* Add store serialization tests for GPTQ

* remove TODO comments

* kernels are now in packed form

Author: JyotinderSingh

keras/src/layers/core/dense_test.py

@@ -831,7 +831,6 @@ def test_int4_kernel_returns_unpacked_form(self):
     def test_legacy_load_own_variables(self):
         # In previous versions, `load_own_variables` accepted a store with
         # numeric keys.
-        # TODO(JyotinderSingh): add gptq_store test.
         float32_store = {
             "0": np.random.random((8, 16)).astype("float32"),
             "1": np.random.random((16,)).astype("float32"),
@@ -862,6 +861,18 @@ def test_legacy_load_own_variables(self):
             # outputs_grad_amax_history.
             "7": np.random.random((1024,)).astype("float32"),
         }
+        gptq_store = {
+            # bias
+            "0": np.random.random((16,)).astype("float32"),
+            # quantized_kernel
+            "1": np.random.randint(0, 16, size=(8, 8), dtype="uint8"),
+            # kernel_scale.
+            "2": np.random.random((16, 1)).astype("float32"),
+            # kernel_zero
+            "3": np.random.random((16, 1)).astype("uint8"),
+            # g_idx
+            "4": np.random.random((8,)).astype("float32"),
+        }
 
         # Test float32 layer.
         layer = layers.Dense(units=16)
@@ -899,6 +910,16 @@ def test_legacy_load_own_variables(self):
         self.assertAllClose(layer.outputs_grad_scale, float8_store["6"])
         self.assertAllClose(layer.outputs_grad_amax_history, float8_store["7"])
 
+        # Test gptq-quantized layer.
+        layer = layers.Dense(units=16, dtype="gptq/4/8_from_float32")
+        layer.build((None, 8))
+        layer.load_own_variables(gptq_store)
+        self.assertAllClose(layer.bias, gptq_store["0"])
+        self.assertAllClose(layer.quantized_kernel, gptq_store["1"])
+        self.assertAllClose(layer.kernel_scale, gptq_store["2"])
+        self.assertAllClose(layer.kernel_zero, gptq_store["3"])
+        self.assertAllClose(layer.g_idx, gptq_store["4"])
+
     def test_int4_gptq_kernel_returns_unpacked_form(self):
         """Test that the `kernel` property returns the unpacked int4 GPTQ
         kernel."""

keras/src/layers/core/einsum_dense_test.py

@@ -1043,7 +1043,6 @@ def test_int4_kernel_returns_unpacked_form(self):
     def test_legacy_load_own_variables(self):
         # In previous versions, `load_own_variables` accepted a store with
         # numeric keys.
-        # TODO(JyotinderSingh): add gptq_store test.
         float32_store = {
             "0": np.random.random((3, 8, 32)).astype("float32"),
             "1": np.random.random((32,)).astype("float32"),
@@ -1074,6 +1073,18 @@ def test_legacy_load_own_variables(self):
             # outputs_grad_amax_history.
             "7": np.random.random((1024,)).astype("float32"),
         }
+        gptq_store = {
+            # bias
+            "0": np.random.random((32,)).astype("float32"),
+            # quantized_kernel
+            "1": np.random.randint(0, 16, size=(16, 24), dtype="uint8"),
+            # kernel_scale.
+            "2": np.random.random((32, 3)).astype("float32"),
+            # kernel_zero
+            "3": np.random.random((32, 3)).astype("uint8"),
+            # g_idx
+            "4": np.random.random((24,)).astype("float32"),
+        }
         config = dict(
             equation="ab,bcd->acd",
             output_shape=(8, 32),
@@ -1116,6 +1127,16 @@ def test_legacy_load_own_variables(self):
         self.assertAllClose(layer.outputs_grad_scale, float8_store["6"])
         self.assertAllClose(layer.outputs_grad_amax_history, float8_store["7"])
 
+        # Test gptq-quantized layer.
+        layer = layers.EinsumDense(**config, dtype="gptq/4/8_from_float32")
+        layer.build((None, 3))
+        layer.load_own_variables(gptq_store)
+        self.assertAllClose(layer.bias, gptq_store["0"])
+        self.assertAllClose(layer.quantized_kernel, gptq_store["1"])
+        self.assertAllClose(layer.kernel_scale, gptq_store["2"])
+        self.assertAllClose(layer.kernel_zero, gptq_store["3"])
+        self.assertAllClose(layer.g_idx, gptq_store["4"])
+
     def test_int4_gptq_kernel_returns_unpacked_form(self):
         """Test that the `kernel` property returns the unpacked int4 GPTQ
         kernel."""