Merge branch 'main' into main

pytorch · Aug 15, 2024 · 2146a5a · 2146a5a
2 parents 93ca471 + ffa88a4
commit 2146a5a
Show file tree

Hide file tree

Showing 61 changed files with 1,061 additions and 588 deletions.
diff --git a/.github/workflows/regression_test.yml b/.github/workflows/regression_test.yml
@@ -31,11 +31,17 @@ jobs:
             torch-spec: 'torch==2.3.0'
             gpu-arch-type: "cuda"
             gpu-arch-version: "12.1"
+          - name: CUDA 2.4
+            runs-on: linux.g5.12xlarge.nvidia.gpu
+            torch-spec: 'torch==2.4.0'
+            gpu-arch-type: "cuda"
+            gpu-arch-version: "12.1"
           - name: CUDA Nightly
             runs-on: linux.g5.12xlarge.nvidia.gpu
             torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cu121'
             gpu-arch-type: "cuda"
             gpu-arch-version: "12.1"
+
           - name: CPU 2.2.2
             runs-on: linux.4xlarge
             torch-spec: 'torch==2.2.2 --index-url https://download.pytorch.org/whl/cpu "numpy<2" '
@@ -46,6 +52,11 @@ jobs:
             torch-spec: 'torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu'
             gpu-arch-type: "cpu"
             gpu-arch-version: ""
+          - name: CPU 2.4
+            runs-on: linux.4xlarge
+            torch-spec: 'torch==2.4.0 --index-url https://download.pytorch.org/whl/cpu'
+            gpu-arch-type: "cpu"
+            gpu-arch-version: ""
           - name: CPU Nightly
             runs-on: linux.4xlarge
             torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cpu'

