Adds get_quantization_layer_structure hooks for GPTQ (#2462)

* Adds get_quantization_layer_structure hooks

* cleanup

* format

Author: JyotinderSingh

keras_hub/src/models/causal_lm.py

@@ -429,3 +429,25 @@ def _post_quantize(self, mode, **kwargs):
         super()._post_quantize(mode, **kwargs)
         # Reset the compiled generate function.
         self.generate_function = None
+
+    def get_quantization_layer_structure(self, mode):
+        if mode != "gptq":
+            return None
+
+        backbone = self.backbone
+        # Check for standard backbone structure.
+        if not hasattr(backbone, "transformer_layers"):
+            return None
+
+        # Check for embedding.
+        embedding = getattr(backbone, "token_embedding", None)
+        if embedding is None:
+            embedding = getattr(backbone, "embedding", None)
+
+        if embedding is None:
+            return None
+
+        return {
+            "pre_block_layers": [embedding],
+            "sequential_blocks": backbone.transformer_layers,
+        }

keras_hub/src/models/gemma/gemma_causal_lm.py

@@ -431,3 +431,19 @@ def default_layer_intercept_fn(x, unused_i):
         )
         per_token_loss = per_token_loss_fn(target_ids, logits)
         return per_token_loss
+
+    def get_quantization_layer_structure(self, mode):
+        if mode != "gptq":
+            return None
+
+        # Wrap embedding + scaling
+        backbone = self.backbone
+        inputs = keras.Input(shape=(None,), dtype="int32")
+        x = backbone.token_embedding(inputs)
+        x = x * ops.cast(ops.sqrt(backbone.hidden_dim), x.dtype)
+        pre_processor = keras.Model(inputs=inputs, outputs=x)
+
+        return {
+            "pre_block_layers": [pre_processor],
+            "sequential_blocks": backbone.transformer_layers,
+        }

keras_hub/src/models/gemma/gemma_causal_lm_test.py

@@ -295,3 +295,17 @@ def layer_intercept_fn_for_testing(x, i):
         # Assert shapes for info exfiltrated into the parent context.
         self.assertEqual(ops.shape(embedded_prompts), expected_embedded_shape)
         self.assertEqual(ops.shape(scores), expected_score_shape)
+
+    def test_get_quantization_layer_structure(self):
+        causal_lm = GemmaCausalLM(**self.init_kwargs)
+        structure = causal_lm.get_quantization_layer_structure("gptq")
+        self.assertIsInstance(structure, dict)
+        self.assertIn("pre_block_layers", structure)
+        self.assertIn("sequential_blocks", structure)
+        self.assertLen(structure["pre_block_layers"], 1)
+        self.assertIsInstance(structure["pre_block_layers"][0], keras.Model)
+        self.assertEqual(
+            structure["sequential_blocks"], self.backbone.transformer_layers
+        )
+
+        self.assertIsNone(causal_lm.get_quantization_layer_structure("int8"))

keras_hub/src/models/gpt2/gpt2_causal_lm.py

@@ -420,3 +420,20 @@ def default_layer_intercept_fn(x, unused_i):
         )
         per_token_loss = per_token_loss_fn(target_ids, logits)
         return per_token_loss
+
+    def get_quantization_layer_structure(self, mode):
+        if mode != "gptq":
+            return None
+
+        backbone = self.backbone
+        token_ids = keras.Input(shape=(None,), dtype="int32")
+        tokens = backbone.token_embedding(token_ids)
+        positions = backbone.position_embedding(tokens)
+        x = backbone.embeddings_add((tokens, positions))
+        x = backbone.embeddings_dropout(x)
+        pre_processor = keras.Model(inputs=token_ids, outputs=x)
+
+        return {
+            "pre_block_layers": [pre_processor],
+            "sequential_blocks": backbone.transformer_layers,
+        }

keras_hub/src/models/gpt2/gpt2_causal_lm_test.py

@@ -1,5 +1,6 @@
 from unittest.mock import patch
 
+import keras
 import pytest
 from keras import ops
 
@@ -199,3 +200,17 @@ def layer_intercept_fn_for_testing(x, i):
         # Assert shapes for info exfiltrated into the parent context.
         self.assertEqual(ops.shape(embedded_prompts), expected_embedded_shape)
         self.assertEqual(ops.shape(scores), expected_score_shape)
+
+    def test_get_quantization_layer_structure(self):
+        causal_lm = GPT2CausalLM(**self.init_kwargs)
+        structure = causal_lm.get_quantization_layer_structure("gptq")
+        self.assertIsInstance(structure, dict)
+        self.assertIn("pre_block_layers", structure)
+        self.assertIn("sequential_blocks", structure)
+        self.assertLen(structure["pre_block_layers"], 1)
+        self.assertIsInstance(structure["pre_block_layers"][0], keras.Model)
+        self.assertEqual(
+            structure["sequential_blocks"], self.backbone.transformer_layers
+        )
+
+        self.assertIsNone(causal_lm.get_quantization_layer_structure("int8"))

keras_hub/src/models/masked_lm.py

@@ -84,3 +84,25 @@ def compile(
             weighted_metrics=weighted_metrics,
             **kwargs,
         )
+
+    def get_quantization_layer_structure(self, mode):
+        if mode != "gptq":
+            return None
+
+        backbone = self.backbone
+        # Check for standard backbone structure.
+        if not hasattr(backbone, "transformer_layers"):
+            return None
+
+        # Check for embedding.
+        embedding = getattr(backbone, "token_embedding", None)
+        if embedding is None:
+            embedding = getattr(backbone, "embedding", None)
+
+        if embedding is None:
+            return None
+
+        return {
+            "pre_block_layers": [embedding],
+            "sequential_blocks": backbone.transformer_layers,
+        }

keras_hub/src/models/mistral/mistral_causal_lm_test.py

@@ -199,3 +199,18 @@ def layer_intercept_fn_for_testing(x, i):
         # Assert shapes for info exfiltrated into the parent context.
         self.assertEqual(ops.shape(embedded_prompts), expected_embedded_shape)
         self.assertEqual(ops.shape(scores), expected_score_shape)
+
+    def test_get_quantization_layer_structure(self):
+        causal_lm = MistralCausalLM(**self.init_kwargs)
+        structure = causal_lm.get_quantization_layer_structure("gptq")
+        self.assertIsInstance(structure, dict)
+        self.assertIn("pre_block_layers", structure)
+        self.assertIn("sequential_blocks", structure)
+        self.assertEqual(
+            structure["pre_block_layers"], [self.backbone.token_embedding]
+        )
+        self.assertEqual(
+            structure["sequential_blocks"], self.backbone.transformer_layers
+        )
+
+        self.assertIsNone(causal_lm.get_quantization_layer_structure("int8"))

keras_hub/src/models/phi3/phi3_causal_lm_test.py

@@ -115,3 +115,19 @@ def test_all_presets(self):
                 preset=preset,
                 input_data=self.input_data,
             )
+
+    def test_get_quantization_layer_structure(self):
+        causal_lm = Phi3CausalLM(**self.init_kwargs)
+        structure = causal_lm.get_quantization_layer_structure("gptq")
+        self.assertIsInstance(structure, dict)
+        self.assertIn("pre_block_layers", structure)
+        self.assertIn("sequential_blocks", structure)
+        self.assertEqual(
+            structure["pre_block_layers"],
+            [self.backbone.token_embedding],
+        )
+        self.assertEqual(
+            structure["sequential_blocks"], self.backbone.transformer_layers
+        )
+
+        self.assertIsNone(causal_lm.get_quantization_layer_structure("int8"))