skip test for pytorch < 2.3

Author: gau-nernst

benchmarks/benchmark_adam_8bit.py

@@ -197,7 +197,7 @@ def evaluate_model(model, args):
 
                 step += 1
 
-                if args.profile and step == 50:
+                if args.profile and step == 20:
                     break
 
         if args.profile:

test/prototype/test_optim_8bit.py

@@ -11,6 +11,7 @@
 )
 from torchao.prototype.optim_8bit import AdamDTQ8bit, AdamWDTQ8bit
 from torchao.prototype.optim_8bit.subclass import quantize_8bit_with_qmap, QMAP_SIGNED
+from torchao.utils import TORCH_VERSION_AFTER_2_3
 
 try:
     import bitsandbytes as bnb
@@ -45,9 +46,10 @@ def test_quantize_8bit_with_qmap_compile(self, device):
         torch.testing.assert_close(actual_scale, expected_scale)
 
 
+@pytest.mark.skipif(bnb is None, reason="bitsandbytes is not availablle")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="bitsandbytes 8-bit Adam only works for CUDA")
+@pytest.mark.xfail(not TORCH_VERSION_AFTER_2_3, reason="torch.compile() fails for PyTorch < 2.3")
 class TestOptim8bit(TestCase):
-    @pytest.mark.skipif(bnb is None, reason="bitsandbytes is not availablle")
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="bitsandbytes 8-bit Adam only works for CUDA")
     @parametrize("optim_cls,bnb_optim_cls", [
         (AdamDTQ8bit, bnb.optim.Adam8bit),
         (AdamWDTQ8bit, bnb.optim.AdamW8bit),

torchao/prototype/optim_8bit/README.md

@@ -17,6 +17,8 @@ You can also change quantization block size (default 2048) by passing `block_siz
 
 **Other optimizers**: AdamW is also available as `AdamWDTQ8bit`.
 
+NOTE: this requires PyTorch >= 2.3
+
 ## Benchmarks
 
 Benchmark script for fine-tuning a [timm](https://github.com/huggingface/pytorch-image-models) model on [resisc45](https://huggingface.co/datasets/timm/resisc45) dataset is available at [benchmarks/benchmark_adam_8bit.py](../../../benchmarks/benchmark_adam_8bit.py).