Negative scale (min > max) issue during QAT

[minseop4898]

Hi, I'm using aimet_torch version 2.7.0.
While performing QAT training on sim.model, I encountered cases where min > max, which caused the computed scale to become negative.

To work around this, I only performed PTQ and saved the state_dict of sim.model after compute_encodings. Then I modified the internal aimet_torch code as shown below and reloaded the model for QAT. After this change, the issue no longer occurred.

• In aimet_torch/v2/quantization/affine/quantizer.py -> AffineQuantizerBase

...

def _is_min_max_quantizer(self):
    # return "min" in self._parameters and "max" in self._parameters
    return "min" in self._parameters

@property
def max(self):
    return self.min + torch.nn.functional.softplus(self.delta)

@max.setter
def max(self, value):
    softplus_delta = value - self.min
    self.delta.data = softplus_delta + torch.log(-torch.expm1(-softplus_delta))

...

def _reparametrize_to_min_max(self):
    # pylint: disable=attribute-defined-outside-init
    if self._is_min_max_quantizer():
        return

    is_initialized = self.is_initialized()

    self.register_quantization_parameter(
        "min", nn.Parameter(-torch.ones(self.shape))
    )
    self.register_quantization_parameter(
        "delta", nn.Parameter(torch.ones(self.shape))
    )
    # self.register_quantization_parameter(
    #     "max", nn.Parameter(torch.ones(self.shape))
    # )

...

• Load PTQ sim.model

# create dummy sim instance

...

model_state_dict = torch.load(model_path, map_location=device)
sim.model.load_state_dict(model_state_dict, strict=False)
for key, value in model_state_dict.items():
    if key.endswith(".max"):
        parts = key.split(".")
        target = sim.model
        for p in parts[:-1]:
            target = getattr(target, p)
        setattr(target, "max", value)

...

# perform QAT

Is there any official way to ensure min < max during QAT without modifying aimet_torch like this?
If not, are there any plans to handle this case in future versions?

Thanks!

[quic-kyunggeu]

Hi @minseop4898, this is a very good point. Your softplus solution looks like a really nice trick to avoid negative scale.
Although for right now, it will involve a lot of negative side effects to apply this solution directly in AIMET, especially in terms of API and backwards compatibility.

Instead, I made a simple fix in 2ed8305 which prevents negative scale without losing backwards compatibility. This fix will be first included in aimet-torch 2.11 (planned Jul 28).
Please let us know if you have any further question!

[minseop4898]

Thank you for your response, @quic-kyunggeu, and sorry for the delayed reply.
I have a follow-up question: could you clarify what you mean by "backwards compatibility"?
In the approach you shared, it looks like you're using abs and minimum, which are non-differentiable.
Wouldn’t it be more stable for QAT to use smooth alternatives like softplus?
Is the reason you didn’t go with that related to the backwards compatibility you mentioned?

[quic-kyunggeu]

@minseop4898 That's right. Theoretically, I don't see any problem in the softplus approach. But the problem is that, in history, many our user code has been assuming that quantizer.max are static parameters (torch.nn.Parameter).
For this historical reason, converting it to a property that is dynamically derived from self.min and self.max can cause many practical problems.
For example, many users are familiar with overwriting quantizer.max with a new value, like quantizer.max.copy_(1.0), and by doing so they expect it to actually update the quantization range as it always did.
This is what I meant by "backwards compatibility" in my original answer.

[minseop4898]

@quic-kyunggeu Got it. Thanks for the reply and for updating the code!