Significant Slowdown in Inference Speed with Quantized Model in PyTorch 2.1 pt2e

siahuat0727: <p>I’ve recently encountered an issue with PyTorch 2.1 where the inference speed of a quantized model is significantly slower than its FP32 counterpart (running on CUDA). The quantized model’s inference is over 10 times slower.</p>
<p>Here’s the code snippet that reproduces this behavior:</p>
<pre><code class="lang-auto">from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
from torch.ao.quantization.quantizer.xnnpack_quantizer import (
    XNNPACKQuantizer,
    get_symmetric_quantization_config,
)
import torch
import torchvision
import time

def loop(model, inputs):
    # warmup
    with torch.no_grad():
        for i in range(20):
            model(inputs)

    s = time.time()
    res = []
    with torch.no_grad():
        for i in range(20):
            s = time.time() 
            model(inputs)
            res.append(time.time() - s)
    return res

   
if __name__ == '__main__':
    inputs = torch.randn(32, 3, 224, 224).cuda()
    model = torchvision.models.vgg16().eval().cuda()
    res_fp32 = loop(model, inputs)
    print(res_fp32)

    model = torch._export.capture_pre_autograd_graph(model, [inputs])
    quantizer = XNNPACKQuantizer().set_global(
        get_symmetric_quantization_config(is_per_channel=True)
    )
    model = prepare_pt2e(model, quantizer)
    with torch.no_grad():
        model(inputs)
    model = convert_pt2e(model)

    res_quant = loop(model, inputs)
    print(res_quant)
</code></pre>
<p>The result on my local machine:</p>
<pre><code class="lang-auto">[..., 0.055181026458740234, 0.055129289627075195]
[..., 1.3574013710021973, 1.354731559753418]
</code></pre>
<p>Here are some profiler results:</p>
<p>fp32 model<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/b/1bdb1cec2b7ebd64b23ba8057f44440bc658e970.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/1bdb1cec2b7ebd64b23ba8057f44440bc658e970" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/b/1bdb1cec2b7ebd64b23ba8057f44440bc658e970_2_690x211.jpeg" alt="image" data-base62-sha1="3Yql1tyETaDZysJvxUlGYpVhLUY" width="690" height="211" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/b/1bdb1cec2b7ebd64b23ba8057f44440bc658e970_2_690x211.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/1/b/1bdb1cec2b7ebd64b23ba8057f44440bc658e970_2_1035x316.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/b/1bdb1cec2b7ebd64b23ba8057f44440bc658e970.jpeg 2x" data-dominant-color="262626"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1280×392 232 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>quantized model<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/8/482febf4438c4036a20a0cc7096c9803baed6c72.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/482febf4438c4036a20a0cc7096c9803baed6c72" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/8/482febf4438c4036a20a0cc7096c9803baed6c72_2_690x211.jpeg" alt="image" data-base62-sha1="aiB34MLUiAL51zH4dLluTYdA5Lc" width="690" height="211" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/8/482febf4438c4036a20a0cc7096c9803baed6c72_2_690x211.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/4/8/482febf4438c4036a20a0cc7096c9803baed6c72_2_1035x316.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/8/482febf4438c4036a20a0cc7096c9803baed6c72.jpeg 2x" data-dominant-color="252525"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1280×392 220 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The performance running on the current Colab is similar. <a href="https://colab.research.google.com/drive/1G3B6yHGtP0AOM7nU_TeAZQ4QtOUozSt2?usp=sharing" rel="noopener nofollow ugc">Link to Colab</a></p>
<p>I’m seeking insights into any potential solutions or optimizations that I might be missing.</p>
<p>My goal is to evaluate the performance of the quantized model after PTQ.</p>
<p>Any advice or guidance would be greatly appreciated. Thank you.</p>

siahuat0727: <p>The bottleneck is per-channel quantization. For details, refer to <a href="https://github.com/pytorch/pytorch/issues/116865" rel="noopener nofollow ugc">here</a>.</p>

HDCharles: <p>so to summarize jerry’s comment in the issue, what you have is kind of a <em>fake</em> quantized model that would then be lowered to the actual backend in question.</p>
<p>see <a href="https://pytorch.org/tutorials/prototype/pt2e_quant_ptq.html#lowering-and-performance-evaluation" class="inline-onebox" rel="noopener nofollow ugc">(prototype) PyTorch 2 Export Post Training Quantization — PyTorch Tutorials 2.2.0+cu121 documentation</a></p>
<blockquote>
<p>The model produced at this point is not the final model that runs on the device, it is a reference quantized model that captures the intended quantized computation from the user, expressed as ATen operators and some additional quantize/dequantize operators, to get a model that runs on real devices, we’ll need to lower the model. For example, for the models that run on edge devices, we can lower with delegation and ExecuTorch runtime operators.</p>
</blockquote>
<p>Further, my understanding is that we (only kind of) have kernels for quantized ops on cuda at the moment so it doesn’t look like this method will work.</p>
<p>There is another option, as described in the recent blog posts (<a href="https://pytorch.org/blog/accelerating-generative-ai/" rel="noopener nofollow ugc">1</a>,<a href="https://pytorch.org/blog/accelerating-generative-ai-2/" rel="noopener nofollow ugc">2</a>,<a href="https://pytorch.org/blog/accelerating-generative-ai-3/" rel="noopener nofollow ugc">3</a>) Where you can apply quantization using the api’s in the <a href="https://github.com/pytorch-labs/ao" rel="noopener nofollow ugc">torchao repository</a>. HOWEVER it currently only works for linear ops, not convolutions so i doubt you’ll see a ton of improvement for VGG which is mainly convolutional. It works great for transformer models though.</p>
<pre><code class="lang-auto">import torch
from torchao.quantization import quant_api

# some user model and example input
model = torch.nn.Sequential(torch.nn.Linear(32, 64)).cuda().to(torch.bfloat16)
input = torch.randn(32,32, dtype=torch.bfloat16, device='cuda')

# convert linear modules to quantized linear modules
quant_api.change_linear_weights_to_int8_dqtensors(model)

# compile the model to improve performance
model_c = torch.compile(model, mode='max-autotune')
model_c(input)
</code></pre>

siahuat0727: <p>Thank you so much for your informative reply!</p>
<p>I have some more questions,</p>
<ol>
<li>I’m still a little confused, does this lowering already work now? I didn’t find tutorials about it, for example how to lowering to a cuda device.</li>
<li>I’m kinda curious about the relationship between torchao and pt2e quantization. Does torchao seem more likely to exist as a standalone project?</li>
</ol>

HDCharles: <ol>
<li>
<p>No it doesn’t.</p>
</li>
<li>
<p>they’re both made by our team (pytorch ao). Pt2e does a lot of things and needs graph information for some of its features, the torchao stuff does not. They may both be folded together at some point in the future. They’re both under active development. Torchao will continue to exist or be moved to core pytorch in the future. For your use case torchao is ready now, pt2e has had more development oriented towards CPU/edge and less focus on cuda.</p>
</li>
</ol>

siahuat0727: <p>Gotcha! Thank you so much!!!</p>