Fix error when trying to access state_dict after activation quantization

[DN6]

What does this PR do?

Trying to access the model state dict after using activation quantization throws an unexpected error.

Traceback (most recent call last):
  File "/home/dhruv/diffusers/../scripts/create_dummy_flux.py", line 42, in <module>
    model.state_dict()
  File "/home/dhruv/miniconda3/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2219, in state_dict
    module.state_dict(
  File "/home/dhruv/miniconda3/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2219, in state_dict
    module.state_dict(
  File "/home/dhruv/miniconda3/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2219, in state_dict
    module.state_dict(
  [Previous line repeated 1 more time]
  File "/home/dhruv/miniconda3/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2216, in state_dict
    self._save_to_state_dict(destination, prefix, keep_vars)
  File "/home/dhruv/miniconda3/envs/diffusers/lib/python3.10/site-packages/optimum/quanto/nn/qmodule.py", line 150, in _save_to_state_dict
    destination[prefix + "weight"] = self.weight if keep_vars else self.weight.detach()
AttributeError: 'NoneType' object has no attribute 'detach'

This is because when replacing the LayerNorms in the model with QLayerNorm the weights are set to None, so this line will throw an error

optimum-quanto/optimum/quanto/nn/qmodule.py

Line 150 in 2202f84

destination[prefix + "weight"] = self.weight if keep_vars else self.weight.detach()

Snippet to test

import torch
from diffusers import FluxTransformer2DModel
from optimum.quanto import freeze, quantize, qint8, Calibration

torch_device = torch.device("cuda")


def get_dummy_inputs():
    return {
        "hidden_states": torch.randn(
            (1, 4096, 64), generator=torch.Generator("cpu").manual_seed(0)
        ).to(torch_device, torch.bfloat16),
        "encoder_hidden_states": torch.randn(
            (1, 512, 4096),
            generator=torch.Generator("cpu").manual_seed(0),
        ).to(torch_device, torch.bfloat16),
        "pooled_projections": torch.randn(
            (1, 768),
            generator=torch.Generator("cpu").manual_seed(0),
        ).to(torch_device, torch.bfloat16),
        "timestep": torch.tensor([1]).to(torch_device, torch.bfloat16),
        "img_ids": torch.randn(
            (4096, 3), generator=torch.Generator("cpu").manual_seed(0)
        ).to(torch_device, torch.bfloat16),
        "txt_ids": torch.randn(
            (512, 3), generator=torch.Generator("cpu").manual_seed(0)
        ).to(torch_device, torch.bfloat16),
        "guidance": torch.tensor([3.5]).to(torch_device, torch.bfloat16),
    }


model = FluxTransformer2DModel.from_pretrained(
    "hf-internal-testing/tiny-flux-transformer"
)
quantize(model, weights=qint8, activations=qint8)
model.to(torch_device)

with Calibration(momentum=0.9):
    model(**get_dummy_inputs())

freeze(model)
model.state_dict()

Fixes # (issue)

Before submitting

• 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 run all tests locally and make sure they pass.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[dacorvo]

Thanks for the pull-request: can you squash your commits and put a meaningful commit message to allow the CI to run ?

git reset --soft HEAD~2
git commit
git push -f

[github-actions]

This PR is stale because it has been open 15 days with no activity. Remove stale label or comment or this will be closed in 5 days.

[dacorvo]

The CI errors are fixed in #373 .

[dacorvo]

LGTM, thanks !