diff --git a/pyproject.toml b/pyproject.toml index e7b84f3c8e..7fa249e64c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,7 +60,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.1.1", + "unsloth_zoo>=2026.1.2", "torchvision", "unsloth[triton]", ] @@ -523,7 +523,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.1.1", + "unsloth_zoo>=2026.1.2", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,<=4.57.3", diff --git a/unsloth/__init__.py b/unsloth/__init__.py index d9633e8ec1..5b571cd456 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -79,7 +79,7 @@ from importlib.metadata import PackageNotFoundError # Check for unsloth_zoo try: unsloth_zoo_version = importlib_version("unsloth_zoo") - if Version(unsloth_zoo_version) < Version("2026.1.1"): + if Version(unsloth_zoo_version) < Version("2026.1.2"): 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`" @@ -126,6 +126,7 @@ from .import_fixes import ( fix_xformers_performance_issue, fix_vllm_aimv2_issue, fix_vllm_guided_decoding_params, + fix_vllm_pdl_blackwell, ignore_logger_messages, patch_ipykernel_hf_xet, patch_trackio, @@ -138,6 +139,7 @@ from .import_fixes import ( fix_xformers_performance_issue() fix_vllm_aimv2_issue() fix_vllm_guided_decoding_params() +fix_vllm_pdl_blackwell() ignore_logger_messages() patch_ipykernel_hf_xet() patch_trackio() @@ -149,6 +151,7 @@ fix_executorch() del fix_xformers_performance_issue del fix_vllm_aimv2_issue del fix_vllm_guided_decoding_params +del fix_vllm_pdl_blackwell del ignore_logger_messages del patch_ipykernel_hf_xet del patch_trackio diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index bb6996a3e3..1e05e462e9 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -556,3 +556,118 @@ def fix_huggingface_hub(): huggingface_hub.is_offline_mode = ( lambda: huggingface_hub.constants.HF_HUB_OFFLINE ) + + +def fix_vllm_pdl_blackwell(): + """ + Fix vLLM PDL (Programmatic Dependent Launch) bug on Blackwell GPUs (SM100). + + The issue: vLLM's LoRA Triton kernels use tl.extra.cuda.gdc_wait() for PDL + optimization on SM90+ GPUs. This fails on SM100 (B200/B100) during CUDA graph + capture because Triton's pipeliner can't handle gdc_wait in complex kernels. + + See: https://github.com/vllm-project/vllm/issues/30872 + """ + if importlib.util.find_spec("vllm") is None: + return + + # Check if any CUDA GPU is SM100 (Blackwell) + try: + import torch + + if not torch.cuda.is_available(): + return + + # Scan all GPUs for SM100 - fix applies globally via env var and monkey-patch + has_sm100 = False + sm100_gpu_name = None + for i in range(torch.cuda.device_count()): + major, minor = torch.cuda.get_device_capability(i) + if major == 10: + has_sm100 = True + sm100_gpu_name = torch.cuda.get_device_name(i) + break + + if not has_sm100: + return + except Exception: + return + + # Helper to check if module spec exists + def _spec_exists(name): + try: + return importlib.util.find_spec(name) is not None + except (ModuleNotFoundError, ValueError): + return False + + # Check if vLLM has the PDL-related modules before doing internet check + has_utils = _spec_exists("vllm.lora.ops.triton_ops.utils") + has_expand_op = _spec_exists("vllm.lora.ops.triton_ops.lora_expand_op") + has_shrink_op = _spec_exists("vllm.lora.ops.triton_ops.lora_shrink_op") + + if not has_utils and not has_expand_op and not has_shrink_op: + # Old vLLM version without PDL support - nothing to patch + return + + # Check if vLLM version includes the fix + VLLM_PDL_FIX_VERSION = "0.13.2" + try: + vllm_version = Version(importlib_version("vllm")) + if vllm_version > Version(VLLM_PDL_FIX_VERSION): + logger.info( + f"Unsloth: SM100 ({sm100_gpu_name}) detected but vLLM {vllm_version} " + f"should include PDL fix - skipping workaround" + ) + return + except Exception as e: + logger.debug( + f"Unsloth: vLLM version check failed ({e}), applying PDL workaround." + ) + + # Apply the PDL fix + os.environ["TRITON_DISABLE_PDL"] = "1" + + def fake_supports_pdl(*args, **kwargs): + return False + + patched = [] + + # First, patch the source module (utils.py) where supports_pdl is defined. + # This is critical because supports_pdl uses @lru_cache - we must clear the + # cache to prevent stale cached results from the original function. + try: + utils_module = importlib.import_module("vllm.lora.ops.triton_ops.utils") + if hasattr(utils_module, "supports_pdl"): + original_fn = utils_module.supports_pdl + if hasattr(original_fn, "cache_clear"): + original_fn.cache_clear() + utils_module.supports_pdl = fake_supports_pdl + patched.append("utils") + except (ImportError, ModuleNotFoundError, AttributeError): + pass + + # Also patch the consumer modules that import supports_pdl from utils. + # This ensures the patched function is used even if the module was already + # imported before this fix runs. + consumer_modules = { + "lora_expand_op": "vllm.lora.ops.triton_ops.lora_expand_op", + "lora_shrink_op": "vllm.lora.ops.triton_ops.lora_shrink_op", + "fused_moe_lora_op": "vllm.lora.ops.triton_ops.fused_moe_lora_op", + } + for name, path in consumer_modules.items(): + try: + module = importlib.import_module(path) + if hasattr(module, "supports_pdl"): + module.supports_pdl = fake_supports_pdl + patched.append(name) + except (ImportError, ModuleNotFoundError, AttributeError): + pass + + if patched: + logger.info( + f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - " + f"patched: {', '.join(patched)}" + ) + else: + # Just set the env var - vLLM might be an older version without supports_pdl + logger.info(f"Unsloth: Set TRITON_DISABLE_PDL=1 for SM100 ({sm100_gpu_name})") diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 912e6f7e3f..fbb14013ff 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -24,7 +24,7 @@ from .utils import ( is_cdna, ) from transformers.models.llama.modeling_llama import logger -from packaging.version import Version +from unsloth_zoo.utils import Version from unsloth_zoo.loss_utils import ( patch_loss_functions as _patch_loss_functions, diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 5952d4af0c..e6c4a12874 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2026.1.1" +__version__ = "2026.1.2" __all__ = [ "SUPPORTS_BFLOAT16", @@ -1197,7 +1197,10 @@ def get_statistics(local_files_only = False): # You can disable this by setting UNSLOTH_DISABLE_STATISTICS import os - if "UNSLOTH_DISABLE_STATISTICS" in os.environ: + if ( + "UNSLOTH_DISABLE_STATISTICS" in os.environ + or os.environ.get("UNSLOTH_USE_MODELSCOPE", "0") == "1" + ): return if local_files_only: return diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py index e9f56763d6..c33317ee02 100644 --- a/unsloth/models/cohere.py +++ b/unsloth/models/cohere.py @@ -15,7 +15,7 @@ from .llama import * from ._utils import __version__ from unsloth_zoo.hf_utils import dtype_from_config -from unsloth_zoo.utils import _get_dtype +from unsloth_zoo.utils import _get_dtype, Version from ..utils.packing import get_packed_info_from_kwargs from ..utils.attention_dispatch import ( AttentionConfig, @@ -35,8 +35,6 @@ try: repeat_kv, ) except: - from packaging.version import Version - transformers_version = Version(transformers_version) if not transformers_version >= Version("4.42"): raise ImportError( diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 291d442673..1789a9cd92 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -14,7 +14,7 @@ from .llama import * from ._utils import __version__ -from unsloth_zoo.utils import _get_dtype +from unsloth_zoo.utils import _get_dtype, Version from unsloth_zoo.hf_utils import dtype_from_config from ..utils.packing import ( build_sdpa_packed_attention_mask, @@ -34,8 +34,6 @@ try: repeat_kv, ) except: - from packaging.version import Version - transformers_version = Version(transformers_version) if not transformers_version >= Version("4.38"): raise ImportError( diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 4b2503b8a1..16d04955d3 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -14,7 +14,7 @@ from .llama import * from ._utils import __version__ -from unsloth_zoo.utils import _get_dtype +from unsloth_zoo.utils import _get_dtype, Version from unsloth_zoo.hf_utils import dtype_from_config from ..utils.packing import get_packed_info_from_kwargs from ..utils.attention_dispatch import ( @@ -41,8 +41,6 @@ try: repeat_kv, ) except: - from packaging.version import Version - transformers_version = Version(transformers_version) if not transformers_version >= Version("4.42"): raise ImportError( diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index f85f1b641f..aae746aed1 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -15,7 +15,7 @@ from .llama import * import os from ._utils import __version__ -from unsloth_zoo.utils import _get_dtype +from unsloth_zoo.utils import _get_dtype, Version from unsloth_zoo.hf_utils import dtype_from_config from ..utils.packing import get_packed_info_from_kwargs from ..utils.attention_dispatch import ( @@ -41,8 +41,6 @@ try: GraniteForCausalLM, ) except: - from packaging.version import Version - transformers_version = Version(transformers_version) if not transformers_version >= Version("4.45.0"): raise ImportError( diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 85332e1116..fe2a89d893 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -28,7 +28,6 @@ from .mapper import ( ) # https://github.com/huggingface/transformers/pull/26037 allows 4 bit loading! -from packaging.version import Version from transformers import __version__ as transformers_version from unsloth.models._utils import TorchAOConfig from unsloth_zoo.utils import Version diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 88aeeda8a1..e945a80354 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -43,10 +43,31 @@ torch_compile_options = { "triton.cudagraphs": False, } -from trl import __version__ as trl_version +# vLLM compatibility shim (TRL expects GuidedDecodingParams even if vLLM doesn't provide it) +try: + import vllm.sampling_params as _unsloth_vllm_sp + + if not hasattr(_unsloth_vllm_sp, "GuidedDecodingParams"): + + class GuidedDecodingParams: + def __init__(self, **kwargs): + self.kwargs = kwargs + + _unsloth_vllm_sp.GuidedDecodingParams = GuidedDecodingParams +except Exception: + pass + +from trl import __version__ as trl_version_raw +from importlib.metadata import version as importlib_version from unsloth_zoo.utils import Version -trl_version = Version(trl_version) +try: + trl_version = Version(trl_version_raw) +except Exception: + try: + trl_version = Version(importlib_version("trl")) + except Exception: + trl_version = Version("0.0.0") def vLLMSamplingParams(**kwargs): @@ -242,12 +263,18 @@ def prepare_for_training_mode(f): @functools.wraps(f) def wrapper(self, *args, **kwargs): # Enable training mode + _was_training = None + if hasattr(self, 'model') and hasattr(self.model, "training"): + _was_training = self.model.training if hasattr(self, 'model') and hasattr(self.model, "for_training"): self.model.for_training() output = f(self, *args, **kwargs) - # Return inference mode + # Restore previous mode when possible if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - self.model.for_inference() + if _was_training is False: + self.model.for_inference() + elif _was_training is True and hasattr(self.model, "for_training"): + self.model.for_training() # Reset gradient checkpointing buffers to free memory while staying ready for next run try: reset_unsloth_gradient_checkpointing_buffers() @@ -332,6 +359,32 @@ pass ''' +def _wrap_grpo_generate_and_score(trainer_cls): + if not hasattr(trainer_cls, "_generate_and_score_completions"): + return + original = trainer_cls._generate_and_score_completions + if getattr(original, "_unsloth_restore_training_wrapped", False): + return + + def wrapped(self, *args, **kwargs): + was_training = getattr(getattr(self, "model", None), "training", None) + try: + return original(self, *args, **kwargs) + finally: + if ( + was_training is False + and hasattr(self, "model") + and hasattr(self.model, "for_inference") + ): + try: + self.model.for_inference() + except Exception: + pass + + wrapped._unsloth_restore_training_wrapped = True + trainer_cls._generate_and_score_completions = wrapped + + def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): # Patch for vLLM and Unsloth PEFT import trl @@ -694,6 +747,19 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): ) RLTrainer_post += training_check + # Sync chat_template from processing_class to vLLM's tokenizer + # This fixes base models that have custom chat templates applied after loading + if "model" in call_args: + vllm_chat_template_sync = ( + "if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'):\n" + " _vllm_tok = self.llm.get_tokenizer()\n" + " _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None)\n" + " if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None:\n" + " _vllm_tok.chat_template = _pc.chat_template\n" + "pass\n" + ) + RLTrainer_post += vllm_chat_template_sync + # Edit optional metrics other_metrics_processor = "" if trainer_file in RL_METRICS_CHANGES: @@ -1059,6 +1125,16 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): globals(), ) + if trainer_file == "grpo_trainer": + try: + _wrap_grpo_generate_and_score( + getattr(created_module, f"Unsloth{RLTrainer_name}") + ) + except Exception as e: + logger.info( + f"Unsloth: Could not wrap _generate_and_score_completions for {RLTrainer_name}: {e}" + ) + def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports): init = inspect.getsource(RLTrainer.__init__) diff --git a/unsloth/trainer.py b/unsloth/trainer.py index 0d98cff305..858dcf2cd3 100644 --- a/unsloth/trainer.py +++ b/unsloth/trainer.py @@ -211,7 +211,7 @@ def _backwards_compatible_trainer(trainer_class, config_class): if "processing_class" in trainer_params and "tokenizer" in kwargs: kwargs["processing_class"] = kwargs.pop("tokenizer") - if ("args" in kwargs) and (Version(trl.__version__) >= Version("0.13.0.dev0")): + if ("args" in kwargs) and (Version(trl) >= Version("0.13.0.dev0")): training_args = kwargs.pop("args", None) # Get parameters that Trainer.__init__ actually expects @@ -412,7 +412,7 @@ def _patch_trl_trainer(): if hasattr(trl, "__UNSLOTH_BACKWARDS_COMPATIBLE__"): return - if Version(trl.__version__) <= Version("0.11.0"): + if Version(trl) <= Version("0.11.0"): return import trl.trainer diff --git a/unsloth/utils/attention_dispatch.py b/unsloth/utils/attention_dispatch.py index 0e5f3c1951..a7620549be 100644 --- a/unsloth/utils/attention_dispatch.py +++ b/unsloth/utils/attention_dispatch.py @@ -32,9 +32,6 @@ from ..utils.packing import ( if HAS_FLASH_ATTENTION: from flash_attn import flash_attn_func, flash_attn_varlen_func HAS_XFORMERS = xformers is not None -BlockDiagonalCausalMask = None -if HAS_XFORMERS: - BlockDiagonalCausalMask = xformers.attn_bias.BlockDiagonalCausalMask SDPA_HAS_GQA = "enable_gqa" in (scaled_dot_product_attention.__doc__ or "") FLASH_VARLEN = "flash_varlen"