MacOS M1 - Torchao binaries are broken since 0.8.0 release

[atalman]

This applies to releases: 0.8.0 and 0.9.0

Here is the log:

pip install torchao==0.8.0         
Collecting torchao==0.8.0
  Downloading torchao-0.8.0-py3-none-any.whl.metadata (14 kB)
Downloading torchao-0.8.0-py3-none-any.whl (676 kB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 676.0/676.0 kB 1.6 MB/s eta 0:00:00
Installing collected packages: torchao
  Attempting uninstall: torchao
    Found existing installation: torchao 0.9.0
    Uninstalling torchao-0.9.0:
      Successfully uninstalled torchao-0.9.0
Successfully installed torchao-0.8.0
(py311) atalman@atalman-mbp aarch64_linux % python3                   
Python 3.11.8 (main, Feb 26 2024, 15:36:12) [Clang 14.0.6 ] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import torchao
import error: No module named 'triton'
>>> exit()

Here is the error log with torchao 0.9.0 and torch 2.6.0:

 File "/Users/atalman/Downloads/test_alloc/pytorch/.ci/aarch64_linux/triton.py", line 4, in do_stuff
    import torchao
  File "/Users/atalman/miniconda3/envs/py311/lib/python3.11/site-packages/torchao/__init__.py", line 46, in <module>
    from torchao.quantization import (
  File "/Users/atalman/miniconda3/envs/py311/lib/python3.11/site-packages/torchao/quantization/__init__.py", line 13, in <module>
    from .autoquant import (
  File "/Users/atalman/miniconda3/envs/py311/lib/python3.11/site-packages/torchao/quantization/autoquant.py", line 6, in <module>
    from torchao.dtypes import (
  File "/Users/atalman/miniconda3/envs/py311/lib/python3.11/site-packages/torchao/dtypes/__init__.py", line 1, in <module>
    from . import affine_quantized_tensor_ops
  File "/Users/atalman/miniconda3/envs/py311/lib/python3.11/site-packages/torchao/dtypes/affine_quantized_tensor_ops.py", line 6, in <module>
    from torchao.dtypes.affine_quantized_tensor import (
  File "/Users/atalman/miniconda3/envs/py311/lib/python3.11/site-packages/torchao/dtypes/affine_quantized_tensor.py", line 12, in <module>
    from torchao.quantization.quant_primitives import (
  File "/Users/atalman/miniconda3/envs/py311/lib/python3.11/site-packages/torchao/quantization/quant_primitives.py", line 202, in <module>
    register_custom_op = _register_custom_op(quant_lib)
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/atalman/miniconda3/envs/py311/lib/python3.11/site-packages/torchao/utils.py", line 217, in _register_custom_op
    from torch._inductor.decomposition import register_decomposition
  File "/Users/atalman/miniconda3/envs/py311/lib/python3.11/site-packages/torch/_inductor/decomposition.py", line 27, in <module>
    from torch._inductor.utils import pad_listlike
  File "/Users/atalman/miniconda3/envs/py311/lib/python3.11/site-packages/torch/_inductor/utils.py", line 50, in <module>
    from torch._inductor.runtime.hints import DeviceProperties
  File "/Users/atalman/miniconda3/envs/py311/lib/python3.11/site-packages/torch/_inductor/runtime/hints.py", line 67, in <module>
    from triton.compiler.compiler import AttrsDescriptor
ModuleNotFoundError: No module named 'triton.compiler'; 'triton' is not a package

[HDCharles]

if you uninstall triton and pytorch-triton, you can reproduce this on linux too, i think we just haven't been careful with triton imports, we should fix this

[atalman]

For linux we should be ok I think since Linux torch whl installs triton package by default

[malfet]

For linux we should be ok I think since Linux torch whl installs triton package by default

Not on the CPU-only version

[HDCharles]

@atalman does anything change going forward since were not publishing M1 binaries anymore?

[jainapurva]

#1747 I think we might need to fix M1 builds, if we're expecting more custom arm kernels

[atalman]

@HDCharles @jainapurva I think we would need to fix this issue

[jainapurva]

Closing the issue, as Pure-Python binaries work on M1.