Skip to content

Commit 1e402b9

Browse files
add test marker to run all tests with @require_bitsandbytes (huggingface#28278)
1 parent f3aa7db commit 1e402b9

File tree

2 files changed

+11
-2
lines changed

2 files changed

+11
-2
lines changed

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,4 +32,5 @@ doctest_optionflags="NUMBER NORMALIZE_WHITESPACE ELLIPSIS"
3232
doctest_glob="**/*.md"
3333
markers = [
3434
"flash_attn_test: marks tests related to flash attention (deselect with '-m \"not flash_attn_test\"')",
35+
"bitsandbytes: select (or deselect with `not`) bitsandbytes integration tests",
3536
]

src/transformers/testing_utils.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -966,9 +966,17 @@ def require_aqlm(test_case):
966966

967967
def require_bitsandbytes(test_case):
968968
"""
969-
Decorator for bits and bytes (bnb) dependency
969+
Decorator marking a test that requires the bitsandbytes library. Will be skipped when the library or its hard dependency torch is not installed.
970970
"""
971-
return unittest.skipUnless(is_bitsandbytes_available(), "test requires bnb")(test_case)
971+
if is_bitsandbytes_available() and is_torch_available():
972+
try:
973+
import pytest
974+
975+
return pytest.mark.bitsandbytes(test_case)
976+
except ImportError:
977+
return test_case
978+
else:
979+
return unittest.skip("test requires bitsandbytes and torch")(test_case)
972980

973981

974982
def require_optimum(test_case):

0 commit comments

Comments
 (0)