Questions on QAT for Wav2Vec

Midora: <p>Hello,</p>
<p>I am currently facing an issue while trying to apply QAT to the pre-trained model retrieved through:</p>
<pre><code class="lang-auto">torchaudio.pipelines.WAV2VEC2_ASR_BASE_100H
</code></pre>
<p>At first I want to only apply QAT sequentially on the attention layers in the encoder, and then when successful apply it as well to the Conv layers in the feature extractor. I do not touch the Position Embedding part nor the LayerNorm. The purpose is to measure the difference between different configuration.</p>
<p>I am using this for each modules corresponding to a linear layer:</p>
<pre><code class="lang-auto">module.qconfig = torch.quantization.QConfig(
        activation=torch.quantization.FakeQuantize.with_args(observer=torch.quantization.MovingAverageMinMaxObserver,
                                                             dtype=torch.qint8),
        weight=torch.quantization.FakeQuantize.with_args(observer=torch.quantization.MovingAverageMinMaxObserver,
                                                         dtype=torch.qint8, qscheme=torch.per_tensor_affine))
</code></pre>
<p>The others <code>qconfig</code> are set to None.</p>
<p>And then on the whole model:</p>
<pre><code class="lang-auto">torch.quantization.prepare_qat(model, inplace=True)
</code></pre>
<p>There is no issue with training, and I did QAT both sequentially or on all the linear layers at the same time.</p>
<p>I am then saving the model using:</p>
<pre><code class="lang-auto">torch.save(torch.quantization.convert(model.to("cpu").eval(), inplace=False).state_dict(), model_save_path)
</code></pre>
<p>I can load the quantized model, and when printing it, it looks as expected.</p>
<p>My issue arises when I am running a validation on the quantized model. The issue is with the forward pass. I tried this for example:</p>
<pre><code class="lang-auto">torch.quantization.QuantStub()
torch.quantization.DeQuantStub()
</code></pre>
<pre><code class="lang-auto">import torch.quantization as quant

class QuantizedWav2VecModel(nn.Module):

    def __init__(self, pretrained_model):
        super(QuantizedWav2VecModel, self).__init__()
        self.pretrained_model = pretrained_model
        self.apply_quant = quant.QuantStub()
        self.apply_dequant = quant.DeQuantStub()

    def forward(self, input_values, input_lengths):
        quantized_input = self.apply_quant(input_values)
        output, output_size = self.pretrained_model(quantized_input, input_lengths)
        dequantized_output = self.apply_dequant(output)
        return dequantized_output, output_size
</code></pre>
<p>As some of my layers are not quantized, if I apply QuantStub then I run in this type of error:</p>
<p><code>NotImplementedError: Could not run 'aten::_slow_conv2d_forward' with arguments from the 'QuantizedCPU' backend.</code></p>
<p>I tried to change my forward pass, to check for each module and its childrens if I should apply QuantStub and DeQuantStub, but then I also need to keep in mind what goes into the forward pass of Wav2Vec, as some layers do not need the input lengths. This raise other errors as well.</p>
<p><strong>My questions are:</strong></p>
<p>1- Is there an other way to use QAT on this pretrained model, rather than the one I am using ? I would like to be able to apply QAT to whatever layer or group of layers I want.</p>
<p>2- Is the cleanest way to reimplement Wav2Vec and then load the weights from the pre-trained model ?</p>
<p>Note: I didn’t use <code>torch.ao.quantization</code> as I had the following error with it:</p>
<p><code>AssertionError:  nnq.Conv1d.from_float only works for Conv1d but got:&lt;class 'torch.nn.utils.parametrize.ParametrizedConv1d'&gt;</code> in the method <code>_ConvNd.from_float()</code></p>
<p>Thank you.</p>

jerryzh168: <p>I’d recommend to start with our new pt2e QAT tutorial: <a href="https://pytorch.org/tutorials/prototype/pt2e_quant_qat.html" class="inline-onebox" rel="noopener nofollow ugc">(prototype) PyTorch 2 Export Quantization-Aware Training (QAT) — PyTorch Tutorials 2.5.0+cu124 documentation</a></p>

Khady_KAMA1: <p>hello <a class="mention" href="/u/jerryzh168">@jerryzh168</a> , tried the pt2e export but it raise raise this issue :</p>
<pre><code class="lang-auto">` File "/opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 967, in call_function
return handler(tx, args, kwargs)
File "/opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 848, in builtin_dispatch
rv = fn(tx, args, kwargs)
File "/opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 766, in call_self_handler
result = self_handler(tx, *args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 1472, in call_getitem
return args[0].call_method(tx, "getitem", args[1:], kwargs)
File "/opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/misc.py", line 1082, in call_method
return super().call_method(tx, name, args, kwargs)
File "/opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/base.py", line 343, in call_method
unimplemented(f"call_method {self} {name} {args} {kwargs}")
File "/opt/conda/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 297, in unimplemented
raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: call_method GetAttrVariable(TupleVariable(length=2), shape) getitem (ConstantVariable(),) {}

from user code:
File "/home/jupyter/Wav2vec2_qat/wav2vec2_with_ai_edge_torch.py", line 111, in forward
attention_mask = torch.ones_like(input_values[:, :hidden_states.shape[1]])`
</code></pre>
<p><a class="mention" href="/u/midora">@Midora</a>  do you finally fixe the issue</p>

jerryzh168: <p>this seems like an torch.export issue, maybe you can post in <a href="https://discuss.pytorch.org/c/compile/41" class="inline-onebox">torch.compile - PyTorch Forums</a> and I can tag relevant people to take a look</p>