generalize torch compatibility check

Author: liangel-02

.github/scripts/validate_binaries.sh

@@ -1,6 +1,67 @@
-pip install ${PYTORCH_PIP_PREFIX} torchao --index-url ${PYTORCH_PIP_DOWNLOAD_URL}
-# Intial smoke test, tries importing torchao
-python  ./test/smoke_tests/smoke_tests.py
-# Now we install dev-requirments and try to run the tests
-pip install -r dev-requirements.txt
-pytest test --verbose -s
+# pip install ${PYTORCH_PIP_PREFIX} torchao --index-url ${PYTORCH_PIP_DOWNLOAD_URL}
+# # Intial smoke test, tries importing torchao
+# python  ./test/smoke_tests/smoke_tests.py
+# # Now we install dev-requirments and try to run the tests
+# pip install -r dev-requirements.txt
+# pytest test --verbose -s
+
+# Testing compatibility
+# We know that torchao .so files built using PyTorch 2.8.0 are not ABI compatible with PyTorch 2.9+. (see #2919)
+# If the version of torch is not compatible with the version of torchao,
+# we expect to skip loading the .so files and a warning should be logged but no error
+
+# Function to check torchao import with configurable expectations
+check_torchao_import() {
+    local expect_warning="$1"
+    local warning_text="$2"
+    output=$(python -c "import torchao" 2>&1)
+    exit_code=$?
+
+    if [ $exit_code -ne 0 ]; then
+        echo "ERROR: Failed to import torchao"
+        echo "Output: $output"
+        exit 1
+    fi
+
+    warning_found=false
+    if [ -n "$warning_text" ] && echo "$output" | grep -i "$warning_text" > /dev/null; then
+        echo "Output: $output"
+        warning_found=true
+    fi
+
+    if [ "$expect_warning" != "$warning_found" ]; then
+        echo echo "FAILURE: expect_warning is $expect_warning but warning_found is $warning_found with message $output"
+        exit 1
+    fi
+}
+
+# Uninstall torch
+pip uninstall torch
+# Install compatible version of torch
+pip install torch==2.8.0
+# Build torchao
+pip install torchao==0.13.0
+# Test 1: Should import successfully without warning
+check_torchao_import "false" ""
+
+# Uninstall torch
+pip uninstall torch
+# Uninstall torchao
+pip uninstall torchao
+# Install compatible version of torch (nightly 2.9.0dev)
+pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu129
+# Build torchao (nightly 0.14.0dev...)
+pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu129
+# Test 1: Should import successfully without warning
+check_torchao_import "false" ""
+
+# Uninstall torch
+pip uninstall torch
+# Uninstall torchao
+pip uninstall torchao
+# Install non-ABI stable torch version with older version of torchao
+pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu129
+# Build torchao
+pip install torchao==0.13.0
+# Test 2: Should import with specific warning
+check_torchao_import "true" "Skipping import of cpp extensions due to incompatible torch version"

torchao/__init__.py

@@ -1,4 +1,5 @@
 import logging
+import re
 
 # torch/nested/_internal/nested_tensor.py:417: UserWarning: Failed to initialize NumPy: No module named 'numpy'
 import warnings
@@ -20,24 +21,50 @@
 except PackageNotFoundError:
     __version__ = "unknown"  # In case this logic breaks don't break the build
 
+
 logger = logging.getLogger(__name__)
 
+
+def parse_version(version_string):
+    """
+    Parses the major and minor of a torch version
+
+    Examples:
+        - 2.9.0dev... (pre-release) becomes [2, 9]
+        - 2.8.0 becomes [2, 8]
+        - 2.8.1 becomes [2, 8]
+    """
+    # Check for pre-release indicators
+    match = re.match(r"(\d+)\.(\d+)", version_string)
+    if match:
+        major, minor = map(int, match.groups())
+        return [major, minor]
+    else:
+        raise ValueError(f"Invalid version string format: {version_string}")
+
+
 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
+    # We know that torchao .so files built using PyTorch 2.8.0 are not ABI compatible with PyTorch 2.9+. (see #2919)
+    # The following code skips importing the .so files if incompatible torch version 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 str(torch.__version__) >= "2.9":
+    if parse_version(torch.__version__) == parse_version("2.8") and parse_version(
+        __version__
+    ) == parse_version("0.13"):
+        pass
+    elif parse_version(torch.__version__) == parse_version("2.9") and parse_version(
+        __version__
+    ) == parse_version("0.14"):
+        pass
+    else:
         logger.warning(
-            f"Skipping import of cpp extensions due to incompatible torch version {torch.__version__} for torchao version {__version__}"
+            f"Skipping import of cpp extensions due to incompatible torch version {torch.__version__} for torchao version {__version__} \
+            Please see GitHub issue #2919 for more info"
         )
         skip_loading_so_files = True
 

torchao/utils.py

@@ -378,6 +378,14 @@ def torch_version_at_least(min_version):
     return parse_version(torch.__version__) >= parse_version(min_version)
 
 
+def _torch_version_greater_than(min_version):
+    if is_fbcode():
+        return True
+
+    # Parser for local identifiers
+    return parse_version(torch.__version__) > parse_version(min_version)
+
+
 def _deprecated_torch_version_at_least(version_str: str) -> str:
     """
     Wrapper for existing TORCH_VERSION_AT_LEAST* variables that will log