From ab049e7aa74e570313a7d61d78aac396640ab8d3 Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Wed, 29 May 2024 21:53:54 -0700 Subject: [PATCH 1/7] Graceful handling of cpp extensions --- README.md | 11 +++++++---- setup.py | 4 +++- test/test_ops.py | 5 +++++ torchao/__init__.py | 10 ++++++++-- torchao/dtypes/__init__.py | 6 +++++- 5 files changed, 28 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index a573df54c8..aff09af24f 100644 --- a/README.md +++ b/README.md @@ -34,17 +34,20 @@ pip install -r dev-requirements.txt There are two options; -If you plan to be developing the library run: ```Shell -python setup.py develop +pip install -e . ``` If you want to install from source run ```Shell -python setup.py install +pip install . ``` ** Note: -Since we are building pytorch c++/cuda extensions by default, running `pip install .` will -not work. +If you are running into any issues while building `ao` cpp extensions you can instead build using + +```shell +USE_CPP=0 pip install . +``` ### Quantization diff --git a/setup.py b/setup.py index f5ff1128fa..014a36832c 100644 --- a/setup.py +++ b/setup.py @@ -19,6 +19,8 @@ def read_requirements(file_path): # Determine the package name based on the presence of an environment variable package_name = "torchao-nightly" if os.environ.get("TORCHAO_NIGHTLY") else "torchao" version_suffix = os.getenv("VERSION_SUFFIX", "") +use_cpp = os.getenv('USE_CPP') + # Version is year.month.date if using nightlies version = current_date if package_name == "torchao-nightly" else "0.2.0" @@ -92,7 +94,7 @@ def get_extensions(): package_data={ "torchao.kernel.configs": ["*.pkl"], }, - ext_modules=get_extensions(), + ext_modules=get_extensions() if use_cpp != "0" else None, install_requires=read_requirements("requirements.txt"), extras_require={"dev": read_requirements("dev-requirements.txt")}, description="Package for applying ao techniques to GPU models", diff --git a/test/test_ops.py b/test/test_ops.py index 4e463b4e26..1dce59f3a3 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -6,7 +6,12 @@ import unittest from parameterized import parameterized import pytest +import os +use_cpp = os.getenv('USE_CPP') + +if use_cpp: + pytest.skip("skipping cpp extensions") # torch.testing._internal.optests.generate_tests.OpCheckError: opcheck(op, ...): # test_faketensor failed with module 'torch' has no attribute '_custom_ops' (scroll up for stack trace) diff --git a/torchao/__init__.py b/torchao/__init__.py index c8f04c1d9e..676f5bdd52 100644 --- a/torchao/__init__.py +++ b/torchao/__init__.py @@ -1,12 +1,18 @@ import torch +import logging + _IS_FBCODE = ( hasattr(torch._utils_internal, "IS_FBSOURCE") and torch._utils_internal.IS_FBSOURCE ) if not _IS_FBCODE: - from . import _C - from . import ops + try: + from . import _C + from . import ops + except: + _C = None + logging.info("Skipping import of cpp extensions") from torchao.quantization import ( apply_weight_only_int8_quant, diff --git a/torchao/dtypes/__init__.py b/torchao/dtypes/__init__.py index d12a6da566..e5cba48b43 100644 --- a/torchao/dtypes/__init__.py +++ b/torchao/dtypes/__init__.py @@ -1,7 +1,11 @@ from .nf4tensor import NF4Tensor, to_nf4 from .uint4 import UInt4Tensor from .aqt import AffineQuantizedTensor, to_aq -from .float6_e3m2 import to_float6_e3m2, from_float6_e3m2 + +try: + from .float6_e3m2 import to_float6_e3m2, from_float6_e3m2 +except RuntimeError: + pass __all__ = [ "NF4Tensor", From 7b81743da42863db060857843983db9296f64427 Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Wed, 29 May 2024 21:58:34 -0700 Subject: [PATCH 2/7] update --- test/test_ops.py | 7 ++++--- torchao/dtypes/float6_e3m2.py | 6 +++++- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/test/test_ops.py b/test/test_ops.py index 1dce59f3a3..3c6a8e3fe2 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -8,10 +8,11 @@ import pytest import os -use_cpp = os.getenv('USE_CPP') +try: + import torchao.ops +except RuntimeError: + pytest.skip("torchao.ops not available") -if use_cpp: - pytest.skip("skipping cpp extensions") # torch.testing._internal.optests.generate_tests.OpCheckError: opcheck(op, ...): # test_faketensor failed with module 'torch' has no attribute '_custom_ops' (scroll up for stack trace) diff --git a/torchao/dtypes/float6_e3m2.py b/torchao/dtypes/float6_e3m2.py index 0c27838d06..416173df60 100644 --- a/torchao/dtypes/float6_e3m2.py +++ b/torchao/dtypes/float6_e3m2.py @@ -3,6 +3,10 @@ from torch.utils._triton import has_triton from torchao.ops import to_float6_e3m2_packed_cpu, to_float6_e3m2_unpacked_cpu, from_float6_e3m2_packed_cpu, from_float6_e3m2_unpacked_cpu +try: + import torchao.ops +except RuntimeError: + pytest.skip("torchao.ops not available") # some useful constants FLOAT6_E3M2_MAX = 28.0 @@ -121,7 +125,7 @@ def to_float6_e3m2(tensor: Tensor, no_bit_packing: bool = False) -> Tensor: if tensor.is_cpu: if no_bit_packing: return to_float6_e3m2_unpacked_cpu(tensor) - + *leading_dims, last_dim = tensor.shape return to_float6_e3m2_packed_cpu(tensor.view(-1, last_dim)).view(*leading_dims, -1) From a38866eab4ce587de4337a428b075bd2db450a54 Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Wed, 29 May 2024 23:12:03 -0700 Subject: [PATCH 3/7] push --- test/dtypes/test_float6_e3m2.py | 7 +++++++ test/test_ops.py | 1 - torchao/dtypes/__init__.py | 12 +++++++----- torchao/dtypes/float6_e3m2.py | 5 ----- 4 files changed, 14 insertions(+), 11 deletions(-) diff --git a/test/dtypes/test_float6_e3m2.py b/test/dtypes/test_float6_e3m2.py index c3365cffeb..304a78c563 100644 --- a/test/dtypes/test_float6_e3m2.py +++ b/test/dtypes/test_float6_e3m2.py @@ -5,6 +5,13 @@ parametrize, run_tests, ) + +try: + import torchao.ops +except RuntimeError: + pytest.skip("torchao.ops not available") + + from torchao.dtypes.float6_e3m2 import to_float6_e3m2, from_float6_e3m2 diff --git a/test/test_ops.py b/test/test_ops.py index 3c6a8e3fe2..b20e029380 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -6,7 +6,6 @@ import unittest from parameterized import parameterized import pytest -import os try: import torchao.ops diff --git a/torchao/dtypes/__init__.py b/torchao/dtypes/__init__.py index e5cba48b43..f5f0eac647 100644 --- a/torchao/dtypes/__init__.py +++ b/torchao/dtypes/__init__.py @@ -2,11 +2,6 @@ from .uint4 import UInt4Tensor from .aqt import AffineQuantizedTensor, to_aq -try: - from .float6_e3m2 import to_float6_e3m2, from_float6_e3m2 -except RuntimeError: - pass - __all__ = [ "NF4Tensor", "to_nf4", @@ -16,3 +11,10 @@ "to_float6_e3m2", "from_float6_e3m2", ] + +# CPP extensions +try: + from .float6_e3m2 import to_float6_e3m2, from_float6_e3m2 + __all__.extend(["to_float6_e3m2", "from_float6_e3m2") +except RuntimeError: + pass diff --git a/torchao/dtypes/float6_e3m2.py b/torchao/dtypes/float6_e3m2.py index 416173df60..609ea4a119 100644 --- a/torchao/dtypes/float6_e3m2.py +++ b/torchao/dtypes/float6_e3m2.py @@ -3,11 +3,6 @@ from torch.utils._triton import has_triton from torchao.ops import to_float6_e3m2_packed_cpu, to_float6_e3m2_unpacked_cpu, from_float6_e3m2_packed_cpu, from_float6_e3m2_unpacked_cpu -try: - import torchao.ops -except RuntimeError: - pytest.skip("torchao.ops not available") - # some useful constants FLOAT6_E3M2_MAX = 28.0 FLOAT6_E3M2_SMALLEST_SUBNORMAL = 0.0625 From 20f6dc2496ab68c9a8a8a1348d4fedd8ab88f891 Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Wed, 29 May 2024 23:12:18 -0700 Subject: [PATCH 4/7] yolo --- torchao/dtypes/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchao/dtypes/__init__.py b/torchao/dtypes/__init__.py index f5f0eac647..ee80c7b324 100644 --- a/torchao/dtypes/__init__.py +++ b/torchao/dtypes/__init__.py @@ -15,6 +15,6 @@ # CPP extensions try: from .float6_e3m2 import to_float6_e3m2, from_float6_e3m2 - __all__.extend(["to_float6_e3m2", "from_float6_e3m2") + __all__.extend(["to_float6_e3m2", "from_float6_e3m2"]) except RuntimeError: pass From 55c540d12a75e4604ede45b22e3c32b0ee936bd3 Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Wed, 29 May 2024 23:15:05 -0700 Subject: [PATCH 5/7] revert some changes' --- torchao/dtypes/float6_e3m2.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torchao/dtypes/float6_e3m2.py b/torchao/dtypes/float6_e3m2.py index 609ea4a119..0c27838d06 100644 --- a/torchao/dtypes/float6_e3m2.py +++ b/torchao/dtypes/float6_e3m2.py @@ -3,6 +3,7 @@ from torch.utils._triton import has_triton from torchao.ops import to_float6_e3m2_packed_cpu, to_float6_e3m2_unpacked_cpu, from_float6_e3m2_packed_cpu, from_float6_e3m2_unpacked_cpu + # some useful constants FLOAT6_E3M2_MAX = 28.0 FLOAT6_E3M2_SMALLEST_SUBNORMAL = 0.0625 @@ -120,7 +121,7 @@ def to_float6_e3m2(tensor: Tensor, no_bit_packing: bool = False) -> Tensor: if tensor.is_cpu: if no_bit_packing: return to_float6_e3m2_unpacked_cpu(tensor) - + *leading_dims, last_dim = tensor.shape return to_float6_e3m2_packed_cpu(tensor.view(-1, last_dim)).view(*leading_dims, -1) From 47e572b80d8a33bf17585ac8e6ad35cb4419adc2 Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Thu, 30 May 2024 11:14:19 -0700 Subject: [PATCH 6/7] Update __init__.py --- torchao/dtypes/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/torchao/dtypes/__init__.py b/torchao/dtypes/__init__.py index ee80c7b324..44077dab65 100644 --- a/torchao/dtypes/__init__.py +++ b/torchao/dtypes/__init__.py @@ -8,8 +8,6 @@ "UInt4Tensor" "AffineQuantizedTensor", "to_aq", - "to_float6_e3m2", - "from_float6_e3m2", ] # CPP extensions From ac44a52d858385245ce84fa681171a49435660a8 Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Thu, 30 May 2024 13:05:27 -0700 Subject: [PATCH 7/7] Update README.md --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index aff09af24f..8da9f67da2 100644 --- a/README.md +++ b/README.md @@ -34,19 +34,19 @@ pip install -r dev-requirements.txt There are two options; -If you plan to be developing the library run: ```Shell -pip install -e . +python setup.py develop ``` If you want to install from source run ```Shell -pip install . +python setup.py install ``` ** Note: If you are running into any issues while building `ao` cpp extensions you can instead build using ```shell -USE_CPP=0 pip install . +USE_CPP=0 python setup.py install ``` ### Quantization