fix consecutive oneshot

src/llmcompressor/pytorch/model_load/helpers.py

@@ -213,7 +213,7 @@ def load_safetensors_state_dict(file_path: str) -> Dict[str, torch.Tensor]:
         return {key: f.get_tensor(key) for key in f.keys()}
 
 
-def copy_python_files_from_model_cache(model, save_path: str):
+def copy_python_files_from_model_cache(model: Module, save_path: str):
     config = model.config
     cache_path = None
     if hasattr(config, "_name_or_path"):

tests/llmcompressor/transformers/obcq/test_consecutive_runs.py

@@ -21,9 +21,13 @@ def _test_consecutive_runs(
         import math
 
         from llmcompressor.core import active_session
-        from llmcompressor.pytorch.model_load.helpers import get_session_model
+        from llmcompressor.pytorch.model_load.helpers import (
+            get_session_model,
+            initialize_recipe,
+        )
         from llmcompressor.pytorch.utils.helpers import tensor_sparsity
         from llmcompressor.transformers import oneshot
+        from llmcompressor.transformers.utils.helpers import resolve_recipe
         from llmcompressor.utils.pytorch import qat_active
 
         # test recipe with 50% sparsity, quantization and smoothquant
@@ -49,9 +53,13 @@ def _test_consecutive_runs(
         self.assertEqual(len(stages), 1)
         session.reset()
 
+        recipe = resolve_recipe(recipe=self.first_recipe, model_path=self.output_first)
+        if recipe:
+            initialize_recipe(model=first_tiny_model, recipe_path=recipe)
+
         # reload saved model and up sparsity to 0.7
         oneshot(
-            model=self.output_first,
+            model=first_tiny_model,
             dataset=self.dataset,
             num_calibration_samples=num_calibration_samples,
             recipe=self.second_recipe,
@@ -64,12 +72,14 @@ def _test_consecutive_runs(
         layer_0_sparse = tensor_sparsity(
             second_tiny_model.model.layers[0].self_attn.k_proj.weight
         )
+
         assert math.isclose(layer_0_sparse.item(), 0.7, rel_tol=tolerance)
         assert qat_active(second_tiny_model)
 
         session = active_session()
         session_recipe = session.lifecycle.recipe_container.compiled_recipe
         stages = [stage.group for stage in session_recipe.stages]
+
         self.assertEqual(len(stages), 2)
 
         recipe_path = self.output_second / "recipe.yaml"