diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 5dd16bae99..99d651ae5f 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -96,6 +96,9 @@ def get_device_type(): raise NotImplementedError("Unsloth currently only works on NVIDIA, AMD and Intel GPUs.") pass DEVICE_TYPE : str = get_device_type() +# HIP fails for autocast and other torch functions. Use CUDA instead +DEVICE_TYPE_TORCH = DEVICE_TYPE +if DEVICE_TYPE_TORCH == "hip": DEVICE_TYPE_TORCH = "cuda" @functools.cache def get_device_count(): @@ -146,7 +149,9 @@ pass # OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536. Reducing block sizes or `num_stages` if (major_torch >= 2 and minor_torch >= 8) or (major_torch > 2): os.environ["UNSLOTH_ENABLE_CCE"] = "0" -pass +elif DEVICE_TYPE == "hip": + # CCE also fails in HIP / AMD + os.environ["UNSLOTH_ENABLE_CCE"] = "0" # Fix other issues import importlib.util diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index bf7d441c38..e787f55532 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -87,7 +87,7 @@ import functools import warnings, subprocess, re, inspect, psutil, os, math from unsloth_zoo.utils import Version from importlib.metadata import version as importlib_version -from unsloth import DEVICE_TYPE, DEVICE_COUNT +from unsloth import DEVICE_TYPE, DEVICE_COUNT, DEVICE_TYPE_TORCH from unsloth_zoo.log import logger from unsloth_zoo.tokenizer_utils import ( patch_tokenizer as _patch_tokenizer, diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 8ff74872a3..596042288d 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -27,7 +27,7 @@ from transformers import __version__ as transformers_version from unsloth_zoo.utils import Version, _get_dtype from unsloth_zoo.hf_utils import dtype_from_config, add_dtype_kwargs, fix_lora_auto_mapping from unsloth_zoo.peft_utils import SKIP_QUANTIZATION_MODULES -from unsloth import DEVICE_TYPE, DEVICE_COUNT +from unsloth import DEVICE_TYPE, DEVICE_COUNT, DEVICE_TYPE_TORCH transformers_version = Version(transformers_version) # Transformers moved rotary embeddings out of all attention layers @@ -732,7 +732,7 @@ def LlamaModel_fast_forward( position_ids = torch.arange( past_key_values_length, seq_length + past_key_values_length, dtype = torch.int32, - device = f"{DEVICE_TYPE}:0", + device = f"{DEVICE_TYPE_TORCH}:0", ) position_ids = position_ids.unsqueeze(0).view(-1, seq_length) elif position_ids is not None: @@ -905,13 +905,13 @@ def LlamaModel_fast_forward( is_causal = True, sliding_window = self.config.sliding_window, )\ - .to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = DEVICE_TYPE,)\ + .to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = DEVICE_TYPE_TORCH,)\ .squeeze(0).squeeze(0) self.GA_mask = AttentionMaskConverter( is_causal = True, )\ - .to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = DEVICE_TYPE,)\ + .to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = DEVICE_TYPE_TORCH,)\ .squeeze(0).squeeze(0) pass pass @@ -1028,11 +1028,11 @@ def _LlamaModel_fast_forward_inference(attention_fast_forward_inference=LlamaAtt bsz, q_len, hd = X.shape assert(q_len == 1) # Get saved buffers to reduce memory movement - residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = f"{DEVICE_TYPE}:0") - _XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = f"{DEVICE_TYPE}:0") + residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = f"{DEVICE_TYPE_TORCH}:0") + _XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = f"{DEVICE_TYPE_TORCH}:0") XX, XX2 = _XX[0], _XX[1] - variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = f"{DEVICE_TYPE}:0") - temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = f"{DEVICE_TYPE}:0") + variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = f"{DEVICE_TYPE_TORCH}:0") + temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = f"{DEVICE_TYPE_TORCH}:0") temp_gates, temp_ups = tuple(temp_mlp[0].to(torch.device(x)) for x in range(DEVICE_COUNT)), tuple(temp_mlp[1].to(torch.device(x)) for x in range(DEVICE_COUNT)) seq_len = past_key_values[0][0].shape[-2] @@ -1196,10 +1196,14 @@ def CausalLM_fast_forward(fast_forward_inference): else: RETURN_LOGITS = os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1" # < 1024 Normal Unsloth uses less VRAM! - if bsz*q_len <= 1024: RETURN_LOGITS = True + if DEVICE_TYPE == "hip": + # [TODO] AMD GPUs fail on chunked_cross_entropy loss! + # RuntimeError: Triton Error [HIP]: Code: 1, Messsage: invalid argument + RETURN_LOGITS = False + elif bsz*q_len <= 1024: + RETURN_LOGITS = True if not RETURN_LOGITS and labels is not None: - n_items = kwargs.get("num_items_in_batch", None) if n_items is None: n_items = kwargs.get("n_items", None) @@ -1374,7 +1378,7 @@ class LlamaRotaryEmbedding(torch.nn.Module): partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 dim = getattr(config, "head_dim", None) if dim is None: dim = int((config.hidden_size // config.num_attention_heads)) - device = DEVICE_TYPE + device = DEVICE_TYPE_TORCH max_position_embeddings = config.max_position_embeddings pass @@ -1486,7 +1490,7 @@ class LlamaExtendedRotaryEmbedding(torch.nn.Module): base = config.rope_theta partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 dim = int((config.hidden_size // config.num_attention_heads)) - device = DEVICE_TYPE + device = DEVICE_TYPE_TORCH max_position_embeddings = config.max_position_embeddings pass @@ -1606,7 +1610,7 @@ class LongRopeRotaryEmbedding(torch.nn.Module): base = config.rope_theta partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 dim = int((config.hidden_size // config.num_attention_heads)) - device = DEVICE_TYPE + device = DEVICE_TYPE_TORCH max_position_embeddings = config.max_position_embeddings pass @@ -1760,7 +1764,7 @@ def unsloth_fast_generate( kwargs["pad_token_id"] = kwargs.pop("pad_token_id", model_eos_token_id) # Mixed precision autocast - with torch.inference_mode(), torch.autocast(device_type = DEVICE_TYPE, dtype = dtype): + with torch.inference_mode(), torch.autocast(device_type = DEVICE_TYPE_TORCH, dtype = dtype): output = self._old_generate(*args, **kwargs) pass @@ -2384,7 +2388,7 @@ class FastLlamaModel: pass model.get_input_embeddings().modules_to_save.default\ - .to(device = DEVICE_TYPE, dtype = new_dtype, non_blocking = True) + .to(device = DEVICE_TYPE_TORCH, dtype = new_dtype, non_blocking = True) model.get_input_embeddings().modules_to_save.default.requires_grad_(True) # [TODO] Move old embed_tokens to CPU - should be disk! @@ -2404,7 +2408,7 @@ class FastLlamaModel: pass model.get_output_embeddings().modules_to_save.default\ - .to(device = DEVICE_TYPE, dtype = new_dtype, non_blocking = True) + .to(device = DEVICE_TYPE_TORCH, dtype = new_dtype, non_blocking = True) model.get_output_embeddings().modules_to_save.default.requires_grad_(True) # [TODO] Move old lm_head to CPU - should be disk! @@ -2673,7 +2677,7 @@ class FastLlamaModel: pass model.get_input_embeddings().modules_to_save.default\ - .to(device = DEVICE_TYPE, dtype = new_dtype, non_blocking = True) + .to(device = DEVICE_TYPE_TORCH, dtype = new_dtype, non_blocking = True) model.get_input_embeddings().modules_to_save.default.requires_grad_(True) pass @@ -2689,7 +2693,7 @@ class FastLlamaModel: pass model.get_output_embeddings().modules_to_save.default\ - .to(device = DEVICE_TYPE, dtype = new_dtype, non_blocking = True) + .to(device = DEVICE_TYPE_TORCH, dtype = new_dtype, non_blocking = True) model.get_output_embeddings().modules_to_save.default.requires_grad_(True) pass diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index faab2d30b1..b547739df2 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -298,7 +298,12 @@ def MistralForCausalLM_fast_forward( else: RETURN_LOGITS = os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1" # < 1024 Normal Unsloth uses less VRAM! - if bsz * q_len <= 1024: RETURN_LOGITS = True + if DEVICE_TYPE == "hip": + # [TODO] AMD GPUs fail on chunked_cross_entropy loss! + # RuntimeError: Triton Error [HIP]: Code: 1, Messsage: invalid argument + RETURN_LOGITS = False + elif bsz*q_len <= 1024: + RETURN_LOGITS = True if not RETURN_LOGITS and labels is not None: n_items = kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index b90ad00cf8..b2704876e3 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -71,7 +71,7 @@ except: # Old HF Hub versions <= 0.0.25 from huggingface_hub.utils._token import get_token pass -from unsloth import DEVICE_TYPE, DEVICE_COUNT +from unsloth import DEVICE_TYPE, DEVICE_COUNT, DEVICE_TYPE_TORCH __all__ = [ "FastBaseModel", @@ -204,10 +204,10 @@ def unsloth_base_fast_generate( # Mixed precision autocast if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": - autocaster = torch.autocast(device_type = "cuda", dtype = torch.float16) + autocaster = torch.autocast(device_type = DEVICE_TYPE_TORCH, dtype = torch.float16) dtype = torch.float16 else: - autocaster = torch.autocast(device_type = "cuda", dtype = dtype) + autocaster = torch.autocast(device_type = DEVICE_TYPE_TORCH, dtype = dtype) # Prepare LoRA # state_dict = convert_lora_modules(self, dtype = dtype)