huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-09 19:34:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(server): check cuda capability when importing flash models

Browse Source
This commit is contained in:
OlivierDehaene 2023-04-19 12:48:04 +02:00
parent 2475aede61
commit 7c16352d1e
1 changed files with 12 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
server/text_generation_server/models/__init__.py
View File

@ -24,7 +24,18 @@ try:
FlashSantacoderSharded, FlashSantacoderSharded,
) )
FLASH_ATTENTION = torch.cuda.is_available() if torch.cuda.is_available():
major, minor = torch.cuda.get_device_capability()
is_sm75 = major == 7 and minor == 5
is_sm8x = major == 8 and minor >= 0
is_sm90 = major == 9 and minor == 0
supported = is_sm75 or is_sm8x or is_sm90
if not supported:
raise ImportError(f"GPU with CUDA capability {major} {minor} is not supported")
FLASH_ATTENTION = supported
else:
FLASH_ATTENTION = False
except ImportError: except ImportError:
logger.opt(exception=True).warning("Could not import Flash Attention enabled models") logger.opt(exception=True).warning("Could not import Flash Attention enabled models")
FLASH_ATTENTION = False FLASH_ATTENTION = False