AssertionError: There are 1 incoming edges to node scaled_dot_product_attention

[ueno-hiroshi002]

Issue Type

Bug

Source

pip (model-compression-toolkit)

MCT Version

2.6.0

OS Platform and Distribution

Linux Ubuntu 22.04

Python version

3.10

Describe the issue

I tried post-training quantization on a DeiT model from timm(1.0.24).

from timm import create_model
float_model = create_model("deit_tiny_patch16_224.fb_in1k", pretrained=True)

However, I got the following error:

AssertionError: There are 1 incoming edges to node scaled_dot_product_attention:scaled_dot_product_attention, and they should be removed before deleting the node from the graph. 

The following model causes the same error:

create_model("maxvit_rmlp_pico_rw_256.sw_in1k", pretrained=True)

Expected behaviour

Success PTQ.

Code to reproduce the issue

import numpy as np
from timm import create_model
from timm.data.transforms_factory import create_transform
from timm.data import resolve_data_config
import model_compression_toolkit as mct

# Prepare Model
float_model = create_model("deit_tiny_patch16_224.fb_in1k", pretrained=True)
transform = create_transform(**resolve_data_config(float_model.pretrained_cfg, model=float_model), is_training=False)

# Prepare Data
def representative_data_gen():
    yield [np.random.random((1, 3, 224, 224))]

# MCT Quantization
quantized_model, quantization_info = mct.ptq.pytorch_post_training_quantization(
                in_module=float_model,
                representative_data_gen=representative_data_gen)

Log output

Traceback (most recent call last):
  File "/home/psnrdu/sss/model_augment/minimum_deit.py", line 23, in <module>
    quantized_model, quantization_info = mct.ptq.pytorch_post_training_quantization(
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/ptq/pytorch/quantization_facade.py", line 123, in pytorch_post_training_quantization
    tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_module,
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/runner.py", line 99, in core_runner
    graph = graph_preparation_runner(in_model,
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/graph_prep_runner.py", line 82, in graph_preparation_runner
    transformed_graph = get_finalized_graph(graph,
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/graph_prep_runner.py", line 127, in get_finalized_graph
    graph = substitute(initial_graph, fw_impl.get_substitutions_prepare_graph(fw_info))
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/common/substitutions/apply_substitutions.py", line 36, in substitute
    graph = substitution.substitute(graph, idn)
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/scaled_dot_product_attention.py", line 229, in substitute
    graph.remove_node(attention_node, new_graph_outputs=graph_outputs)
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/common/graph/base_graph.py", line 56, in wrapper
    result = method(self, *args, **kwargs)
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/common/graph/base_graph.py", line 511, in remove_node
    assert len(
AssertionError: There are 1 incoming edges to node scaled_dot_product_attention:scaled_dot_product_attention, and they should be removed before deleting the node from the graph.