Micro-benchmark inference

[jainapurva]

This is the first PR, in the benchmarking effort. It provides the outline to setup inference microbenchmarking for quant api's in torchao.

The different inputs like quantization techniques, matrix sizes, compile, sparsity etc, will be given as input to the python script. The options are re-defined in the scipt for quantization techniques, and a developer can add new quant technique. The script will generate a csv with performance numbers, and that'll be used to plot charts and as an input to dashboard . The script performs the following tasks:

• Take input as .yml
• Performs benchmarking for quantize_ APIs eval time for configurations
• Record all the config params and their respective time in csv file.
• Test cases

Future PRs will include more config options and process the generated results.

benchmark_config.yml

# Sample configuration for inference kernel benchmarks
quantization_config_recipe_names:
  - "baseline"  # String format is same as llama/generate.py
  - "int8wo"
  - "int4wo-128"
output_dir: "benchmarks/microbenchmarks/test/results"  # Directory for results and plots
model_params:
  matrix_shapes:
    - name: "custom"
      shapes: [
        [1024, 1024, 1024],  # [m, k, n]
        [2048, 4096, 1024],
        [4096, 4096, 1024]
      ]
  high_precision_dtype: "torch.bfloat16"
  use_torch_compile: true
  torch_compile_mode: "max-autotune"
  device: "cuda"  # Change this to "cuda", "mps", "xpu", or "cpu" as needed
  model_type: "linear"

Run command:

python benchmarks/microbenchmarks/benchmark_runner.py --config benchmarks/microbenchmarks/benchmark_config.yml

Output will be stored in

benchmarks/microbenchmarks/results/results.csv

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1759

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[HDCharles]

lgtm