diff --git a/.gitignore b/.gitignore index 014a60a024..b7b6522364 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ __pycache__/ *.py[cod] *.class +unsloth_compiled_cache/ # C extensions *.so diff --git a/unsloth/__init__.py b/unsloth/__init__.py index f3fef871b6..a1a39fd192 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -50,20 +50,6 @@ os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" # "pinned_use_cuda_host_register:True,"\ # "pinned_num_register_threads:8" -# Hugging Face Hub faster downloads -if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ: - os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" -pass - -# Disable XET Cache for now -os.environ["HF_XET_HIGH_PERFORMANCE"] = "1" -os.environ["HF_XET_CHUNK_CACHE_SIZE_BYTES"] = "0" -os.environ["HF_XET_RECONSTRUCT_WRITE_SEQUENTIALLY"] = "0" -os.environ["HF_XET_NUM_CONCURRENT_RANGE_GETS"] = "64" -# More verbose HF Hub info -if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1": - os.environ["HF_HUB_VERBOSITY"] = "info" - # Log Unsloth is being used os.environ["UNSLOTH_IS_PRESENT"] = "1" @@ -225,12 +211,11 @@ elif DEVICE_TYPE == "xpu": # Check for unsloth_zoo try: unsloth_zoo_version = importlib_version("unsloth_zoo") - if Version(unsloth_zoo_version) < Version("2025.4.1"): - pass - # print( - # "Unsloth: Updating Unsloth-Zoo utilies to the latest version.\n"\ - # "To disable this, set `os.environ['UNSLOTH_DISABLE_AUTO_UPDATES'] = '1'`" - # ) + if Version(unsloth_zoo_version) < Version("2025.7.1"): + print( + "Unsloth: Please update Unsloth and Unsloth-Zoo to the latest version!\n"\ + "Do this via `pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth unsloth_zoo`" + ) # if os.environ.get("UNSLOTH_DISABLE_AUTO_UPDATES", "0") == "0": # try: # os.system("pip install --upgrade --no-cache-dir --no-deps unsloth_zoo") diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 5c955a3c8d..1c65246b31 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -78,7 +78,7 @@ HAS_CUDA_STREAM = False # INTEL GPU specific logic if DEVICE_TYPE == "xpu": # TODO: Changed here after adding XPU BNB support - HAS_XPU_STREAM = False + HAS_XPU_STREAM = True def get_ptr(x: Optional[torch.Tensor]): raise RuntimeError("XPU BNB support is not implemented yet. This function should not be called.") else: diff --git a/unsloth/models/__init__.py b/unsloth/models/__init__.py index 3378049014..89d3fd7630 100644 --- a/unsloth/models/__init__.py +++ b/unsloth/models/__init__.py @@ -19,7 +19,11 @@ from .qwen2 import FastQwen2Model from .qwen3 import FastQwen3Model from .qwen3_moe import FastQwen3MoeModel from .granite import FastGraniteModel -from .falcon_h1 import FastFalconH1Model +try: + from .falcon_h1 import FastFalconH1Model +except: + # transformers_version < 4.53.0 does not have falcon_h1 so silenty skip it for now + pass from .dpo import PatchDPOTrainer, PatchKTOTrainer from ._utils import is_bfloat16_supported, is_vLLM_available, __version__ from .rl import PatchFastRL, vLLMSamplingParams \ No newline at end of file diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 76eefc3c3c..f576d17dc1 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -142,6 +142,12 @@ warnings.filterwarnings(action = "ignore", category = RuntimeWarning, module = " import logging logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.CRITICAL+1) +def get_device_num(): + if DEVICE_TYPE == "xpu": + return torch.xpu.device_count() + else: + return torch.cuda.device_count() + # Ignore logging messages class HideLoggingMessage(logging.Filter): __slots__ = "text", @@ -740,7 +746,7 @@ def get_statistics(): pass pass try: - devices = torch.cuda.device_count() + devices = get_device_num() _get_statistics(f"{devices if devices <= 8 else 9}") except: pass @@ -767,7 +773,7 @@ BitsAndBytesConfig__init__ = BitsAndBytesConfig__init__.replace( ) exec(BitsAndBytesConfig__init__, globals()) -if torch.cuda.device_count() == 1: +if get_device_num() == 1: from accelerate.utils.dataclasses import DistributedType def _prepare_backend(self, *args, **kwargs): return None, DistributedType.NO import accelerate.state diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py index 25704301d1..bfa833be9b 100644 --- a/unsloth/models/cohere.py +++ b/unsloth/models/cohere.py @@ -128,7 +128,7 @@ def CohereAttention_fast_forward( past_key_value = (K, V) if use_cache else None # Attention module - if (not HAS_FLASH_ATTENTION and attention_mask is None): + if (not HAS_FLASH_ATTENTION and HAS_XFORMERS and attention_mask is None): # Xformers memory efficient attention # Also has Flash Attention v2 dispatching Q = Q.transpose(1, 2) diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index 0f88bbc55e..243922fc13 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -115,7 +115,7 @@ def GraniteAttention_fast_forward( past_key_value = (K, V) if use_cache else None # Attention module - if (not HAS_FLASH_ATTENTION and attention_mask is None): + if (not HAS_FLASH_ATTENTION and HAS_XFORMERS and attention_mask is None): # Xformers memory efficient attention Q = Q.transpose(1, 2) K = K.transpose(1, 2) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 7d56dac2ec..1253e641e7 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -85,6 +85,11 @@ from triton import __version__ as triton_version HAS_XFORMERS = xformers is not None BlockDiagonalCausalMask = xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFORMERS else None +def clean_gpu_cache(): + if DEVICE_TYPE == "xpu": + torch.xpu.empty_cache() + else: + torch.cuda.empty_cache() def original_apply_qkv(self, X): Q = self.q_proj(X) @@ -319,6 +324,13 @@ def LlamaAttention_fast_forward_inference( # Knn, Vnn = Knn, Vnn # pass + # when qlen==vlen and attn_mask is None, we should use causal attention + Q_len = Qn.shape[-2] + K_len = Knn.shape[-2] + if attention_mask is None and Q_len == K_len: + is_causal = True + else: + is_causal = False # Attention if bsz == 1: Qn *= self.scalar # See https://github.com/ggerganov/llama.cpp/issues/7805#issuecomment-2153349963 @@ -519,11 +531,18 @@ def LlamaAttention_fast_forward( V = V.transpose(1, 2) A = flash_attn_func(Q, K, V, causal = True) else: + # when qlen==vlen and attn_mask is None, we should use causal attention + Q_len = Q.shape[-2] + K_len = K.shape[-2] + if attention_mask is None and Q_len == K_len: + is_causal = True + else: + is_causal = False # Grouped query attention if SDPA_HAS_GQA: # Needs (batch_size, n_heads, seq_len, head_dim) # is_casual and attention_mask must not be both set! - A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False, enable_gqa = n_groups != 1) + A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = is_causal, enable_gqa = n_groups != 1) # Go back to (batch_size, seq_len, n_heads, head_dim) A = A.transpose(1, 2)#.contiguous() else: @@ -538,7 +557,7 @@ def LlamaAttention_fast_forward( Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous() # Needs (batch_size, n_heads, seq_len, head_dim) # is_casual and attention_mask must not be both set! - A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False) + A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = is_causal) # Go back to (batch_size, seq_len, n_heads, head_dim) A = A.transpose(1, 2).contiguous() pass @@ -1275,9 +1294,8 @@ def PeftModel_fast_forward( logits_to_keep = 0, **kwargs, ): - is_classification = "Classification" in str(type( self.base_model.model)) + is_classification = "Classification" in str(type(self.base_model.model)) if is_classification: - #causal_mask = causal_mask, return self.base_model( input_ids = input_ids, attention_mask = attention_mask, @@ -1287,7 +1305,7 @@ def PeftModel_fast_forward( output_hidden_states = output_hidden_states, return_dict = return_dict, **kwargs, - ) + ) else: return self.base_model( input_ids = input_ids, @@ -1752,10 +1770,11 @@ class FastLlamaModel: if not is_vLLM_available(): print("Unsloth: vLLM is not installed! Will use Unsloth inference!") fast_inference = False - major_version, minor_version = torch.cuda.get_device_capability() - if major_version < 7: - print("Unsloth: vLLM does not work on older GPUs - will switch to Unsloth inference!") - fast_inference = False + if DEVICE_TYPE == "cuda": + major_version, minor_version = torch.cuda.get_device_capability() + if major_version < 7: + print("Unsloth: vLLM does not work on older GPUs - will switch to Unsloth inference!") + fast_inference = False 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 @@ -1779,8 +1798,8 @@ class FastLlamaModel: num_gpus = torch.xpu.device_count() gpu_stats_snippet = f"Intel Toolkit: {gpu_version}." - # TODO: After adding vLLM support for XPU, changed this - vllm_version = "" + try: vllm_version = f" vLLM: {importlib_version('vllm')}." + except: vllm_version = "" else: raise ValueError(f"Unsloth: Unsupported device type: {DEVICE_TYPE}") @@ -2020,7 +2039,10 @@ class FastLlamaModel: import gc for _ in range(3): gc.collect() - torch.cuda.empty_cache()""" + if DEVICE_TYPE == "xpu": + torch.xpu.empty_cache() + else: + torch.cuda.empty_cache()""" debug_info = debug_info.split('\n') debug_info = "\n".join([debug_info[0]] + [spaces + x[8:] for x in debug_info[1:]]) @@ -2508,7 +2530,7 @@ class FastLlamaModel: # Remove old items to save VRAM for _ in range(3): gc.collect() - torch.cuda.empty_cache() + clean_gpu_cache() pass if train_lm_head: @@ -2519,7 +2541,7 @@ class FastLlamaModel: # Remove old items to save VRAM for _ in range(3): gc.collect() - torch.cuda.empty_cache() + clean_gpu_cache() pass pass @@ -2580,7 +2602,7 @@ class FastLlamaModel: # Clear deleted GPU items for _ in range(3): gc.collect() - torch.cuda.empty_cache() + clean_gpu_cache() pass # Patch for fast inference @@ -2695,10 +2717,21 @@ class FastLlamaModel: if lora_dropout == 0 and bias == "none": for idx, layer in enumerate(model.model.model.layers): + # Determine MLP module name (falcon_h1 has feed_forward, llama style has mlp) + if hasattr(layer, "mlp"): + mlp_module_name = "mlp" + elif hasattr(layer, "feed_forward"): + mlp_module_name = "feed_forward" + else: + logger.warning_once(f"Unsloth: No MLP module found in layer {idx} so skipping peft mlp patching") + continue + + mlp_module = getattr(layer, mlp_module_name) + # MLP patching - gate_proj = layer.mlp.gate_proj - up_proj = layer.mlp. up_proj - down_proj = layer.mlp.down_proj + gate_proj = mlp_module.gate_proj + up_proj = mlp_module. up_proj + down_proj = mlp_module.down_proj if hasattr(gate_proj, "lora_A") and \ hasattr( up_proj, "lora_A") and \ @@ -2711,7 +2744,7 @@ class FastLlamaModel: (len(getattr(down_proj, "lora_magnitude_vector", []) or []) == 0): # https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module - layer.mlp.forward = types.MethodType(_apply_lora_mlp, layer.mlp) + mlp_module.forward = types.MethodType(_apply_lora_mlp, mlp_module) n_mlp += 1 else: logger.warning_once( @@ -2796,7 +2829,7 @@ class FastLlamaModel: # Clear deleted GPU items for _ in range(3): gc.collect() - torch.cuda.empty_cache() + clean_gpu_cache() pass # Patch for fast inference diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index ef71c89c4a..a3e07e3b0f 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -99,7 +99,7 @@ def MistralAttention_fast_forward( past_key_value = (K, V) if use_cache else None # Attention module - if (not HAS_FLASH_ATTENTION and attention_mask is None): + if (not HAS_FLASH_ATTENTION and HAS_XFORMERS and attention_mask is None): # Xformers memory efficient attention Q = Q.transpose(1, 2) K = K.transpose(1, 2) @@ -191,15 +191,35 @@ def MistralForCausalLM_fast_forward( if causal_mask is None and past_key_values is None: bsz, q_len = input_ids.shape sliding_window = getattr(self.config, "sliding_window", None) - if sliding_window is None or sliding_window == "null" or sliding_window <= 0: - causal_mask = xformers.attn_bias.LowerTriangularMask() - elif q_len <= sliding_window: - causal_mask = xformers.attn_bias.LowerTriangularMask() - else: - causal_mask = xformers.attn_bias.BlockDiagonalCausalMask\ - .from_seqlens([q_len]*bsz)\ - .make_local_attention(window_size = sliding_window) - pass + + if HAS_XFORMERS and attention_mask is None: + if sliding_window is None or sliding_window == "null" or sliding_window <= 0: + causal_mask = xformers.attn_bias.LowerTriangularMask() + elif q_len <= sliding_window: + causal_mask = xformers.attn_bias.LowerTriangularMask() + else: + causal_mask = xformers.attn_bias.BlockDiagonalCausalMask\ + .from_seqlens([q_len]*bsz)\ + .make_local_attention(window_size = sliding_window) + + elif not HAS_XFORMERS and attention_mask is None: + if sliding_window is None or sliding_window == "null" or sliding_window <= 0 or q_len <= sliding_window: + # Fully causal mask + mask = torch.full((q_len, q_len), -torch.inf, device=input_ids.device) + mask = torch.triu(mask, diagonal=1) + attention_mask = mask.expand(bsz, 1, q_len, q_len) + else: + # Sliding window attention + q_indices = torch.arange(q_len, device=input_ids.device).view(-1, 1) + k_indices = torch.arange(q_len, device=input_ids.device).view(1, -1) + + causal_bool_mask = k_indices <= q_indices + window_bool_mask = (q_indices - k_indices) < sliding_window + + mask = torch.where(causal_bool_mask & window_bool_mask, 0.0, -torch.inf) + attention_mask = mask[None, None, :, :].expand(bsz, 1, q_len, q_len) + + attention_mask = attention_mask.to(dtype=_get_dtype(self.config.torch_dtype)) output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py index 83c9dbea0a..80bd6ee7c9 100644 --- a/unsloth/models/qwen3.py +++ b/unsloth/models/qwen3.py @@ -126,7 +126,7 @@ def Qwen3Attention_fast_forward( past_key_value = (K, V) if use_cache else None # Attention module - if (not HAS_FLASH_ATTENTION and attention_mask is None): + if (not HAS_FLASH_ATTENTION and HAS_XFORMERS and attention_mask is None): # Xformers memory efficient attention Q = Q.transpose(1, 2) K = K.transpose(1, 2) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 9b0f4e4aef..002e4d1863 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -26,6 +26,8 @@ import torch import inspect from collections import defaultdict from unsloth_zoo.rl_replacements import RL_REPLACEMENTS +from unsloth import DEVICE_TYPE + RL_EXTRA_ARGS = defaultdict(list) RL_FUNCTIONS = defaultdict(list) RL_PRE_ITEMS = defaultdict(list) @@ -258,7 +260,7 @@ def grpo_trainer__get_per_token_logps(function_name, function): if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1': self._autocast_dtype = torch.float16 os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - with torch.amp.autocast(device_type = 'cuda', dtype = self._autocast_dtype): + with torch.amp.autocast(device_type = DEVICE_TYPE, dtype = self._autocast_dtype): # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded logits = model( input_ids = input_ids, @@ -291,6 +293,58 @@ def grpo_trainer__get_per_token_logps(function_name, function): pass RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps) +def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): + if function_name != "_get_per_token_logps_and_entropies": return function + + # Just copy over from _get_per_token_logps replacement function above. For now this returns None anyway + def _get_per_token_logps_and_entropies(self, model, input_ids, attention_mask, logits_to_keep, batch_size = None, compute_entropy = False): + if True: # os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0': + return {"logps": None, "entropies": None} # Unsloth efficient GRPO + # Otherwise, calculate normally: + if not hasattr(self, '_autocast_dtype'): + self._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16 + if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1': self._autocast_dtype = torch.float16 + + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + with torch.amp.autocast(device_type = 'cuda', dtype = self._autocast_dtype): + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + logits = model( + input_ids = input_ids, + attention_mask = attention_mask, + logits_to_keep = logits_to_keep + 1, + ).logits + + entropies = None + if compute_entropy: + from trl.trainer.utils import entropy_from_logits + entropies = entropy_from_logits(logits) + + # logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred + return {"logps": logits, "entropies": entropies} + # input_ids = input_ids[:, -logits_to_keep:] + # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. + # See https://github.com/huggingface/trl/issues/2770 + # logits = logits[:, -logits_to_keep:] + # return logits + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + # logits = logits / self.temperature + # logps = selective_log_softmax(logits, input_ids) + + # row_indices, col_indices = torch.where(logps < -20) + + # # Method 1: Check if tensors have elements + # if len(row_indices) > 0 and len(col_indices) > 0: + # breakpoint() # Breakpoint triggered here + # print("Found high values!") + # return logps # compute logprobs for the input tokens + pass + pass + + function = inspect.getsource(_get_per_token_logps_and_entropies) + return function +pass +RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps_and_entropies) + grpo_compute_loss = RL_REPLACEMENTS["grpo_compute_loss"] grpo_compute_loss_slow = RL_REPLACEMENTS["grpo_compute_loss_slow"] UnslothEfficientGRPO = RL_REPLACEMENTS["UnslothEfficientGRPO"] @@ -319,14 +373,16 @@ def grpo_trainer_compute_loss(function_name, function): _input_ids = input_ids _logits_to_keep = logits_to_keep - per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep) + get_logps_func = lambda model, input_ids, attention_mask, logits_to_keep, batch_size=None, compute_entropy=False: self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep, batch_size) if hasattr(self, "_get_per_token_logps") else self._get_per_token_logps_and_entropies(model, input_ids, attention_mask, logits_to_keep, batch_size, compute_entropy)['logps'] + + per_token_logps = get_logps_func(model, input_ids, attention_mask, logits_to_keep) # Compute the KL divergence between the model and the reference model # _prepare_inputs doesn't return reference log probs anymore. We need to calculate it ourselves. # https://github.com/huggingface/trl/blob/05bc43e960396581e458195b8388efe6b82cae1f/trl/trainer/grpo_trainer.py#L1328 if self.beta != 0.0: with torch.inference_mode(), model.disable_adapter(): - ref_per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep) + ref_per_token_logps = per_token_logps = get_logps_func(model, input_ids, attention_mask, logits_to_keep) else: ref_per_token_logps = None # per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1