disable bfloat16

This commit is contained in:
Daniel Han 2025-03-17 18:28:46 -07:00
commit 8766b15f9b
2 changed files with 10 additions and 10 deletions

View file

@ -462,11 +462,11 @@ class FastModel(FastBaseModel):
if token is None: token = get_token()
SUPPORTS_BFLOAT16 = is_bfloat16_supported()
if dtype is None:
dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16
elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16:
logger.warning_once("Device does not support bfloat16. Will change to float16.")
dtype = torch.float16
# if dtype is None:
# dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16
# elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16:
# logger.warning_once("Device does not support bfloat16. Will change to float16.")
# dtype = torch.float16
assert(dtype in (torch.float16, torch.bfloat16, torch.float32))
patch_compiled_autograd()

View file

@ -202,11 +202,11 @@ class FastBaseModel:
get_statistics() # For debugging - we use a download counter to see if environments are not breaking
if dtype is None:
dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16
elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16:
logger.warning_once("Device does not support bfloat16. Will change to float16.")
dtype = torch.float16
# if dtype is None:
# dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16
# elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16:
# logger.warning_once("Device does not support bfloat16. Will change to float16.")
# dtype = torch.float16
assert(dtype in (torch.float16, torch.bfloat16, torch.float32))