torchao init: do not load .so files for known incompatible torch version (#2908)

Summary:

Short term fix for #2901 to unblock
the 0.13.0 release.

Long version:
1. torchao's c++ kernels are not using libtorch and therefore are not
   guaranteed to work across different PyTorch versions
2. looks like we got lucky with (1) as torchao kernels just happened to
   work across PyTorch versions <= 2.8, but PyTorch nightlies in 2.9
   introduce a breaking ABI change (I don't know what specifically).
   Therefore, if we build torchao with torch 2.8, and then import it in
   an environment with torch 2.9+, the Python process will crash with
   `Aborted (core dumped)`.

For now, I just gate out the "known broken" case where we detect
that the torch version used to build torchao is < 2.9, and the torch
version in the environment when torchao is imported is >= 2.9. If this
is detected, this PR skips importing the `.so` files and logs a warning,
to at least have most of the torchao Python API still work and give the
user some information about how to get the custom kernels working.

For future releases, we'll need to make this more robust - leaving that
for future PRs.

Test Plan:

```bash
// install the 0.13.0 RC, built with PyTorch 2.8
with-proxy pip install torchao==0.13.0 --extra-index-url https://download.pytorch.org/whl/test/cu128

// copy over these changes to the local __init__.py file in the installation:
// ~/.conda/envs/pytorch_nightly/lib/python3.11/site-packages/torchao/__init__.py

// install PyTorch 2.9.x nightly
with-proxy pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128

// import torchao, verify no more crash and the warning message is emitted
(pytorch_nightly) [[email protected] ~/local]$ python -X faulthandler -c "import torch; print(torch.__version__); import torchao"
2.9.0.dev20250829+cu128
Skipping import of cpp extensions due to incompatible torch version 2.9.0.dev20250829+cu128 for torchao version 0.13.0+cu128
```

Reviewers:

Subscribers:

Tasks:

Tags:

Author: vkuzo

torchao/__init__.py

@@ -22,23 +22,43 @@
 
 logger = logging.getLogger(__name__)
 
-try:
-    from pathlib import Path
-
-    so_files = list(Path(__file__).parent.glob("_C*.so"))
-    if len(so_files) > 0:
-        for file in so_files:
-            torch.ops.load_library(str(file))
-        from . import ops
-
-    # The following library contains CPU kernels from torchao/experimental
-    # They are built automatically by ao/setup.py if on an ARM machine.
-    # They can also be built outside of the torchao install process by
-    # running the script `torchao/experimental/build_torchao_ops.sh <aten|executorch>`
-    # For more information, see https://github.com/pytorch/ao/blob/main/torchao/experimental/docs/readme.md
-    from torchao.experimental.op_lib import *  # noqa: F403
-except Exception as e:
-    logger.debug(f"Skipping import of cpp extensions: {e}")
+skip_loading_so_files = False
+# if torchao version has "+git", assume it's locally built and we don't know
+#   anything about the PyTorch version used to build it
+# otherwise, assume it's prebuilt by torchao's build scripts and we can make
+#   assumptions about the PyTorch version used to build it.
+if (not "+git" in __version__) and not ("unknown" in __version__):
+    # torchao v0.13.0 is built with PyTorch 2.8.0. We know that torchao .so
+    # files built using PyTorch 2.8.0 are not ABI compatible with PyTorch 2.9+.
+    # The following code skips importing the .so files if PyTorch 2.9+ is
+    # detected, to avoid crashing the Python process with "Aborted (core
+    # dumped)".
+    # TODO(#2901, and before next torchao release): make this generic for
+    # future torchao and torch versions
+    if __version__.startswith("0.13.0") and torch.__version__ > "2.8":
+        logger.warning(
+            f"Skipping import of cpp extensions due to incompatible torch version {torch.__version__} for torchao version {__version__}"
+        )
+        skip_loading_so_files = True
+
+if not skip_loading_so_files:
+    try:
+        from pathlib import Path
+
+        so_files = list(Path(__file__).parent.glob("_C*.so"))
+        if len(so_files) > 0:
+            for file in so_files:
+                torch.ops.load_library(str(file))
+            from . import ops
+
+        # The following library contains CPU kernels from torchao/experimental
+        # They are built automatically by ao/setup.py if on an ARM machine.
+        # They can also be built outside of the torchao install process by
+        # running the script `torchao/experimental/build_torchao_ops.sh <aten|executorch>`
+        # For more information, see https://github.com/pytorch/ao/blob/main/torchao/experimental/docs/readme.md
+        from torchao.experimental.op_lib import *  # noqa: F403
+    except Exception as e:
+        logger.warning(f"Skipping import of cpp extensions: {e}")
 
 from torchao.quantization import (
     autoquant,