diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index e6178e60f..984057e9f 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -26,6 +26,7 @@ already_imported = [mod for mod in critical_modules if mod in sys.modules] # Fix some issues before importing other packages from .import_fixes import ( fix_message_factory_issue, + fix_torch_check_is_size, check_fbgemm_gpu_version, disable_broken_causal_conv1d, disable_broken_vllm, @@ -72,6 +73,7 @@ fix_bitsandbytes_rocm_arch_detection() disable_broken_causal_conv1d() disable_broken_vllm() fix_message_factory_issue() +fix_torch_check_is_size() check_fbgemm_gpu_version() torchvision_compatibility_check() fix_diffusers_warnings() @@ -81,6 +83,7 @@ del fix_bitsandbytes_rocm_arch_detection del disable_broken_causal_conv1d del disable_broken_vllm del fix_message_factory_issue +del fix_torch_check_is_size del check_fbgemm_gpu_version del torchvision_compatibility_check del fix_diffusers_warnings diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index c5300ed4d..09de248c7 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -172,6 +172,10 @@ if not UNSLOTH_ENABLE_LOGGING: # Deprecation warnings from torchao warnings.filterwarnings("ignore", message = "`int4_weight_only` is deprecated") warnings.filterwarnings("ignore", message = "`int8_weight_only` is deprecated") + # torch._check_is_size FutureWarning (called by bitsandbytes 4-bit dequant) + warnings.filterwarnings( + "ignore", message = r"_check_is_size will be removed", category = FutureWarning + ) # TorchAO deprecated import paths (https://github.com/pytorch/ao/issues/2752) warnings.filterwarnings( @@ -253,6 +257,30 @@ if not UNSLOTH_ENABLE_LOGGING: ) +def fix_torch_check_is_size(): + """Shim torch._check_is_size if a future torch removes it (bitsandbytes 4-bit + dequant calls it). The FutureWarning is silenced in suppress_cuda_printf.""" + try: + import torch + + if hasattr(torch, "_check_is_size"): + return + + def _check_is_size( + i, + message = None, + *, + max = None, + ): + torch._check(i >= 0, message) + if max is not None: + torch._check(i <= max, message) + + torch._check_is_size = _check_is_size + except Exception: + return + + # Fix up AttributeError: 'MessageFactory' object has no attribute 'GetPrototype' # MUST do this at the start primarily due to tensorflow causing issues def fix_message_factory_issue():