Add skip_sparsity_compression_stats option

Author: anmarques

src/automation/tasks/llmcompressor.py

@@ -32,7 +32,8 @@ def __init__(
         vision_samples: Optional[int]=None,
         max_seq_len: int=8192,
         trust_remote_code: bool=False,
-        max_memory_per_gpu: str="hessian",
+        max_memory_per_gpu: str="auto",
+        skip_sparsity_compression_stats=True,
         tracing_class: Optional[str]=None,
         tags: Union[str, List[str]]=None,
         task_type: str="training",
@@ -94,6 +95,7 @@ def __init__(
         self.model_class = model_class
         self.dataset_loader = dataset_loader
         self.data_collator = data_collator
+        self.skip_sparsity_compression_stats = skip_sparsity_compression_stats
 
         if tags is not None:
             tags = list(set(config_kwargs.pop("tags", []).extend(tags)))
@@ -140,6 +142,7 @@ def get_arguments(self):
                 "trust_remote_code": self.trust_remote_code,
                 "max_memory_per_gpu": self.max_memory_per_gpu,
                 "tracing_class": self.tracing_class,
+                "skip_sparsity_compression_stats": self.skip_sparsity_compression_stats,
                 "tags": self.tags,
             },
         }

src/automation/tasks/scripts/llmcompressor_script.py

@@ -27,6 +27,7 @@ def llmcompressor_main(
     max_seq_len,
     text_samples,
     vision_samples,
+    skip_sparsity_compression_stats,
     save_directory,
     data_collator,
 ):
@@ -129,7 +130,7 @@ def llmcompressor_main(
     )
 
     # Save model compressed
-    model.save_pretrained(save_directory, save_compressed=True)
+    model.save_pretrained(save_directory, save_compressed=True, skip_sparsity_compression_stats=skip_sparsity_compression_stats)
     processor.save_pretrained(save_directory)
 
     return recipe
@@ -155,6 +156,7 @@ def main(configurations=None):
     recipe = args.get("recipe", None)
     recipe_args = args.get("recipe_args", None)
     tags = args.get("tags", None)
+    skip_sparsity_compression_stats = parse_argument(args["skip_sparsity_compression_stats"], bool)
 
     dataset_loader_fn = load_callable_configuration("dataset loader", configurations)
     data_collator_fn = load_callable_configuration("data collator", configurations)
@@ -175,6 +177,7 @@ def main(configurations=None):
         max_seq_len,
         text_samples,
         vision_samples,
+        skip_sparsity_compression_stats,
         save_directory,
         data_collator_fn,
     )