Skip to content

Commit

Permalink
Small fix for micro benchmark code (#711)
Browse files Browse the repository at this point in the history
Summary:
There seems to be some problems running benchmark_aq.py:
```
torch._dynamo.exc.Unsupported: torch.* op returned non-Tensor int call_function <method 'size' of 'torch._C.TensorBase' objects>
```
when we run the benchmark with multiple shapes sometimes. But the problem will be gone if we reset the dynamo caches before
each benchmark run

Test Plan:
python benchmarks/benchmark_aq.py

Reviewers:

Subscribers:

Tasks:

Tags:
  • Loading branch information
jerryzh168 authored Aug 22, 2024
1 parent b09307a commit 9860194
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion benchmarks/benchmark_aq.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
_replace_with_custom_fn_if_matches_filter,
)
import copy
from torchao.utils import unwrap_tensor_subclass

def _int8wo_api(mod, **kwargs):
if TORCH_VERSION_AT_LEAST_2_4:
Expand Down Expand Up @@ -133,15 +134,17 @@ def _bench_quantized_tensor_subclass_perf(api, ref_api, M, N, K, kwargs=None):
WARMUP = 20
RUNS = 100

torch._dynamo.reset()
m_ref = torch.compile(m_ref, mode='max-autotune', fullgraph=True)
benchmark_model(m_ref, WARMUP, example_inputs)
ref_elapsed_time = benchmark_model(m_ref, RUNS, example_inputs)

torch._dynamo.reset()
m = torch.compile(m, mode='max-autotune', fullgraph=True)
benchmark_model(m, WARMUP, example_inputs)
elapsed_time = benchmark_model(m, RUNS, example_inputs)


torch._dynamo.reset()
m_bf16 = torch.compile(m_bf16, mode='max-autotune', fullgraph=True)
benchmark_model(m_bf16, WARMUP, example_inputs)
bf16_elapsed_time = benchmark_model(m_bf16, RUNS, example_inputs)
Expand Down

0 comments on commit 9860194

Please sign in to comment.