update

Author: a-r-r-o-w

tests/quantization/bnb/test_4bit.py

@@ -45,6 +45,7 @@
     require_peft_backend,
     require_torch,
     require_torch_accelerator,
+    require_torch_version_greater,
     require_transformers_version_greater,
     slow,
     torch_device,
@@ -860,7 +861,7 @@ def test_fp4_double_safe(self):
         self.test_serialization(quant_type="fp4", double_quant=True, safe_serialization=True)
 
 
-# @require_torch_version_greater("2.7.1")
+@require_torch_version_greater("2.7.1")
 class Bnb4BitCompileTests(QuantCompileTests):
     quantization_config = PipelineQuantizationConfig(
         quant_backend="bitsandbytes_8bit",

tests/quantization/test_torch_compile_utils.py

@@ -74,7 +74,6 @@ def _test_torch_compile_with_group_offload_leaf(
             "onload_device": torch.device("cuda"),
             "offload_device": torch.device("cpu"),
             "offload_type": "leaf_level",
-            "num_blocks_per_group": 1,
             "use_stream": use_stream,
         }
         pipe.transformer.enable_group_offload(**group_offload_kwargs)

tests/quantization/torchao/test_torchao.py

@@ -670,7 +670,7 @@ def test_torch_compile_with_group_offload_leaf(self):
         # weight.tensor_impl.{data,scale,zero_point}.device will be cpu
 
         # For use_stream=True:
-        # # NotImplementedError: AffineQuantizedTensor dispatch: attempting to run unimplemented operator/function: func=<OpOverload(op='aten.is_pinned', overload='default')>, types=(<class 'torchao.dtypes.affine_quantized_tensor.AffineQuantizedTensor'>,), arg_types=(<class 'torchao.dtypes.affine_quantized_tensor.AffineQuantizedTensor'>,), kwarg_types={}
+        # NotImplementedError: AffineQuantizedTensor dispatch: attempting to run unimplemented operator/function: func=<OpOverload(op='aten.is_pinned', overload='default')>, types=(<class 'torchao.dtypes.affine_quantized_tensor.AffineQuantizedTensor'>,), arg_types=(<class 'torchao.dtypes.affine_quantized_tensor.AffineQuantizedTensor'>,), kwarg_types={}
         super()._test_torch_compile_with_group_offload_leaf(quantization_config=self.quantization_config)