[Test Fix] Add Quantization then finetune tests (#964)

~~Contingent on merge of
huggingface/transformers#34719
^ has been merged not yet released

SUMMARY:
Add test to 
* Given a model, oneshot quantize, then run ptq - training. 
Model must be run_compressed = False to run

Note:
* When running finetune on an already optimized (one-shotted) mode, the
model needs to be decompressed explicitly using
`CompressedTensorsConfig`. See
https://github.com/vllm-project/llm-compressor/pull/964/files#diff-e480ed475c0a5b2beb4052c1dd2aca671999634ace41a5ea017fdff1ce68be0bR130-R135
* Tests using x2 H100s passed

Also fix a bug where in log_sparsification, the layer name is not being
recognized so fails. Here nothting is being sparsified, so num params is
set to zero

TEST PLAN:
ran the test using transformers main
must pass
tests/llmcompressor/transformers/finetune/test_oneshot_then_finetune.py

---------

Co-authored-by: Dipika Sikka <[email protected]>

Author: George

src/llmcompressor/pytorch/utils/sparsification.py

@@ -105,15 +105,14 @@ def params_quantized(self) -> int:
         """
         :return: number of parameters across quantized layers
         """
-        return sum(
-            torch.numel(self.trainable_params[f"{name}.weight"])
-            + (
-                torch.numel(self.trainable_params[f"{name}.bias"])
-                if hasattr(layer, "bias") and layer.bias is not None
-                else 0
-            )
-            for (name, layer) in get_quantized_layers(self.module)
-        )
+        num_params = 0
+        for name, layer in get_quantized_layers(self.module):
+            if getattr(layer, "weight", None) is not None:
+                num_params += torch.numel(layer.weight)
+            if getattr(layer, "bias", None) is not None:
+                num_params += torch.numel(layer.bias)
+
+        return num_params
 
     @property
     def params_quantized_percent(self) -> float:

tests/llmcompressor/transformers/finetune/test_oneshot_then_finetune.py

@@ -1,28 +1,23 @@
-import os
 import shutil
 import unittest
 from pathlib import Path
 
 import pytest
+from transformers import AutoModelForCausalLM
+from transformers.utils.quantization_config import CompressedTensorsConfig
+
+from llmcompressor.core import create_session
+from llmcompressor.modifiers.quantization import QuantizationModifier
+from llmcompressor.transformers import oneshot, train
 
 
 @pytest.mark.unit
-@pytest.mark.skipif(
-    "CADENCE" in os.environ
-    and (os.environ["CADENCE"] == "weekly" or os.environ["CADENCE"] == "nightly"),
-    reason="Don't run for weekly and nightly tests as those use multi gpu "
-    "runners and this test fails when ngpu>1",
-)
 class TestOneshotThenFinetune(unittest.TestCase):
     def setUp(self):
         self.output = Path("./finetune_output")
+        self.quantization_config = CompressedTensorsConfig(run_compressed=False)
 
-    def test_oneshot_then_finetune(self):
-        from transformers import AutoModelForCausalLM
-
-        from llmcompressor.core import create_session
-        from llmcompressor.transformers import oneshot, train
-
+    def test_oneshot_sparsification_then_finetune(self):
         recipe_str = "tests/llmcompressor/transformers/obcq/recipes/test_tiny2.yaml"
         model = AutoModelForCausalLM.from_pretrained(
             "Xenova/llama2.c-stories15M", device_map="auto"
@@ -47,8 +42,12 @@ def test_oneshot_then_finetune(self):
         recipe_str = (
             "tests/llmcompressor/transformers/finetune/test_finetune_recipe.yaml"
         )
+
+        # Explictly decompress the model for training using quantization_config
         model = AutoModelForCausalLM.from_pretrained(
-            self.output / "oneshot_out", device_map="auto"
+            self.output / "oneshot_out",
+            device_map="auto",
+            quantization_config=self.quantization_config,
         )
         distill_teacher = AutoModelForCausalLM.from_pretrained(
             "Xenova/llama2.c-stories15M", device_map="auto"
@@ -73,7 +72,12 @@ def test_oneshot_then_finetune(self):
             )
 
         # test reloading checkpoint and final model
-        model = AutoModelForCausalLM.from_pretrained(output_dir, device_map="auto")
+        # verify checkpoint reloading and can carry out finetune
+        # with the saved model
+        # Explictly decompress the model for training using quantization_config
+        model = AutoModelForCausalLM.from_pretrained(
+            output_dir, device_map="auto", quantization_config=self.quantization_config
+        )
         with create_session():
             train(
                 model=model,
@@ -88,5 +92,71 @@ def test_oneshot_then_finetune(self):
                 resume_from_checkpoint=True,  # use last checkpoint
             )
 
+    def test_oneshot_quantization_then_finetune(self):
+        recipe = QuantizationModifier(
+            targets="Linear", scheme="FP8_DYNAMIC", ignore=["lm_head"]
+        )
+
+        model = AutoModelForCausalLM.from_pretrained(
+            "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+            device_map="auto",
+        )
+        dataset = "open_platypus"
+        concatenate_data = False
+        num_calibration_samples = 64
+        output_dir = self.output / "oneshot_out"
+        splits = {"calibration": "train[:10%]"}
+
+        with create_session():
+            oneshot(
+                model=model,
+                dataset=dataset,
+                output_dir=output_dir,
+                num_calibration_samples=num_calibration_samples,
+                recipe=recipe,
+                concatenate_data=concatenate_data,
+                splits=splits,
+            )
+
+        from transformers.utils.quantization_config import CompressedTensorsConfig
+
+        quantization_config = CompressedTensorsConfig(run_compressed=False)
+        model = AutoModelForCausalLM.from_pretrained(
+            output_dir,
+            device_map="auto",
+            quantization_config=quantization_config,
+        )
+        dataset = "open_platypus"
+        concatenate_data = False
+        output_dir = self.output / "finetune_out"
+        splits = {"calibration": "train[:10%]", "train": "train[:10%]"}
+
+        with create_session():
+            train(
+                model=model,
+                dataset=dataset,
+                output_dir=output_dir,
+                num_calibration_samples=num_calibration_samples,
+                recipe=recipe,
+                concatenate_data=concatenate_data,
+                splits=splits,
+            )
+
+        # test reloading checkpoint and final model
+        model = AutoModelForCausalLM.from_pretrained(
+            output_dir, device_map="auto", quantization_config=quantization_config
+        )
+        with create_session():
+            train(
+                model=model,
+                dataset=dataset,
+                output_dir=output_dir,
+                num_calibration_samples=num_calibration_samples,
+                recipe=recipe,
+                concatenate_data=concatenate_data,
+                splits=splits,
+                resume_from_checkpoint=True,  # use last checkpoint
+            )
+
     def tearDown(self):
         shutil.rmtree(self.output)