|
2 | 2 | import timeit |
3 | 3 | import multiprocessing |
4 | 4 |
|
5 | | -from piquant import QuantDtype |
| 5 | +os.environ['CUDA_VISIBLE_DEVICES'] = '' # Disable CUDA |
| 6 | +os.environ['OMP_NUM_THREADS'] = str(multiprocessing.cpu_count()) # OpenMP |
| 7 | +os.environ['MKL_NUM_THREADS'] = str(multiprocessing.cpu_count()) # MKL |
6 | 8 |
|
7 | | -os.environ['CUDA_VISIBLE_DEVICES'] = '' # Force CPU usage |
8 | | -os.environ['OMP_NUM_THREADS'] = str(multiprocessing.cpu_count()) |
9 | | -os.environ['MKL_NUM_THREADS'] = str(multiprocessing.cpu_count()) |
10 | | - |
11 | | -import torch |
12 | | -from torch.ao.quantization.fx._decomposed import quantize_per_tensor |
13 | 9 | import piquant |
| 10 | +import torch |
| 11 | +import numpy as np |
14 | 12 | import matplotlib.pyplot as plt |
15 | 13 |
|
16 | | -num_runs = 1000 |
17 | | -numel = 27264000 # Value from realistic test |
18 | | - |
| 14 | +torch.set_num_threads(multiprocessing.cpu_count()) |
19 | 15 |
|
20 | | -def quantize_torch_fx(tensor: torch.Tensor, scale: int, zero_point: float) -> None: |
21 | | - return quantize_per_tensor( |
22 | | - tensor, scale=scale, zero_point=zero_point, quant_min=0, quant_max=255, dtype=torch.uint8 |
23 | | - ) |
24 | | - |
25 | | - |
26 | | -def quantize_torch_builtin(tensor: torch.Tensor, scale: int, zero_point: float) -> None: |
27 | | - return torch.quantize_per_tensor(tensor, scale=scale, zero_point=zero_point, dtype=torch.quint8).int_repr() |
| 16 | +NUM_RUNS: int = 1_000 |
| 17 | +NUMEL: int = 1000000 |
28 | 18 |
|
| 19 | +QUANT_DTYPES_TO_BENCH: list[torch.dtype] = [ |
| 20 | + torch.quint8, |
| 21 | + torch.quint4x2, |
| 22 | +] |
29 | 23 |
|
30 | | -def quantize_fast(tensor: torch.Tensor, scale: int, zero_point: float) -> None: |
31 | | - return piquant.quantize_torch( |
32 | | - tensor, config=piquant.QuantConfig(output_dtype=piquant.QuantDtype.UINT8, scale=scale, zero_point=zero_point) |
33 | | - ) |
| 24 | +def quantize_torch(t: torch.Tensor, scale: float, zp: int, dtype: torch.dtype) -> torch.tensor: |
| 25 | + return torch.quantize_per_tensor(t, scale=scale, zero_point=zp, dtype=dtype) |
34 | 26 |
|
35 | 27 |
|
36 | | -tensor = torch.rand(numel, device='cpu') |
37 | | -scale, zero_point = piquant.compute_quant_config_torch(tensor, target_quant_dtype=QuantDtype.UINT8) |
| 28 | +def quantize_piquant(t: torch.Tensor, scale: float, zp: int, dtype: torch.dtype) -> torch.tensor: |
| 29 | + return piquant.torch.quantize(t, scale=scale, zero_point=zp, dtype=dtype) |
38 | 30 |
|
39 | 31 |
|
40 | | -def benchmark_torch_fx_quant() -> None: |
41 | | - quantize_torch_fx(tensor, scale, zero_point) |
| 32 | +dtype_labels: list[str] = [] |
| 33 | +torch_times: list[float] = [] |
| 34 | +piquant_times: list[float] = [] |
42 | 35 |
|
| 36 | +for torch_d in QUANT_DTYPES_TO_BENCH: |
| 37 | + tensor = torch.rand(NUMEL, dtype=torch.float32, device='cpu') |
| 38 | + torch_results = [] |
| 39 | + results_piquant = [] |
43 | 40 |
|
44 | | -def benchmark_torch_quant() -> None: |
45 | | - quantize_torch_builtin(tensor, scale, zero_point) |
| 41 | + scale, zp = piquant.torch.compute_quant_params(tensor, dtype=torch_d) |
| 42 | + zp = int(zp) |
46 | 43 |
|
| 44 | + def _bench_torch() -> None: |
| 45 | + torch_results.append(quantize_torch(tensor, scale, zp, torch_d)) |
47 | 46 |
|
48 | | -def benchmark_fast_quant() -> None: |
49 | | - quantize_fast(tensor, scale, zero_point) |
| 47 | + def _bench_piquant() -> None: |
| 48 | + results_piquant.append(quantize_piquant(tensor, scale, zp, torch_d)) |
50 | 49 |
|
| 50 | + # Warmup runs |
| 51 | + _bench_torch() |
| 52 | + _bench_piquant() |
51 | 53 |
|
52 | | -time_torch_fx = timeit.timeit(benchmark_torch_fx_quant, number=num_runs) |
53 | | -time_torch = timeit.timeit(benchmark_torch_quant, number=num_runs) |
54 | | -time_fast = timeit.timeit(benchmark_fast_quant, number=num_runs) |
| 54 | + torch_time = timeit.timeit(_bench_torch, number=NUM_RUNS) |
| 55 | + piquant_time = timeit.timeit(_bench_piquant, number=NUM_RUNS) |
| 56 | + dtype_labels.append(str(torch_d).replace('torch.', '')) |
| 57 | + torch_times.append(torch_time) |
| 58 | + piquant_times.append(piquant_time) |
55 | 59 |
|
56 | | -print(f'Torch FX quantization time for {num_runs} runs: {time_torch_fx:.6f} seconds') |
57 | | -print(f'Torch quantization time for {num_runs} runs: {time_torch:.6f} seconds') |
58 | | -print(f'Fast quantization time for {num_runs} runs: {time_fast:.6f} seconds') |
| 60 | + # Verify that the results are the same |
| 61 | + for i in range(NUM_RUNS): # We compare dequantized results, because .int_repr() is implemented for packed types in torch |
| 62 | + dq_torch = torch_results[i].dequantize() |
| 63 | + dq_piquant = piquant.torch.dequantize(results_piquant[i], scale=scale, zero_point=zp, dtype=torch.float32) |
| 64 | + assert dq_torch.numel() == dq_piquant.numel() |
| 65 | + assert dq_torch.dtype == dq_piquant.dtype |
| 66 | + assert torch.allclose(dq_torch, dq_piquant, atol=1e-1) |
| 67 | + print(f'{dtype_labels[-1]:<10} | torch: {torch_time:.6f}s | piquant: {piquant_time:.6f}s') |
59 | 68 |
|
60 | | -labels = [ |
61 | | - 'Torch FX Quantize', |
62 | | - 'Torch Quantize', |
63 | | - 'piquant', |
64 | | -] |
65 | | -times = [time_torch_fx, time_torch, time_fast] |
66 | 69 |
|
67 | | -plt.figure(figsize=(6, 4)) |
68 | | -plt.bar(labels, times) |
69 | | -plt.ylabel(f'Time (seconds) for {num_runs} runs') |
70 | | -plt.title('Quantization Benchmark') |
71 | | -plt.savefig('benchmark.png', dpi=300) |
| 70 | +x = np.arange(len(dtype_labels)) |
| 71 | +width = 0.35 |
| 72 | +plt.figure(figsize=(8, 5)) |
| 73 | +plt.bar(x - width / 2, torch_times, width, label='torch') |
| 74 | +plt.bar(x + width / 2, piquant_times, width, label='piquant') |
| 75 | +plt.ylabel(f'Total time for {NUM_RUNS} runs (s)') |
| 76 | +plt.xticks(x, dtype_labels) |
| 77 | +plt.title('Quantization Benchmark: PyTorch vs. piquant') |
| 78 | +plt.legend() |
| 79 | +plt.tight_layout() |
| 80 | +plt.savefig('quant_benchmark.png', dpi=300) |
72 | 81 | plt.show() |
0 commit comments