Feat (utils): replace weights with quantized ones

[Giuseppe5]

Reason for this PR

In certain instances, like learned round, the weight tensors has extra parameters attached to it that could make export somewhat complicated.

Changes Made in this PR

We perform a destructive replacement of the original weight tensor with its quantized counterparts. This allows for easier exports in many scenarios.

There is an optional flag to keeep track of the original weights.

After merging, we reset the quantizers, including setting rounding mode to round, which is what is most commonly supported during export process.

What is missing:

• Integration in the LLM entrypoint.
• Handling the WeightWithInputQuantizer (i.e., the one required by A2Q and A2Q+)

Testing Summary

New tests added.

[nickfraser]

See comments

[pablomlago]

This pattern of overriding values in config and then restoring to the original values appears multiple times. Can we extract this common functionality? E.g.:

from contextlib import contextmanager
@contextmanager
def override_config(**overrides):
    old = {}
    try:
        for k, v in overrides.items():
            old[k] = getattr(config, k)
            setattr(config, k, v)
        yield
    finally:
        for k, v in old.items():
            setattr(config, k, v)

and then use it like:

with override_config(
        REINIT_ON_STATE_DICT_LOAD=False,
        IGNORE_MISSING_KEYS=True,
    ):

[Giuseppe5]

I'm going to open this as an issue and do it specifically in its own PR

[pablomlago]

LGTM, feel free to ignore the following change.