Review updates

Author: jainapurva

benchmarks/microbenchmarks/benchmark_inference.py

@@ -45,7 +45,7 @@ def run(config: BenchmarkConfig) -> BenchmarkResult:
 
     # Use quantize_ to apply each quantization function to the model
     m_copy = deepcopy(base_model).eval().to(config.device)
-    aoBaseConfig = string_to_config(
+    ao_base_config = string_to_config(
         config.quantization,
         config.sparsity,
         high_precision_dtype=config.high_precision_dtype,
@@ -59,7 +59,7 @@ def run(config: BenchmarkConfig) -> BenchmarkResult:
     ):
         if is_cuda:
             print(f"Applying {config.sparsity} sparsity to model")
-            sparsify_(m_copy, aoBaseConfig)
+            sparsify_(m_copy, ao_base_config)
         else:
             print(
                 f"Warning: Skipping {config.sparsity} sparsity as it requires CUDA, but device is {config.device}"
@@ -70,7 +70,7 @@ def run(config: BenchmarkConfig) -> BenchmarkResult:
         pass  # No quantization or sparsity specified, do nothing
     else:
         print("Quantizing model....")
-        quantize_(m_copy, aoBaseConfig)
+        quantize_(m_copy, ao_base_config)
 
     if config.use_torch_compile:
         print("Compiling model....")

benchmarks/microbenchmarks/benchmark_runner.py

@@ -84,14 +84,6 @@ def get_quantization_sparsity_recipes(
     """
     config_recipes = set()
 
-    # Handle edge cases
-    if sparsity_recipes is None and quantization_recipes is None:
-        return {("baseline", None)}
-    if sparsity_recipes is None:
-        return {(quant, None) for quant in quantization_recipes}
-    if quantization_recipes is None:
-        return {("baseline", sparse) for sparse in sparsity_recipes}
-
     # Always include baseline without sparsity
     config_recipes.add(("baseline", None))
 
@@ -134,8 +126,8 @@ def load_benchmark_configs(cli_args: argparse.Namespace) -> List[BenchmarkConfig
     # Create all possible combinations
     configs = []
     quantization_sparsity_recipes = get_quantization_sparsity_recipes(
-        config.get("quantization_config_recipe_names", None),
-        config.get("sparsity_config_recipe_names", None),
+        config.get("quantization_config_recipe_names", []),
+        config.get("sparsity_config_recipe_names", []),
     )
     for model_param in config["model_params"]:
         shapes, params = get_param_combinations(model_param)

benchmarks/microbenchmarks/test/benchmark_config.yml

@@ -1,11 +1,11 @@
 # Sample configuration for inference benchmarks
 benchmark_mode: "inference"
 quantization_config_recipe_names:
-  # - "baseline" Will always run a baseline instatance
+  # Will run a baseline inference for model by default, without quantization for comparison
   - "int4wo-32"
   - "marlin"
 sparsity_config_recipe_names:
-  # - "none" Will always run a without sparsity instance
+  # Will run a baseline inference for model by default, without sparsity for comparison
   - "semi-sparse"
   - "block"
 output_dir: "benchmarks/microbenchmarks/results"

benchmarks/microbenchmarks/test/test_benchmark_inference.py

@@ -49,16 +49,16 @@ def test_run_inference(self, mock_string_to_config):
         self.assertTrue(hasattr(result, "model_inference_time_in_ms"))
 
     @patch("benchmarks.microbenchmarks.benchmark_inference.string_to_config")
-    def test_run_inference_with_sparsity(self, mock_string_to_config):
+    def test_run_inference_with_semi_sparse_marlin(self, mock_string_to_config):
         """Test running inference with sparsity configurations"""
         # Mock string_to_config to return valid configs
+        from torchao.dtypes import MarlinSparseLayout
         from torchao.quantization import Int4WeightOnlyConfig
-        from torchao.sparsity.sparse_api import (
-            BlockSparseWeightConfig,
-        )
 
         # Test with semi-sparse config
-        mock_string_to_config.return_value = Int4WeightOnlyConfig()
+        mock_string_to_config.return_value = Int4WeightOnlyConfig(
+            layout=MarlinSparseLayout()
+        )
         config = BenchmarkConfig(
             quantization="marlin",
             sparsity="semi-sparse",
@@ -77,6 +77,14 @@ def test_run_inference_with_sparsity(self, mock_string_to_config):
         self.assertIsInstance(result, BenchmarkResult)
         self.assertTrue(hasattr(result, "model_inference_time_in_ms"))
 
+    @patch("benchmarks.microbenchmarks.benchmark_inference.string_to_config")
+    def test_run_inference_with_block_sparsity(self, mock_string_to_config):
+        """Test running inference with sparsity configurations"""
+        # Mock string_to_config to return valid configs
+        from torchao.sparsity.sparse_api import (
+            BlockSparseWeightConfig,
+        )
+
         # Test with block sparsity
         mock_string_to_config.return_value = BlockSparseWeightConfig()
         config = BenchmarkConfig(