fix_tests (#17)

drisspg · Feb 11, 2024 · e2a3869 · e2a3869
1 parent 54a17d7
commit e2a3869
Showing 1 changed file with 18 additions and 8 deletions.
diff --git a/test/test_qlora.py b/test/test_qlora.py
@@ -38,7 +38,10 @@ def test_reconstruction(inpt_size: int, block_size: int, scaler_block_size: int)
 
 @unittest.skipIf(not bnb_available, "Bitsandbytes not available")
 @pytest.mark.parametrize("embed_dim", [256, 4096, 5120, 6656, 8192])
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not bnb_available,
+    reason="CUDA is not available or bitsandbytes not available",
+)
 def test_reconstruction_qlora_vs_bnb(embed_dim: int):
     torch.manual_seed(0)
     device = "cuda"
@@ -58,9 +61,11 @@ def test_reconstruction_qlora_vs_bnb(embed_dim: int):
     assert (nugs_diff - bnb_diff).abs() < 2e-1
 
 
-@pytest.mark.skipIf(not bnb_available, "Bitsandbytes not available")
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
 @pytest.mark.parametrize("embed_dim", [256, 4096, 5120, 6656, 8192])
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not bnb_available,
+    reason="CUDA is not available or bitsandbytes not available",
+)
 def test_binning_distribution(embed_dim: int):
     device = "cuda"
     input_weight = qlora.build_input_weight(embed_dim, device)
@@ -86,7 +91,7 @@ def test_binning_distribution(embed_dim: int):
 @pytest.mark.parametrize("embed_dim", [256, 4096, 5120, 6656, 8192])
 @pytest.mark.parametrize("compile", [True, False])
 @pytest.mark.parametrize("requires_grad", [True, False])
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
+@pytest.mark.xfail(reason="TORCH COMPILE No longer works here")
 def test_autograd_func_to_eager(embed_dim: int, compile: bool, requires_grad: bool):
     torch.manual_seed(0)
     device = "cuda"
@@ -103,10 +108,12 @@ def test_autograd_func_to_eager(embed_dim: int, compile: bool, requires_grad: bo
         out.sum().backward()
 
 
-@pytest.mark.skipIf(not bnb_available, "Bitsandbytes not available")
 @pytest.mark.parametrize("embed_dim", [256, 4096, 5120, 6656, 8192])
 @pytest.mark.parametrize("compile", [True, False])
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not bnb_available,
+    reason="CUDA is not available or bitsandbytes not available",
+)
 def test_bitsandbytes_linear_parity(embed_dim, compile):
     device = torch.device("cuda:0")
     input_weight = qlora.build_input_weight(embed_dim, device)
@@ -132,10 +139,12 @@ def qlora_linear(
     assert bnb_difference.max() < 0.5 * embed_dim
 
 
-@pytest.mark.skipIf(not bnb_available, "Bitsandbytes not available")
 @pytest.mark.parametrize("embed_dim", [256, 4096, 5120, 6656, 8192])
 @pytest.mark.parametrize("compile", [True, False])
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not bnb_available,
+    reason="CUDA is not available or bitsandbytes not available",
+)
 def test_bitsandbytes_mlp_parity(embed_dim, compile):
     device = torch.device("cuda:0")
     weights = qlora.get_mlp_weights(embed_dim, device)
@@ -167,6 +176,7 @@ def test_bitsandbytes_mlp_parity(embed_dim, compile):
 @pytest.mark.parametrize("run_backward", [True, False])
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
 def test_qlora_linear(embed_dim: int, compile: bool, r: int, dropout: float, run_backward: bool):
+    torch._dynamo.reset()
     torch.manual_seed(0)
     device = "cuda:0"
     # Analog for replacing first linear in MLP