Fix weight loading of weight_g_idx compressed-tensors parameters

[dsikka]

What does this PR do?

• There is currently a check in modeling_utils.py which replaces parameters with the weight_g substring in the parameter key with parametrizations.weight.original0. This incorrectly replaces the substring of compressed-tensors model parameters with the key weight_g_idx, resulting in a loading error.

transformers/src/transformers/modeling_utils.py

Line 4380 in 02a492a

if "weight_g" in key:

• Updates the src/transformers/quantizers/quantizer_compressed_tensors.py logic to temporarily turn off the weight_norm attribute such that the substring replacement does not occur in modeling_utils.py during weight loading. It then turns it back on in _process_model_after_weight_loading. This allows the models to be loaded correctly

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@SunMarc
@ArthurZucker

[SunMarc]

Thanks for the PR @dsikka ! Instead of a hotfix inside this quantization method, we can maybe solve the issue at its core. I think that you are having this issue because of a recent PR merged to add support for Timm models cc @qubvel. Before, we were only modifying the loaded keys but now, we are also fixing the state_dict which led to this situation. Also the real issue is that the condition to perform the modification of the keys is not strict enough.

    @staticmethod
    def _fix_state_dict_key_on_load(key):
        """Replace legacy parameter names with their modern equivalents. E.g. beta -> bias, gamma -> weight."""

        if "beta" in key:
            return key.replace("beta", "bias")
        if "gamma" in key:
            return key.replace("gamma", "weight")

        # to avoid logging parametrized weight norm renaming
        if hasattr(nn.utils.parametrizations, "weight_norm"):
            if "weight_g" in key:
                return key.replace("weight_g", "parametrizations.weight.original0")
            if "weight_v" in key:
                return key.replace("weight_v", "parametrizations.weight.original1")
        else:
            if "parametrizations.weight.original0" in key:
                return key.replace("parametrizations.weight.original0", "weight_g")
            if "parametrizations.weight.original1" in key:
                return key.replace("parametrizations.weight.original1", "weight_v")
        return key

A potential workaround would be to check if key.endswith("weight_g") or key.endswith("weight_v") instead

[qubvel]

This is also discussed in the latest fix for gamma/ beta renaming #35615

• 🚨🚨🚨 An attempt to fix #29554. Include 'LayerNorm.' in gamma/beta rename scope, optimize string search. #35615 (comment)

And already should be fixed, because the above PR introduced .endswith("weight_v") change.

@dsikka please update to the latest source code from main and let us know if you still face any issues, thanks 🤗

[dsikka]

perfect timing, that PR fixes our issue.

Thanks for linking!