diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 2bfeed5751..5529e4eeef 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -69,8 +69,13 @@ except Exception as exception: raise exception pass +def is_hip(): + return bool(getattr(getattr(torch, "version", None), "hip", None)) + def get_device_type(): if hasattr(torch, "cuda") and torch.cuda.is_available(): + if is_hip(): + return "hip" return "cuda" elif hasattr(torch, "xpu") and torch.xpu.is_available(): return "xpu" @@ -79,7 +84,7 @@ pass DEVICE_TYPE : str = get_device_type() def get_device_count(): - if DEVICE_TYPE == "cuda": + if DEVICE_TYPE in ("cuda", "hip"): return torch.cuda.device_count() elif DEVICE_TYPE == "xpu": return torch.xpu.device_count() @@ -91,11 +96,12 @@ DEVICE_COUNT : int = get_device_count() # Reduce VRAM usage by reducing fragmentation # And optimize pinning of memory -if (DEVICE_TYPE == "cuda") and (os.environ.get("UNSLOTH_VLLM_STANDBY", "0")=="0"): +# TODO(billishyahao): need to add hip related optimization... +if (DEVICE_TYPE in ("cuda", "hip")) and (os.environ.get("UNSLOTH_VLLM_STANDBY", "0")=="0"): os.environ["PYTORCH_CUDA_ALLOC_CONF"] = \ "expandable_segments:True,"\ "roundup_power2_divisions:[32:256,64:128,256:64,>:32]" -elif (DEVICE_TYPE == "cuda") and (os.environ.get("UNSLOTH_VLLM_STANDBY", "0")=="1") and \ +elif (DEVICE_TYPE in ("cuda", "hip")) and (os.environ.get("UNSLOTH_VLLM_STANDBY", "0")=="1") and \ ("expandable_segments:True" in os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")): warnings.warn( "Unsloth: `UNSLOTH_VLLM_STANDBY` is on, but requires `expandable_segments` to be off.\n"\ @@ -153,6 +159,8 @@ if DEVICE_TYPE == "cuda": def is_bf16_supported(): return SUPPORTS_BFLOAT16 torch.cuda.is_bf16_supported = is_bf16_supported pass +elif DEVICE_TYPE == "hip": + SUPPORTS_BFLOAT16 = torch.cuda.is_bf16_supported() elif DEVICE_TYPE == "xpu": # torch.xpu.is_bf16_supported() does not have including_emulation # set SUPPORTS_BFLOAT16 as torch.xpu.is_bf16_supported() @@ -218,6 +226,9 @@ if DEVICE_TYPE == "cuda": "Unsloth will still run for now, but maybe it might crash - let's hope it works!" ) pass +elif DEVICE_TYPE == "hip": + # NO-OP for rocm device + pass elif DEVICE_TYPE == "xpu": # currently intel xpu will not support bnb, will add support in the future # TODO: check triton for intel installed properly. diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 0fa156d412..3813c3809c 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -90,7 +90,7 @@ else: if DEVICE_COUNT > 1: - if DEVICE_TYPE == "cuda": + if DEVICE_TYPE in ("cuda", "hip"): torch_gpu_device = torch.cuda.device elif DEVICE_TYPE == "xpu": torch_gpu_device = torch.xpu.device @@ -312,7 +312,7 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM: return out.t() if is_transposed else out pass # NVIDIA GPU Default Logic -elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: +elif DEVICE_TYPE in ("cuda", "hip") and HAS_CUDA_STREAM: @torch.inference_mode def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False): if quant_state is None: return W @@ -513,7 +513,7 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM: return out pass -elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: +elif DEVICE_TYPE in ("cuda", "hip") and HAS_CUDA_STREAM: def fast_gemv(X, W, quant_state, out = None): if quant_state is None: return torch_matmul(X, W, out = out) # For fast X @ W where seq_len == 1 diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 0ec3836d16..e3b70d93a5 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -452,7 +452,7 @@ pass # ============================================= # torch.cuda.amp.custom_fwd is deprecated >= 2.4 torch_version = torch.__version__ -if DEVICE_TYPE == "cuda": +if DEVICE_TYPE in ("cuda", "hip"): if Version(torch_version) < Version("2.4.0"): torch_amp_custom_fwd = torch.cuda.amp.custom_fwd torch_amp_custom_bwd = torch.cuda.amp.custom_bwd @@ -506,7 +506,7 @@ pass # ============================================= # Get Flash Attention v2 if Ampere (RTX 30xx, A100) -if DEVICE_TYPE == "cuda": +if DEVICE_TYPE in ("cuda", "hip"): import bitsandbytes as bnb from transformers import AutoTokenizer @@ -565,6 +565,44 @@ if DEVICE_TYPE == "cuda": # Tri Dao's benchmark shows xformers is faster for now. HAS_FLASH_ATTENTION = False pass +elif DEVICE_TYPE == "hip": + SUPPORTS_BFLOAT16 = True + if _is_package_available("flash_attn"): + # Check for CUDA linking errors "undefined symbol: _ZNK3c106SymIntltEl" + try: + try: + # See https://github.com/unslothai/unsloth/issues/1437 + from flash_attn.flash_attn_interface import flash_attn_gpu + except: + from flash_attn.flash_attn_interface import flash_attn_cuda + HAS_FLASH_ATTENTION = True + + # Also check for softcapping + from flash_attn import __version__ as flash_attn_version + HAS_FLASH_ATTENTION_SOFTCAPPING = Version(flash_attn_version) >= Version("2.6.3") + if not HAS_FLASH_ATTENTION_SOFTCAPPING: + print( + "Unsloth: If you want to finetune Gemma 2, upgrade flash-attn to version 2.6.3 or higher!\n"\ + "Newer versions support faster and less memory usage kernels for Gemma 2's attention softcapping!\n"\ + "To update flash-attn, do the below:\n"\ + '\npip install --no-deps --no-build-isolation --upgrade "flash-attn>=2.6.3"' + ) + except: + print( + "Unsloth: Your Flash Attention 2 installation seems to be broken?\n"\ + "A possible explanation is you have a new CUDA version which isn't\n"\ + "yet compatible with FA2? Please file a ticket to Unsloth or FA2.\n"\ + "We shall now use Xformers instead, which does not have any performance hits!\n"\ + "We found this negligible impact by benchmarking on 1x A100." + ) + + # Stop Flash Attention from importing! + import transformers.utils.import_utils + transformers.utils.import_utils.is_flash_attn_2_available = lambda *args, **kwargs: False + import transformers.utils + transformers.utils.is_flash_attn_2_available = lambda *args, **kwargs: False + + HAS_FLASH_ATTENTION = False elif DEVICE_TYPE == "xpu": SUPPORTS_BFLOAT16 = True diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 4143db93b8..af3ab8d3da 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1853,6 +1853,8 @@ class FastLlamaModel: if major_version < 7: print("Unsloth: vLLM does not work on older GPUs - will switch to Unsloth inference!") fast_inference = False + elif DEVICE_TYPE == "hip": + fast_inference = True if unsloth_vllm_standby and os.environ.get("UNSLOTH_VLLM_STANDBY", "0") == "0": raise RuntimeError("Unsloth: `unsloth_vllm_standby` is True, but environment variable `UNSLOTH_VLLM_STANDBY` is not set to 1!") pass @@ -1866,6 +1868,14 @@ class FastLlamaModel: gpu_version = torch.version.cuda gpu_stats_snippet = f"CUDA: {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit: {gpu_version}." + from importlib.metadata import version as importlib_version + try: vllm_version = f" vLLM: {importlib_version('vllm')}." + except: vllm_version = "" + elif DEVICE_TYPE == "hip": + gpu_stats = torch.cuda.get_device_properties(0) + gpu_version = torch.version.hip + gpu_stats_snippet = f"ROCm Toolkit: {gpu_version}." + from importlib.metadata import version as importlib_version try: vllm_version = f" vLLM: {importlib_version('vllm')}." except: vllm_version = "" diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 471ebb6a31..6c986a3bd3 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -278,6 +278,14 @@ class FastBaseModel: gpu_version = torch.version.cuda gpu_stats_snippet = f"CUDA: {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit: {gpu_version}." + from importlib.metadata import version as importlib_version + try: vllm_version = f" vLLM: {importlib_version('vllm')}." + except: vllm_version = "" + elif DEVICE_TYPE == "hip": + gpu_stats = torch.cuda.get_device_properties(0) + gpu_version = torch.version.hip + gpu_stats_snippet = f"ROCm Toolkit: {gpu_version}." + from importlib.metadata import version as importlib_version try: vllm_version = f" vLLM: {importlib_version('vllm')}." except: vllm_version = "" @@ -463,7 +471,7 @@ class FastBaseModel: # Clear deleted GPU items for _ in range(3): gc.collect() - if DEVICE_TYPE == "cuda": torch.cuda.empty_cache() + if DEVICE_TYPE in ("cuda", "hip"): torch.cuda.empty_cache() elif DEVICE_TYPE == "xpu": torch.xpu.empty_cache() pass @@ -558,7 +566,7 @@ class FastBaseModel: # Clear deleted GPU items for _ in range(3): gc.collect() - if DEVICE_TYPE == "cuda": + if DEVICE_TYPE in ("cuda", "hip"): torch.cuda.empty_cache() elif DEVICE_TYPE == "xpu": torch.xpu.empty_cache() @@ -627,7 +635,7 @@ class FastBaseModel: # Clear deleted GPU items for _ in range(3): gc.collect() - if DEVICE_TYPE == "cuda": + if DEVICE_TYPE in ("cuda", "hip"): torch.cuda.empty_cache() elif DEVICE_TYPE == "xpu": torch.xpu.empty_cache() @@ -663,7 +671,7 @@ class FastBaseModel: # Clear deleted GPU items for _ in range(3): gc.collect() - if DEVICE_TYPE == "cuda": + if DEVICE_TYPE in ("cuda", "hip"): torch.cuda.empty_cache() elif DEVICE_TYPE == "xpu": torch.xpu.empty_cache() @@ -728,7 +736,7 @@ class FastBaseModel: # Clear deleted GPU items for _ in range(3): gc.collect() - if DEVICE_TYPE == "cuda": + if DEVICE_TYPE in ("cuda", "hip"): torch.cuda.empty_cache() elif DEVICE_TYPE == "xpu": torch.xpu.empty_cache()