From b107b432bb70628d76c2c25fde240a0bfde495ae Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 22 May 2024 20:40:57 +1000 Subject: [PATCH] Fix `is_bfloat16_supported` missing (#510) * Update llama.py * offload * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * continued pretraining trainer * Update trainer.py * Update trainer.py * Update trainer.py * Update trainer.py * is_bfloat16_supported * Update __init__.py * Update README.md * Update llama.py * is_bfloat16_supported * Update __init__.py --- unsloth/__init__.py | 7 +++++++ unsloth/trainer.py | 1 + 2 files changed, 8 insertions(+) diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 2dcf1e6a43..c8f4ca1041 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -41,6 +41,13 @@ try: except: raise ImportError("Pytorch is not installed. Go to https://pytorch.org/.\n"\ "We have some installation instructions on our Github page.") +pass + +# Fix up is_bf16_supported https://github.com/unslothai/unsloth/issues/504 +major_version, minor_version = torch.cuda.get_device_capability() +SUPPORTS_BFLOAT16 = (major_version >= 8) +def is_bf16_supported(): return SUPPORTS_BFLOAT16 +torch.cuda.is_bf16_supported = is_bf16_supported # We support Pytorch 2 # Fixes https://github.com/unslothai/unsloth/issues/38 diff --git a/unsloth/trainer.py b/unsloth/trainer.py index b234a98d8b..c8e00be231 100644 --- a/unsloth/trainer.py +++ b/unsloth/trainer.py @@ -16,6 +16,7 @@ from dataclasses import dataclass, field from typing import Optional from transformers import TrainingArguments from trl import SFTTrainer +from . import is_bfloat16_supported __all__ = [ "UnslothTrainingArguments",