diff --git a/benchmarks/benchmark_aq.py b/benchmarks/benchmark_aq.py
@@ -6,7 +6,7 @@
     Int4WeightOnlyQuantizedLinearWeight,
 )
 from torchao.utils import (
-    TORCH_VERSION_AFTER_2_4,
+    TORCH_VERSION_AT_LEAST_2_4,
 )
 from torchao.quantization.quant_api import (
     _replace_with_custom_fn_if_matches_filter,
@@ -105,7 +105,7 @@ def _bench_quantized_tensor_subclass_perf(api, ref_api, kwargs=None):
     print(f"elapsed time: {elapsed_time}, ref elapsed time: {ref_elapsed_time}")
     assert elapsed_time < 1.05 * ref_elapsed_time
 
-if __name__ == "__main__" and TORCH_VERSION_AFTER_2_4 and torch.cuda.is_available():
+if __name__ == "__main__" and TORCH_VERSION_AT_LEAST_2_4 and torch.cuda.is_available():
     from torchao.quantization.quant_api import change_linear_weights_to_int8_dqtensors
     _bench_quantized_tensor_subclass_perf(change_linear_weights_to_int8_dqtensors, _ref_change_linear_weights_to_int8_dqtensors)
 

diff --git a/benchmarks/intmm.py b/benchmarks/intmm.py
@@ -6,7 +6,7 @@
 import pathlib
 
 import torch
-from torchao.utils import TORCH_VERSION_AFTER_2_4, TORCH_VERSION_AFTER_2_2
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_4, TORCH_VERSION_AT_LEAST_2_2
 
 
 # Check if CUDA is available, if not, exit the script

diff --git a/test/dtypes/test_affine_quantized.py b/test/dtypes/test_affine_quantized.py
@@ -13,7 +13,7 @@
 import unittest
 import tempfile
 from torchao.utils import (
-    TORCH_VERSION_AFTER_2_5,
+    TORCH_VERSION_AT_LEAST_2_5,
 )
 
 
@@ -46,7 +46,7 @@ def test_weights_only(self):
                 torch.save(ql.state_dict(), f)
                 f.seek(0)
                 # `weights_only=True` is enabled for torch 2.5+
-                if TORCH_VERSION_AFTER_2_5:
+                if TORCH_VERSION_AT_LEAST_2_5:
                     _ = torch.load(f, weights_only=True)
                 else:
                     _ = torch.load(f, weights_only=False)

diff --git a/test/dtypes/test_bitnet.py b/test/dtypes/test_bitnet.py
@@ -4,9 +4,9 @@
 from torchao.prototype.dtypes import BitnetTensor
 from torchao.prototype.dtypes.uint2 import unpack_uint2
 from torchao.quantization.quant_api import _replace_with_custom_fn_if_matches_filter
-from torchao.utils import TORCH_VERSION_AFTER_2_4
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_4
 
-if not TORCH_VERSION_AFTER_2_4:
+if not TORCH_VERSION_AT_LEAST_2_4:
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
 @pytest.fixture(autouse=True)

diff --git a/test/dtypes/test_uint2.py b/test/dtypes/test_uint2.py
@@ -3,9 +3,9 @@
 import torch.nn as nn
 from torchao.prototype.dtypes import UInt2Tensor
 from torchao.prototype.dtypes.uint2 import unpack_uint2
-from torchao.utils import TORCH_VERSION_AFTER_2_4
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_4
 
-if not TORCH_VERSION_AFTER_2_4:
+if not TORCH_VERSION_AT_LEAST_2_4:
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
 @pytest.fixture

diff --git a/test/dtypes/test_uintx.py b/test/dtypes/test_uintx.py
@@ -6,7 +6,7 @@
 
 from torchao.dtypes.uintx.Uintx import to_uintx
 from torchao.quantization.quant_api import quantize_, uintx_weight_only
-from torchao.utils import TORCH_VERSION_AFTER_2_5
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
 from torchao.quantization.quant_primitives import (
     MappingType,
@@ -40,7 +40,7 @@ def forward(self, x):
 @pytest.mark.parametrize("group_size", group_sizes)
 @pytest.mark.parametrize("device", devices)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-@pytest.mark.skipif(not TORCH_VERSION_AFTER_2_5, reason="only works with fix in the nightly build")
+@pytest.mark.skipif(not TORCH_VERSION_AT_LEAST_2_5, reason="only works with fix in the nightly build")
 def test_uintx_weight_only_model_quant(bit_width, group_size, device):
     scale = 512
     fp16 = Linear16(scale, device)
@@ -54,7 +54,7 @@ def test_uintx_weight_only_model_quant(bit_width, group_size, device):
 @pytest.mark.parametrize("group_size", group_sizes)
 @pytest.mark.parametrize("device", devices)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-@pytest.mark.skipif(not TORCH_VERSION_AFTER_2_5, reason="only works with fix in the nightly build")
+@pytest.mark.skipif(not TORCH_VERSION_AT_LEAST_2_5, reason="only works with fix in the nightly build")
 def test_uintx_weight_only_quant(bit_width, group_size, device):
     input_float = torch.randn((1, 256), dtype=torch.float16, device = device)
     mapping_type = MappingType.SYMMETRIC

diff --git a/test/float8/test_base.py b/test/float8/test_base.py
@@ -16,9 +16,9 @@
 import torch
 import torch.nn as nn
 
-from torchao.utils import TORCH_VERSION_AFTER_2_4
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
-if not TORCH_VERSION_AFTER_2_4:
+if not TORCH_VERSION_AT_LEAST_2_5:
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
 

diff --git a/test/float8/test_compile.py b/test/float8/test_compile.py
@@ -11,9 +11,9 @@
 
 import pytest
 
-from torchao.utils import TORCH_VERSION_AFTER_2_4
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
-if not TORCH_VERSION_AFTER_2_4:
+if not TORCH_VERSION_AT_LEAST_2_5:
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
 import torch

diff --git a/test/float8/test_dtensor.py b/test/float8/test_dtensor.py
@@ -19,9 +19,9 @@
 
 import pytest
 
-from torchao.utils import TORCH_VERSION_AFTER_2_4
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
-if not TORCH_VERSION_AFTER_2_4:
+if not TORCH_VERSION_AT_LEAST_2_5:
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
 from torchao.float8 import Float8LinearConfig

diff --git a/test/float8/test_fsdp.py b/test/float8/test_fsdp.py
@@ -18,9 +18,9 @@
 
 import fire
 
-from torchao.utils import TORCH_VERSION_AFTER_2_4
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
-if not TORCH_VERSION_AFTER_2_4:
+if not TORCH_VERSION_AT_LEAST_2_5:
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
 import torch

diff --git a/test/float8/test_fsdp2/test_fsdp2.py b/test/float8/test_fsdp2/test_fsdp2.py
@@ -5,9 +5,9 @@
 import unittest
 from typing import Any, List
 
-from torchao.utils import TORCH_VERSION_AFTER_2_4
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
-if not TORCH_VERSION_AFTER_2_4:
+if not TORCH_VERSION_AT_LEAST_2_5:
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
 

diff --git a/test/float8/test_fsdp_compile.py b/test/float8/test_fsdp_compile.py
@@ -15,9 +15,9 @@
 
 import pytest
 
-from torchao.utils import TORCH_VERSION_AFTER_2_4
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
-if not TORCH_VERSION_AFTER_2_4:
+if not TORCH_VERSION_AT_LEAST_2_5:
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
 import torch

diff --git a/test/float8/test_inference_flows.py b/test/float8/test_inference_flows.py
@@ -12,11 +12,11 @@
 import pytest
 from unittest.mock import patch
 from torchao.utils import (
-    TORCH_VERSION_AFTER_2_4,
+    TORCH_VERSION_AT_LEAST_2_5,
     unwrap_tensor_subclass,
 )
 
-if not TORCH_VERSION_AFTER_2_4:
+if not TORCH_VERSION_AT_LEAST_2_5:
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
 import torch

diff --git a/test/float8/test_numerics_integration.py b/test/float8/test_numerics_integration.py
@@ -11,9 +11,9 @@
 
 import pytest
 
-from torchao.utils import TORCH_VERSION_AFTER_2_4
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
-if not TORCH_VERSION_AFTER_2_4:
+if not TORCH_VERSION_AT_LEAST_2_5:
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
 import torch