pytorch / pytorch

Tensors and Dynamic neural networks in Python with strong GPU acceleration
https://pytorch.org
Other
84.4k stars 22.73k forks source link

Cannot export a quantized ResNet-18 model #139718

Open corehalt opened 2 weeks ago

corehalt commented 2 weeks ago

🐛 Describe the bug

After quantizing ResNet-18 model with PyTorch 2 Export Post Training Quantization it is not possible to export the model.

import torch

from torchvision.models.resnet import resnet18
from torch.ao.quantization.quantizer.xnnpack_quantizer import (
  XNNPACKQuantizer,
  get_symmetric_quantization_config,
)
from torch.ao.quantization.quantize_pt2e import (
  prepare_pt2e,
  convert_pt2e,
)
from executorch import exir

model = resnet18(pretrained=True)
model.to("cpu").eval()
example_inputs = (torch.rand(1, 3, 224, 224),)
exported_model = torch.export.export_for_training(model, example_inputs).module()

quantizer = XNNPACKQuantizer()
quantizer.set_global(get_symmetric_quantization_config(
    is_per_channel = True,
    is_dynamic = False,
    act_qmin = -128,
    act_qmax = 127,
    weight_qmin = -127,
    weight_qmax = 127)
)
prepared_model = prepare_pt2e(exported_model, quantizer)
quantized_model = convert_pt2e(prepared_model, use_reference_representation=True)

aten_dialect_program = torch.export.export(quantized_model, example_inputs)
edge_dialect_program = exir.to_edge(aten_dialect_program)
executorch_program = edge_dialect_program.to_executorch()
with open(f"exported/r18_quantized.pte", "wb") as f:
    f.write(executorch_program.buffer)

Get the error:

