diff --git a/tests/llmcompressor/pytorch/utils/test_helpers.py b/tests/llmcompressor/pytorch/utils/test_helpers.py index 9e2bb373fe..d4e73c66e3 100644 --- a/tests/llmcompressor/pytorch/utils/test_helpers.py +++ b/tests/llmcompressor/pytorch/utils/test_helpers.py @@ -14,6 +14,7 @@ tensors_to_device, tensors_to_precision, ) +from tests.testing_utils import requires_gpu @pytest.mark.unit @@ -55,6 +56,7 @@ def test_tensors_to_device_cpu(tensors): os.getenv("NM_ML_SKIP_PYTORCH_TESTS", False), reason="Skipping pytorch tests", ) +@requires_gpu @pytest.mark.parametrize( "tensors", [ @@ -69,7 +71,6 @@ def test_tensors_to_device_cpu(tensors): [[torch.randn(1, 8)], torch.randn(8, 8)], ], ) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda availability") def test_tensors_to_device_cuda(tensors): out = tensors_to_device(tensors, "cuda") @@ -364,6 +365,7 @@ def test_tensors_module_forward(module, tensors, check_feat_lab_inp): os.getenv("NM_ML_SKIP_PYTORCH_TESTS", False), reason="Skipping pytorch tests", ) +@requires_gpu @pytest.mark.parametrize( "module,tensors,check_feat_lab_inp", [ @@ -417,7 +419,6 @@ def test_tensors_module_forward(module, tensors, check_feat_lab_inp): ), ], ) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda availability") def test_tensors_module_forward_cuda(module, tensors, check_feat_lab_inp): module = module.to("cuda") tensors = tensors_to_device(tensors, "cuda") @@ -471,6 +472,7 @@ def test_tensor_sparsity(tensor, dim, expected_sparsity): os.getenv("NM_ML_SKIP_PYTORCH_TESTS", False), reason="Skipping pytorch tests", ) +@requires_gpu @pytest.mark.parametrize( "tensor,dim,expected_sparsity", [ @@ -490,7 +492,6 @@ def test_tensor_sparsity(tensor, dim, expected_sparsity): ), ], ) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda availability") def test_tensor_sparsity_cuda(tensor, dim, expected_sparsity): tensor = tensor.to("cuda") sparsity = tensor_sparsity(tensor, dim) diff --git a/tests/llmcompressor/transformers/compression/test_has_gpu.py b/tests/llmcompressor/transformers/compression/test_has_gpu.py new file mode 100644 index 0000000000..aaa2f8c254 --- /dev/null +++ b/tests/llmcompressor/transformers/compression/test_has_gpu.py @@ -0,0 +1,13 @@ +import os + +import pytest +import torch + + +@pytest.mark.skipif(os.getenv("GITHUB_ACTIONS") != "true", reason="Only run for GHA") +def test_has_gpu(): + """ + This test exists purely to raise an error if + a runner performs transformers tests without a GPU + """ + assert torch.cuda.is_available() diff --git a/tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py b/tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py index d03ac3cd87..d9a6d0d317 100644 --- a/tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py +++ b/tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py @@ -30,6 +30,7 @@ modify_save_pretrained, patch_tied_tensors_bug, ) +from tests.testing_utils import requires_gpu @pytest.mark.parametrize( @@ -275,7 +276,7 @@ def test_model_reload(offload, torch_dtype, tie_word_embeddings, device_map, tmp assert torch.equal(model_dict[key].cpu(), reloaded_dict[key].cpu()) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires gpu") +@requires_gpu @pytest.mark.parametrize( "offload,torch_dtype,tie_word_embeddings,device_map", [ @@ -340,7 +341,7 @@ def test_model_shared_tensors( assert not torch.equal(lm_head, embed_tokens) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires gpu") +@requires_gpu @pytest.mark.parametrize( "offload,torch_dtype,tie_word_embeddings,device_map", [ @@ -356,6 +357,7 @@ def test_model_shared_tensors_gpu( ) +@requires_gpu @pytest.mark.parametrize( "model_stub, recipe, sparse_format, quant_format", [