QAT with Aimet: Gradient & Training Behavior of param_quantizers min/max under two quant_schemes

[piantou]

When performing QAT with Aimet, assuming my model consists of conv + BN + Silu, I use the following code:

sim = QuantizationSimModel(model=model,
                           quant_scheme=QuantScheme.post_training_tf_enhanced,
                           dummy_input=dummy_input,
                           default_output_bw=8,
                           default_param_bw=8)

Then I set model.train() and define the optimizer:

optimizer = optim.SGD(
            model.parameters(),
            lr=learning_rate,
            momentum=0.9,
            weight_decay=weight_decay,
        )

(1) After setting quant_scheme=QuantScheme.post_training_tf_enhanced:
model.conv.weight obviously has requires_grad=True and is trainable. But do model.conv.param_quantizers.weight.min and model.conv.param_quantizers.weight.max also have requires_grad=True and are trainable?
(2) If I set quant_scheme=QuantScheme.training_range_learning_with_tf_init, what will the requires_grad status of model.conv.param_quantizers.weight.min andmodel.conv.param_quantizers.weight.maxbe?

What is the actual difference in the parameters being updated during QAT training between these two different quant_scheme settings?

[piantou]

I have observed that both model.conv.param_quantizers.weight.min and model.conv.param_quantizers.weight.max remain requires_grad=True under both quantization schemes. So what exactly is the effect of different quant_scheme settings on parameter updates in QAT? What are the concrete differences between them? Thank you very much for your help.

[quic-mtuttle]

Hi @piantou, good question. QuantScheme.post_training_tf and QuantScheme.training_range_learning_with_tf_init will both behave identically in current versions of aimet (aimet_torch 2.x). The "training_range_learning" versions exist only for legacy reasons, as they control whether min/max are learnable in legacy aimet_torch.v1.quantsim models.

In mainline aimet-torch >= 2.0, the quant_scheme ("tf" or "tf_enhanced") controls only how the min/max parameters are computed during sim.compute_encodings, and backprop to these is just controlled via requires_grad = True/False like any other torch parameter

[piantou]

Hi @piantou, good question. QuantScheme.post_training_tf and QuantScheme.training_range_learning_with_tf_init will both behave identically in current versions of aimet (aimet_torch 2.x). The "training_range_learning" versions exist only for legacy reasons, as they control whether min/max are learnable in legacy aimet_torch.v1.quantsim models.

In mainline aimet-torch >= 2.0, the quant_scheme ("tf" or "tf_enhanced") controls only how the min/max parameters are computed during sim.compute_encodings, and backprop to these is just controlled via requires_grad = True/False like any other torch parameter

Thanks for your help! Just to confirm: is sim.compute_encodings only performed once before training, and not needed during or after training?

[quic-mtuttle]

Thanks for your help! Just to confirm: is sim.compute_encodings only performed once before training, and not needed during or after training?

Yes that's right, you only need to compute encodings once before training. Recomputing after training will likely degrade accuracy