update

Author: sayakpaul

tests/quantization/bnb/test_4bit.py

@@ -880,8 +880,5 @@ def test_torch_compile(self):
     def test_torch_compile_with_cpu_offload(self):
         super()._test_torch_compile_with_cpu_offload(quantization_config=self.quantization_config)
 
-    @pytest.mark.xfail(
-        reason="Test fails because of an illegal memory access.",
-    )
     def test_torch_compile_with_group_offload(self):
         super()._test_torch_compile_with_group_offload(quantization_config=self.quantization_config)

tests/quantization/test_torch_compile_utils.py

@@ -71,8 +71,7 @@ def _test_torch_compile_with_group_offload(self, quantization_config, torch_dtyp
         group_offload_kwargs = {
             "onload_device": torch.device("cuda"),
             "offload_device": torch.device("cpu"),
-            "offload_type": "block_level",
-            "num_blocks_per_group": 1,
+            "offload_type": "leaf_level",
             "use_stream": True,
             "non_blocking": True,
         }