[Examples] [Bugfix] skip sparsity stats when saving checkpoints (#1528)

## Purpose ##
* Fix example failure where sparsity stats were attempted to be computed
on a compressed model
* This happened because the `get_model_compressor` cannot compute
sparsity stats of a compressed model

## Prerequisites ##
* neuralmagic/compressed-tensors#345

## Changes ##
* When saving checkpoints, call `get_model_compressor` on the
decompressed model (before it is compressed, not after)
* Do not save checkpoints with sparsity stats by default, due to high
runtime

## Testing ##
* Ran example to completion with reduced training duration

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/llmcompressor/pytorch/model_load/helpers.py

@@ -45,6 +45,16 @@ def save_checkpoint(
         get_model_compressor,  # avoid circular import
     )
 
+    # used for decompression
+    # unfortunately, if skip_sparsity_compression_stats==True, sparsity stats
+    # are computed twice. In the future, track sparsity from recipe or
+    # share recipe between compression and decompression
+    compressor = get_model_compressor(
+        model=model,
+        save_compressed=save_compressed,
+        skip_sparsity_compression_stats=skip_sparsity_compression_stats,
+    )
+
     # saving the model also saves the recipe
     model.save_pretrained(
         save_path,
@@ -55,13 +65,8 @@ def save_checkpoint(
     if processor is not None:
         processor.save_pretrained(save_path)
 
-    # saving the model modifies the model strcuture
+    # decompression: saving the model modifies the model strcuture
     # as this is only a checkpoint, decompress model to enable future training/oneshot
-    compressor = get_model_compressor(
-        model=model,
-        save_compressed=save_compressed,
-        skip_sparsity_compression_stats=skip_sparsity_compression_stats,
-    )
     if compressor is not None:
         compressor.decompress_model(model)
 

src/llmcompressor/transformers/finetune/session_mixin.py

@@ -363,7 +363,7 @@ def save_model(
         self,
         output_dir: str,
         _internal_call: bool = False,
-        skip_sparsity_compression_stats: Optional[bool] = False,
+        skip_sparsity_compression_stats: Optional[bool] = True,
     ):
         """
         Override of the save_model function and expects it to exist in the parent.
@@ -388,6 +388,8 @@ def save_model(
             self.model.prepare_for_save()  # TODO: move to finalize
 
         # save checkpoint
+        # note that skip_sparsity_compression_stats
+        # is True by default to avoid high runtime cost
         self.save_state()
         if self.accelerator.is_main_process:
             processor = getattr(self, "processing_class", self.tokenizer)