Traceback (most recent call last):
  File "repro.py", line 34, in <module>
    aten_dialect_program = torch.export.export(quantized_model, example_inputs)
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/export/__init__.py", line 270, in export
    return _export(
           ^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/export/_trace.py", line 1017, in wrapper
    raise e
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/export/_trace.py", line 990, in wrapper
    ep = fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/export/exported_program.py", line 114, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/export/_trace.py", line 1880, in _export
    export_artifact = export_func(  # type: ignore[operator]
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/export/_trace.py", line 1224, in _strict_export
    return _strict_export_lower_to_aten_ir(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/export/_trace.py", line 1333, in _strict_export_lower_to_aten_ir
    aten_export_artifact = lower_to_aten_callback(
                           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/export/_trace.py", line 637, in _export_to_aten_ir
    gm, graph_signature = transform(aot_export_module)(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 1246, in aot_export_module
    fx_g, metadata, in_spec, out_spec = _aot_export_function(
                                        ^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 1480, in _aot_export_function
    fx_g, meta = create_aot_dispatcher_function(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 522, in create_aot_dispatcher_function
    return _create_aot_dispatcher_function(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 759, in _create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
                               ^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 106, in aot_dispatch_export
    graph, _, _ = aot_dispatch_base_graph(
                  ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 154, in aot_dispatch_base_graph
    fw_module = _create_graph(
                ^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 54, in _create_graph
    fx_g = make_fx(
           ^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 2110, in wrapped
    return make_fx_tracer.trace(f, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 2048, in trace
    return self._trace_inner(f, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 2034, in _trace_inner
    t = dispatch_trace(
        ^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_compile.py", line 32, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 632, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1127, in dispatch_trace
    graph = tracer.trace(root, concrete_args)  # type: ignore[arg-type]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1631, in trace
    res = super().trace(root, concrete_args)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 632, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py", line 823, in trace
    (self.create_arg(fn(*args)),),
                     ^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1182, in wrapped
    out = f(*tensors)
          ^^^^^^^^^^^
  File "<string>", line 1, in <lambda>
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 693, in inner_fn
    outs = fn(*args)
           ^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 413, in _functionalized_f_helper
    f_outs = fn(*f_args)
             ^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 78, in inner_fn
    outs = fn(*args)
           ^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/utils.py", line 182, in flat_fn
    tree_out = fn(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 859, in functional_call
    out = PropagateUnbackedSymInts(mod).run(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/interpreter.py", line 146, in run
    self.env[node] = self.run_node(node)
                     ^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/experimental/symbolic_shapes.py", line 5498, in run_node
    result = super().run_node(n)
             ^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/interpreter.py", line 203, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/interpreter.py", line 275, in call_function
    return target(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_ops.py", line 716, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1230, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_ops.py", line 716, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1258, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_ops.py", line 716, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_ops.py", line 833, in handler
    return torch._library.utils.handle_dispatch_mode(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_library/utils.py", line 282, in handle_dispatch_mode
    return curr_mode.__torch_dispatch__(op_overload, overload_types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torch/_subclasses/functional_tensor.py", line 534, in __torch_dispatch__
    outs_unwrapped = func._op_dk(
                     ^^^^^^^^^^^^
RuntimeError: cannot mutate tensors with frozen storage

While executing %add_ : [num_users=1] = call_function[target=torch.ops.aten.add_.Tensor](args = (%to_27, %to_25), kwargs = {})
Original traceback:
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torchvision/models/resnet.py", line 285, in forward
    return self._forward_impl(x)
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torchvision/models/resnet.py", line 273, in _forward_impl
    x = self.layer1(x)
  File "/pyenv/versions/executorch-cpu/lib/python3.12/site-packages/torchvision/models/resnet.py", line 102, in forward
    out += identity

Versions

PyTorch version: 2.5.0+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 17.0.6 (https://github.com/llvm/llvm-project.git 6009708b4367171ccdbf4b5905cb6a803753fe18)
CMake version: version 3.30.5
Libc version: glibc-2.35

Python version: 3.12.2 (main, May 16 2024, 10:08:02) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-91-generic-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] executorch==0.4.0a0+6a085ff
[pip3] numpy==1.26.4
[pip3] onnx==1.17.0
[pip3] onnxruntime==1.19.2
[pip3] onnxslim==0.1.36
[pip3] torch==2.5.0+cpu
[pip3] torchaudio==2.5.0+cpu
[pip3] torchsr==1.0.4
[pip3] torchvision==0.20.0+cpu
[conda] Could not collect

@jerryzh168

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel @msaroufim

janeyx99 commented 2 weeks ago

@jerryzh168 @HDCharles Is there another API people should use from torchao?

corehalt commented 2 weeks ago

@janeyx99 @jerryzh168

Error mentioned above happens quite often with many different models, to mention a couple of them:

Model: fastvit_mci0.apple_mclip
Error: cannot mutate tensors with frozen storage

Model: darknet53.c2ns_in1k
Error: cannot mutate tensors with frozen storage

Model: torchvision.resnet50
Error: cannot mutate tensors with frozen storage

Model: torchvision.mobilenet_v2
Error: cannot mutate tensors with frozen storage

Another common error I encounter during quantization of other models is the following (seems that aten.sub ends up with pairs of inputs of data type [int32, int64])

Model: resnet50_clip.openai
Error: These operators are taking Tensor inputs with mismatched dtypes: defaultdict(<class 'dict'>, {<EdgeOpOverload: aten.sub.Tensor>: schema = aten::sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor: {'self': torch.int32, 'other': torch.int64, '__ret_0': torch.int32}})

Model: beit_base_patch16_224.in22k_ft_in22k
Error: These operators are taking Tensor inputs with mismatched dtypes: defaultdict(<class 'dict'>, {<EdgeOpOverload: aten.sub.Tensor>: schema = aten::sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor: {'self': torch.int32, 'other': torch.int64, '__ret_0': torch.int32}})

Model: caformer_b36.sail_in1k
Error: These operators are taking Tensor inputs with mismatched dtypes: defaultdict(<class 'dict'>, {<EdgeOpOverload: aten.sub.Tensor>: schema = aten::sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor: {'self': torch.int32, 'other': torch.int64, '__ret_0': torch.int32}})

Model: deit3_base_patch16_384.fb_in1k
Error: These operators are taking Tensor inputs with mismatched dtypes: defaultdict(<class 'dict'>, {<EdgeOpOverload: aten.sub.Tensor>: schema = aten::sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor: {'self': torch.int32, 'other': torch.int64, '__ret_0': torch.int32}})

Model: vit_base_patch14_dinov2.lvd142m
Error: These operators are taking Tensor inputs with mismatched dtypes: defaultdict(<class 'dict'>, {<EdgeOpOverload: aten.sub.Tensor>: schema = aten::sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor: {'self': torch.int32, 'other': torch.int64, '__ret_0': torch.int32}})