tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Flex Attention + Packing with BlockMask support (#2363)

Browse Source
This commit is contained in:
Sung Ching Liu
2025-04-05 18:02:57 -04:00
committed by GitHub
parent e7e0cd97ce
commit a8f38c367c
7 changed files with 281 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/e2e/utils.py
View File

@@ -67,9 +67,21 @@ def require_torch_2_5_1(test_case):
return unittest.skipUnless(is_min_2_5_1(), "test requires torch>=2.5.1")(test_case)
def require_torch_2_6_0(test_case):
"""
Decorator marking a test that requires torch >= 2.6.0
"""
def is_min_2_6_0():
torch_version = version.parse(torch.__version__)
return torch_version >= version.parse("2.6.0")
return unittest.skipUnless(is_min_2_6_0(), "test requires torch>=2.6.0")(test_case)
def require_torch_lt_2_6_0(test_case):
"""
Decorator marking a test that requires torch >= 2.5.1
Decorator marking a test that requires torch < 2.6.0
"""
def is_max_2_6_0():