[Bug] Panic for the pytorch module on the shark-turbine 0.9.3 and 0.9.4

Peefy commented 8 months ago

My shark-turbine version is 0.9.3 and 0.9.4, it will panic for the following code

import torch

print("\nInstalled PyTorch, version:", torch.__version__)

torch.manual_seed(0)

class LinearModule(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.randn(in_features, out_features))
        self.bias = torch.nn.Parameter(torch.randn(out_features))

    def forward(self, input):
        return (input @ self.weight) + self.bias

linear_module = LinearModule(4, 3)
opt_linear_module = (
    torch.compile(linear_module, backend="turbine_cpu")
)
print("Compiled module using Turbine. New module type is", type(opt_linear_module))
args = torch.randn(4)
turbine_output = opt_linear_module(args)

print("Weight:", linear_module.weight)
print("Bias:", linear_module.bias)
print("Args:", args)
print("Output:", turbine_output)

However, shark-turbine 0.9.2 works well

ScottTodd commented 8 months ago

Can you clarify what you mean by "it will panic"? Does your Python interpreter crash? Is there an error message? (I usually associate "panic" with "kernel panic", which would be very unexpected here)

Peefy commented 8 months ago

Sorry, my python version is Python 3.11.8, the PyTorch version: 2.2.1 and the iree version 20240228.815. The error message is as follows:

Installed PyTorch, version: 2.2.1
Traceback (most recent call last):
  File "/Users/lingzhi/_Code/KCLOpenSource/kcl/a.py", line 20, in <module>
    torch.compile(linear_module, backend="turbine_cpu")
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/__init__.py", line 1824, in compile
    backend = _TorchCompileWrapper(backend, mode, options, dynamic)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/__init__.py", line 1692, in __init__
    self.compiler_fn = lookup_backend(backend)
                       ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/_dynamo/backends/registry.py", line 58, in lookup_backend
    _lazy_import_entry_point(compiler_fn)
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/_dynamo/backends/registry.py", line 110, in _lazy_import_entry_point
    compiler_fn = backend_eps[backend_name].load()
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/importlib/metadata/__init__.py", line 202, in load
    module = import_module(match.group('module'))
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<frozen importlib._bootstrap>", line 1204, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1176, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1147, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 690, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 940, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/shark_turbine/dynamo/backends/cpu.py", line 40, in <module>
    from ..passes import turbine_cpu_pass_pipeline
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/shark_turbine/dynamo/passes.py", line 56, in <module>
    @register_decomposition(torch.ops.aten._scaled_dot_product_flash_attention.default)
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/_decomp/__init__.py", line 185, in decomposition_decorator
    pytree.tree_map_(register, aten_op)
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/utils/_pytree.py", line 607, in tree_map_
    deque(map(func, flat_args), maxlen=0)  # consume and exhaust the iterable
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/_decomp/__init__.py", line 182, in register
    _add_op_to_registry(registry, op, fn)
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/_decomp/__init__.py", line 55, in _add_op_to_registry
    raise RuntimeError(f"duplicate registrations for {op_overload}")
RuntimeError: duplicate registrations for aten._scaled_dot_product_flash_attention.default

stellaraccident commented 8 months ago

@aviator19941 another example. Can you please verify this does not happen at head, and I can push a release tomorrow.