[Bugfix] Prepare KD Models when Saving

src/llmcompressor/modifiers/distillation/utils/pytorch/kd_wrapper.py

@@ -24,6 +24,7 @@ def __init__(
         super(KDModuleWrapper, self).__init__()
 
         self.layer = layer
+        self._save_active = False
         self._fsdp_active = fsdp_active
         self.offload_output = offload_output
         self.kd_transforms = transforms
@@ -88,16 +89,28 @@ def named_modules(
         prefix: str = "",
         remove_duplicate: bool = True,
     ):
-        # we want the full names of modules in two cases
+        # outside of saving, we want the full names of modules in two cases:
         # 1. trainer initialization, so teacher is moved to the correct device. This is
         # caught by the kd_enabled flag, which is set when the modifier is started
         # 2. running in DataParallel (non-FSDP) mode so the replicate function can pick
         # up the teacher.
-        if not self.kd_enabled or not self._fsdp_active:
-            return super().named_modules(
+        if self._save_active or (self.kd_enabled and self._fsdp_active):
+            return self.layer.named_modules(
                 memo=memo, prefix=prefix, remove_duplicate=remove_duplicate
             )
 
-        return self.layer.named_modules(
+        return super().named_modules(
             memo=memo, prefix=prefix, remove_duplicate=remove_duplicate
         )
+
+    def prepare_for_save(self):
+        """
+        Prepare model structure to be saved, specifically `self.named_modules`
+        """
+        self._save_active = True
+
+    def finish_save(self):
+        """
+        Finish saving model
+        """
+        self._save_active = False

src/llmcompressor/modifiers/distillation/utils/pytorch/model_wrapper.py

@@ -23,6 +23,7 @@ def __init__(
         self.teacher_model = teacher_model
         self.wrappers = wrappers
         self.kd_comparison = comparison
+        self._save_active = False
         self._fsdp_active = fsdp_active
         self.kd_enabled = False
         self.register_buffer(self.KD_LAST_COMPARISON, torch.zeros(1, device="cpu"))
@@ -88,17 +89,17 @@ def named_modules(
         prefix: str = "",
         remove_duplicate: bool = True,
     ):
-        # we want the full names of modules in two cases
+        # outside of saving, we want the full names of modules in two cases:
         # 1. trainer initialization, so teacher is moved to the correct device. This is
         # caught by the kd_enabled flag, which is set when the modifier is started
         # 2. running in DataParallel (non-FSDP) mode so the replicate function can pick
         # up the teacher.
-        if not self.kd_enabled or not self._fsdp_active:
-            return super().named_modules(
+        if self._save_active or (self.kd_enabled and self._fsdp_active):
+            return self.student_model.named_modules(
                 memo=memo, prefix=prefix, remove_duplicate=remove_duplicate
             )
 
-        return self.student_model.named_modules(
+        return super().named_modules(
             memo=memo, prefix=prefix, remove_duplicate=remove_duplicate
         )
 
@@ -109,6 +110,24 @@ def train(self, mode: bool = True):
         self.student_model.train(mode)
         return self
 
+    def prepare_for_save(self):
+        """
+        Prepare model structure to be saved, specifically `self.named_modules`
+        """
+        self._save_active = True
+        for student_wrapper, teacher_wrapper in self.wrappers.values():
+            student_wrapper.prepare_for_save()
+            teacher_wrapper.prepare_for_save()
+
+    def finish_save(self):
+        """
+        Finish saving model
+        """
+        self._save_active = False
+        for student_wrapper, teacher_wrapper in self.wrappers.values():
+            student_wrapper.finish_save()
+            teacher_wrapper.finish_save()
+
     def __getattr__(self, name: str) -> Any:
         try:
             return super().__getattr__(name)

src/llmcompressor/transformers/finetune/session_mixin.py

@@ -21,6 +21,9 @@
     pre_initialize_structure,
 )
 from llmcompressor.metrics import LoggerManager
+from llmcompressor.modifiers.distillation.utils.pytorch.model_wrapper import (
+    KDModelWrapper,
+)
 from llmcompressor.pytorch.model_load.helpers import RECIPE_FILE_NAME, get_session_model
 from llmcompressor.pytorch.utils import ModuleSparsificationInfo
 from llmcompressor.transformers.finetune.callbacks import (
@@ -341,13 +344,25 @@ def train(self, *args, stage: Optional[str] = None, **kwargs):
         :param kwargs: keyword args to pass to super().train()
         :return: the output from super.train()
         """
+
+        # lifecycle
         checkpoint, epoch = self._calculate_checkpoint_info(kwargs)
         self.initialize_session(epoch=epoch, checkpoint=checkpoint, stage=stage)
+
+        # do not save checkpoints as compressed
+        original_save_compressed = self.args.save_compressed
+        self.args.save_compressed = False
+
+        # train with accelerator
         self.accelerator.wait_for_everyone()
         output = super().train(*args, **kwargs)
         self.accelerator.wait_for_everyone()
-        self.finalize_session()
 
+        # restore original setting for saving final model
+        self.args.save_compressed = original_save_compressed
+
+        # lifecycle
+        self.finalize_session()
         self.accelerator.wait_for_everyone()
 
         # log model sparsity
@@ -430,6 +445,10 @@ def save_model(
         if output_dir is None:
             output_dir = self.args.output_dir
 
+        # knowledge distillation requires making wrappers transparent during
+        if isinstance(self.model, KDModelWrapper):
+            self.model.prepare_for_save()
+
         if not is_fsdp_model(self.model):
             self.model.save_pretrained(
                 output_dir,
@@ -467,6 +486,9 @@ def save_model(
 
         self.accelerator.wait_for_everyone()
 
+        if isinstance(self.model, KDModelWrapper):
+            self.model.finish_save()
+
     def maybe_log_model_sparsification(self):
         """
         Log info on model sparsity and quantization if possible. Only print logs on the

tests/llmcompressor/transformers/finetune/test_oneshot_then_finetune.py

@@ -62,7 +62,7 @@ def test_oneshot_then_finetune(self):
         concatenate_data = False
         output_dir = self.output / "finetune_out"
         splits = "train[:50%]"
-        max_steps = 50
+        max_steps = 25
 
         with create_session():
             train(
@@ -77,5 +77,23 @@ def test_oneshot_then_finetune(self):
                 max_steps=max_steps,
             )
 
+        # test reloading checkpoint and final model
+        model = SparseAutoModelForCausalLM.from_pretrained(
+            output_dir, device_map="auto"
+        )
+        with create_session():
+            train(
+                model=model,
+                distill_teacher=distill_teacher,
+                dataset=dataset,
+                output_dir=output_dir,
+                num_calibration_samples=num_calibration_samples,
+                recipe=recipe_str,
+                concatenate_data=concatenate_data,
+                splits=splits,
+                max_steps=max_steps,
+                resume_from_checkpoint=True,  # use last checkpoint
+            )
+
     def tearDown(self):
         shutil.rmtree(self.output)