reload quantized model after saving

[lucienfostier]

Hello, I'm trying to reload a model after quantizing it and saving it to a checkpoint file + the quantization map:

import json
import torch

import diffusers
import optimum.quanto
from optimum.quanto import quantization_map
from safetensors.torch import load_file
from optimum.quanto import requantize

vae = diffusers.AutoencoderKLTemporalDecoder.from_pretrained("hhyangcs/depth-any-video", subfolder="vae")
optimum.quanto.quantize(vae, weights=optimum.quanto.qfloat8)
optimum.quanto.freeze(vae)

vae.save_pretrained("./vae_quantized")

with open('quantization_map.json', 'w') as f:
  json.dump(quantization_map(vae), f)


state_dict = load_file('/content/vae_quantized/diffusion_pytorch_model.safetensors')
with open('quantization_map.json', 'r') as f:
  quantization_map = json.load(f)

# Create an empty model from your modeling code and requantize it
with torch.device('meta'):
  new_model = diffusers.AutoencoderKLTemporalDecoder()
requantize(new_model, state_dict, quantization_map, device=torch.device('cuda'))

I get the following error:

Click to Expand

``` [/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in load_state_dict(self, state_dict, strict, assign) 2582 2583 if len(error_msgs) > 0: -> 2584 raise RuntimeError( 2585 "Error(s) in loading state_dict for {}:\n\t{}".format( 2586 self.__class__.__name__, "\n\t".join(error_msgs)

RuntimeError: Error(s) in loading state_dict for AutoencoderKLTemporalDecoder:
size mismatch for encoder.conv_in.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.down_blocks.0.resnets.0.norm1.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.down_blocks.0.resnets.0.norm1.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.down_blocks.0.resnets.0.conv1.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.down_blocks.0.resnets.0.norm2.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.down_blocks.0.resnets.0.norm2.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.down_blocks.0.resnets.0.conv2.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.attentions.0.group_norm.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.attentions.0.group_norm.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.attentions.0.to_q.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.attentions.0.to_k.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.attentions.0.to_v.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.attentions.0.to_out.0.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.resnets.0.norm1.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.resnets.0.norm1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.resnets.0.conv1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.resnets.0.norm2.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.resnets.0.norm2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.resnets.0.conv2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.resnets.1.norm1.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.resnets.1.norm1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.resnets.1.conv1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.resnets.1.norm2.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.resnets.1.norm2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.mid_block.resnets.1.conv2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.conv_norm_out.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for encoder.conv_norm_out.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.conv_in.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.attentions.0.group_norm.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.attentions.0.group_norm.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.attentions.0.to_q.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.attentions.0.to_k.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.attentions.0.to_v.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.attentions.0.to_out.0.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.resnets.0.spatial_res_block.norm1.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.resnets.0.spatial_res_block.norm1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.resnets.0.spatial_res_block.conv1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.resnets.0.spatial_res_block.norm2.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.resnets.0.spatial_res_block.norm2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.resnets.0.spatial_res_block.conv2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.resnets.0.temporal_res_block.norm1.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.resnets.0.temporal_res_block.norm1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.resnets.0.temporal_res_block.conv1.weight: copying a param with shape torch.Size([512, 512, 3, 1, 1]) from checkpoint, the shape in current model is torch.Size([64, 64, 3, 1, 1]).
size mismatch for decoder.mid_block.resnets.0.temporal_res_block.conv1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.resnets.0.temporal_res_block.norm2.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.resnets.0.temporal_res_block.norm2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.mid_block.resnets.0.temporal_res_block.conv2.weight: copying a param with shape torch.Size([512, 512, 3, 1, 1]) from checkpoint, the shape in current model is torch.Size([64, 64, 3, 1, 1]).
size mismatch for decoder.mid_block.resnets.0.temporal_res_block.conv2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.0.spatial_res_block.norm1.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.0.spatial_res_block.norm1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.0.spatial_res_block.conv1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.0.spatial_res_block.norm2.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.0.spatial_res_block.norm2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.0.spatial_res_block.conv2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.0.temporal_res_block.norm1.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.0.temporal_res_block.norm1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.0.temporal_res_block.conv1.weight: copying a param with shape torch.Size([512, 512, 3, 1, 1]) from checkpoint, the shape in current model is torch.Size([64, 64, 3, 1, 1]).
size mismatch for decoder.up_blocks.0.resnets.0.temporal_res_block.conv1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.0.temporal_res_block.norm2.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.0.temporal_res_block.norm2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.0.temporal_res_block.conv2.weight: copying a param with shape torch.Size([512, 512, 3, 1, 1]) from checkpoint, the shape in current model is torch.Size([64, 64, 3, 1, 1]).
size mismatch for decoder.up_blocks.0.resnets.0.temporal_res_block.conv2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.1.spatial_res_block.norm1.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.1.spatial_res_block.norm1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.1.spatial_res_block.conv1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.1.spatial_res_block.norm2.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.1.spatial_res_block.norm2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.1.spatial_res_block.conv2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.1.temporal_res_block.norm1.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.1.temporal_res_block.norm1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.1.temporal_res_block.conv1.weight: copying a param with shape torch.Size([512, 512, 3, 1, 1]) from checkpoint, the shape in current model is torch.Size([64, 64, 3, 1, 1]).
size mismatch for decoder.up_blocks.0.resnets.1.temporal_res_block.conv1.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.1.temporal_res_block.norm2.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.1.temporal_res_block.norm2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.up_blocks.0.resnets.1.temporal_res_block.conv2.weight: copying a param with shape torch.Size([512, 512, 3, 1, 1]) from checkpoint, the shape in current model is torch.Size([64, 64, 3, 1, 1]).
size mismatch for decoder.up_blocks.0.resnets.1.temporal_res_block.conv2.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.conv_norm_out.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([64]).
size mismatch for decoder.conv_norm_out.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([64]).

</p>
</details>

What am I missing?

[dacorvo]

@lucienfostier you cannot use save_pretrained from the original model to save the quantized model.
Instead, you must save the corresponding state_dict (see the README, step 5 in the quantization workflow).

[github-actions]

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

[github-actions]

This issue was closed because it has been stalled for 5 days with no activity.