Merge remote-tracking branch 'origin/bench_structure' into bench_sparsity

Author: jainapurva

benchmarks/__init__.py

@@ -47,15 +47,15 @@ model_params:
         [4096, 4096, 1024]
       ]
   high_precision_dtype: "torch.bfloat16"
-  compile: true
-  compile_mode: "max-autotune"
+  compile: "max-autotune" # Options: "default", "max-autotune", "false"
   device: "cuda"  # Options: "cuda", "mps", "xpu", "cpu"
   model_type: "linear"  # Options: "linear", "ln_linear_sigmoid"
 ```
 
 ## Configuration Options
 
 ### Quantization Methods
+Currently, quantization string is in same format as the one being passed in llama/generate.py.
 - `baseline`: No quantization
 - `int8wo`: 8-bit weight-only quantization
 - `int4wo-{group_size}`: 4-bit weight-only quantization with specified group size

benchmarks/microbenchmarks/README.md

@@ -7,21 +7,21 @@
 
 from copy import deepcopy
 from pathlib import Path
-from typing import Dict
 
 import torch
 
-from utils import (
+from benchmarks.microbenchmarks.utils import (
     BenchmarkConfig,
-    benchmark_model_inference_in_microseconds,
+    BenchmarkResult,
     clean_caches,
     create_model_and_input,
-    quantization_string_to_quantization_config,
+    model_inference_time_in_ms,
+    string_to_config,
 )
 from torchao.quantization import quantize_
 
 
-def run(config: BenchmarkConfig) -> Dict[str, float]:
+def run(config: BenchmarkConfig) -> BenchmarkResult:
     """Run inference benchmarks"""
     clean_caches()  # Clean caches
 
@@ -39,25 +39,22 @@ def run(config: BenchmarkConfig) -> Dict[str, float]:
 
     # Use quantize_ to apply each quantization function to the model
     m_copy = deepcopy(base_model).eval().to(config.device)
-    quantization_config = quantization_string_to_quantization_config(
-        config.quantization,
-        config.sparsity,
-        high_precision_dtype=config.high_precision_dtype
+    quantization_config = string_to_config(
+        config.quantization, high_precision_dtype=config.high_precision_dtype
     )
-    if quantization_config:
+    if quantization_config is not None:
         quantize_(m_copy, quantization_config)
-    if config.compile:
+    if config.use_torch_compile:
         print("Compiling model....")
-        m_copy = torch.compile(m_copy, mode=config.compile_mode, fullgraph=True)
+        m_copy = torch.compile(m_copy, mode=config.torch_compile_mode, fullgraph=True)
 
     # Run benchmarks
-    result = {**config.to_dict()}
+    result = BenchmarkResult(config=config)
 
     # Benchmark time to run an inference call for quantized model
-    model_time = benchmark_model_inference_in_microseconds(
+    result.model_inference_time_in_ms = model_inference_time_in_ms(
         model=m_copy, input_data=input_data
     )
-    result["benchmark_model_inference_in_microseconds"] = model_time
 
     # TODO: Benchmark time using profiler
     # Profile dtype model evaluation

benchmarks/microbenchmarks/benchmark_inference.py

@@ -19,7 +19,7 @@
 
 import yaml
 
-from utils import (
+from benchmarks.microbenchmarks.utils import (
     BenchmarkConfig,
     generate_results_csv,
     print_results,
@@ -66,9 +66,9 @@ def load_benchmark_configs(config_path: str) -> List[BenchmarkConfig]:
     return configs
 
 
-def run_benchmarks_from_config(config_path: str) -> None:
+def run_inference_benchmarks_from_config(config_path: str) -> None:
     """Run benchmarks using configurations from YAML file"""
-    from benchmark_inference import run as run_inference
+    from benchmarks.microbenchmarks.benchmark_inference import run as run_inference
 
     configs = load_benchmark_configs(config_path)
     results = []
@@ -100,5 +100,24 @@ def run_benchmarks_from_config(config_path: str) -> None:
         required=True,
         help="Path to benchmark configuration file",
     )
+    parser.add_argument(
+        "--benchmark_mode",
+        "-m",
+        type=str,
+        default="inference",
+        choices=["inference", "training"],
+        help="Benchmark mode to run: inference or training",
+    )
     args = parser.parse_args()
-    run_benchmarks_from_config(args.config)
+
+    # Run benchmarks
+    if args.benchmark_mode == "inference":
+        run_inference_benchmarks_from_config(args.config)
+    elif args.benchmark_mode == "training":
+        print("Training mode not implemented yet")
+    else:
+        raise ValueError(
+            f"Invalid benchmark mode: {args.benchmark_mode}, choose from inference or training"
+        )
+
+    # TODO: Add support for args to override config values and run smaller benchmarks

benchmarks/microbenchmarks/benchmark_runner.py

@@ -13,7 +13,8 @@ model_params:
         [4096, 4096, 1024]
       ]
   high_precision_dtype: "torch.bfloat16"
-  compile: "max-autotune"
+  use_torch_compile: true
+  torch_compile_mode: "max-autotune"
   device: "cuda"  # Change this to "cuda", "mps", "xpu", or "cpu" as needed
   model_type: "linear"
   sparsity: "2:4"

benchmarks/microbenchmarks/test/benchmark_config.yml

@@ -8,7 +8,7 @@ class TestBenchmarkInference(unittest.TestCase):
     def setUp(self):
         self.params = {
             "high_precision_dtype": "torch.float32",  # Use float32 for testing
-            "compile": False,
+            "use_torch_compile": False,
             "device": "cpu",  # Use CPU for testing
             "model_type": "linear",
         }
@@ -23,13 +23,9 @@ def setUp(self):
     def test_run_inference(self):
         result = run(self.config)
 
-        # Check result contains all config attributes
-        for key in self.config.to_dict():
-            self.assertIn(key, result)
-
         # Check benchmark result is present and reasonable
-        self.assertIn("benchmark_model_inference_in_microseconds", result)
-        self.assertGreater(result["benchmark_model_inference_in_microseconds"], 0)
+        self.assertTrue(hasattr(result, "model_inference_time_in_ms"))
+        self.assertGreater(result.model_inference_time_in_ms, 0)
 
 
 if __name__ == "__main__":

benchmarks/microbenchmarks/test/test_benchmark_inference.py

@@ -7,7 +7,7 @@
 from benchmarks.microbenchmarks.benchmark_runner import (
     get_shapes_for_config,
     load_benchmark_configs,
-    run_benchmarks_from_config,
+    run_inference_benchmarks_from_config,
 )
 
 
@@ -24,7 +24,7 @@ def setUp(self):
                     }
                 ],
                 "high_precision_dtype": "torch.float32",
-                "compile": False,
+                "use_torch_compile": False,
                 "device": "cpu",
                 "model_type": "linear",
             },
@@ -72,8 +72,8 @@ def test_load_benchmark_configs(self):
         self.assertEqual(configs[0].quantization, "baseline")
         self.assertEqual(configs[1].quantization, "int8wo")
 
-    def test_run_benchmarks_from_config(self):
-        run_benchmarks_from_config(self.config_path)
+    def test_run_inference_benchmarks_from_config(self):
+        run_inference_benchmarks_from_config(self.config_path)
         results_file = os.path.join(self.config["output_dir"], "results.csv")
         self.assertTrue(os.path.exists(results_file))
 

benchmarks/microbenchmarks/test/test_benchmark_runner.py

@@ -6,6 +6,7 @@
 
 from benchmarks.microbenchmarks.utils import (
     BenchmarkConfig,
+    BenchmarkResult,
     LNLinearSigmoid,
     ToyLinearModel,
     clean_caches,
@@ -18,8 +19,8 @@ class TestUtils(unittest.TestCase):
     def test_benchmark_config(self):
         params = {
             "high_precision_dtype": "torch.bfloat16",
-            "compile": True,
-            "compile_mode": "max-autotune",
+            "use_torch_compile": True,
+            "torch_compile_mode": "max-autotune",
             "device": "cuda",
             "model_type": "linear",
         }
@@ -36,8 +37,8 @@ def test_benchmark_config(self):
         self.assertEqual(config.k, 1024)
         self.assertEqual(config.n, 1024)
         self.assertEqual(config.high_precision_dtype, torch.bfloat16)
-        self.assertEqual(config.compile, True)
-        self.assertEqual(config.compile_mode, "max-autotune")
+        self.assertEqual(config.use_torch_compile, True)
+        self.assertEqual(config.torch_compile_mode, "max-autotune")
         self.assertEqual(config.device, "cuda")
         self.assertEqual(config.model_type, "linear")
         self.assertEqual(config.output_dir, "test_output")
@@ -88,20 +89,24 @@ def test_create_model_and_input(self):
 
     def test_generate_results_csv(self):
         results = [
-            {
-                "quantization": "int8wo",
-                "m": 1024,
-                "k": 1024,
-                "n": 1024,
-                "time_us": 100.0,
-            },
-            {
-                "quantization": "int4wo",
-                "m": 1024,
-                "k": 1024,
-                "n": 1024,
-                "time_us": 50.0,
-            },
+            BenchmarkResult(
+                BenchmarkConfig(
+                    quantization="int8wo",
+                    params={},
+                    shape_name="custom",
+                    shape=[1024, 1024, 1024],
+                    output_dir="test_output",
+                ),
+            ),
+            BenchmarkResult(
+                BenchmarkConfig(
+                    quantization="int4wo",
+                    params={},
+                    shape_name="custom",
+                    shape=[1024, 1024, 1024],
+                    output_dir="test_output",
+                ),
+            ),
         ]
 
         with tempfile.TemporaryDirectory() as tmp_dir: