Skip to content

Commit

Permalink
[BC][internal-first] Cleanup BC fixes (#641)
Browse files Browse the repository at this point in the history
Update autoquant.py
  • Loading branch information
nmacchioni authored Aug 8, 2024
1 parent 5f35645 commit e11201a
Showing 1 changed file with 7 additions and 14 deletions.
21 changes: 7 additions & 14 deletions torchao/quantization/autoquant.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,15 +20,6 @@

import torch.nn.functional as F

try:
from torch._inductor.utils import do_bench
except ImportError:
try:
from torch._inductor.runtime.runtime_utils import do_bench
except ImportError:
from torch._inductor.runtime.benchmarking import benchmarker
do_bench = benchmarker.benchmark

__all__ = [
"AutoQuantizableLinearWeight",
"autoquant",
Expand Down Expand Up @@ -232,14 +223,16 @@ def do_autoquant_bench(op, *args, **kwargs):
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph, stream=stream):
op(*args, **kwargs)
if TORCH_VERSION_AFTER_2_3 and not TORCH_VERSION_AFTER_2_5:
from torch._inductor.runtime.runtime_utils import do_bench_gpu
res = do_bench_gpu(lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median")
elif TORCH_VERSION_AFTER_2_5 and torch.cuda.is_available():
if TORCH_VERSION_AFTER_2_5:
from torch._inductor.runtime.benchmarking import benchmarker
do_bench_gpu = benchmarker.benchmark_gpu
res = benchmarker.benchmark_gpu(
lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median"
)
elif TORCH_VERSION_AFTER_2_3:
from torch._inductor.runtime.runtime_utils import do_bench_gpu
res = do_bench_gpu(lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median")
else:
from torch._inductor.utils import do_bench
res = do_bench(lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median")
return res

Expand Down

0 comments on commit e11201a

Please sign in to comment.