diff --git a/tests/test_issue_5344_guardrail.py b/tests/test_issue_5344_guardrail.py new file mode 100644 index 0000000000..09599bb071 --- /dev/null +++ b/tests/test_issue_5344_guardrail.py @@ -0,0 +1,177 @@ +"""Unit tests for the unslothai/unsloth#5344 silent-quantization-bypass guardrail. + +Covers two failure modes the helper detects: + 1. total bypass: load_in_4bit was requested but no bnb modules exist. + 2. partial bypass: bnb quantized nn.Linear but a large fraction of weight + bytes live in non-nn.Linear Parameters (e.g. Gemma-4 MoE fused experts). +""" + +import warnings + +import torch +import torch.nn as nn + + +# unsloth must be imported before transformers per its loading order, but +# these tests do not exercise the real loader. Import the helper directly. +from unsloth.models.vision import _warn_if_quantization_silently_dropped + + +class _PretendLinear4bit(nn.Module): + """type(m).__name__ == 'Linear4bit' so the guardrail counts it as quantized.""" + + def __init__(self): + super().__init__() + self.weight = nn.Parameter( + torch.zeros(1, dtype = torch.uint8), + requires_grad = False, + ) + + +_PretendLinear4bit.__name__ = "Linear4bit" + + +def _unquantized_model(): + return nn.Sequential(nn.Linear(4, 4), nn.Linear(4, 4)) + + +def _quantized_model(): + return nn.Sequential(nn.Linear(4, 4), _PretendLinear4bit()) + + +def test_fires_when_4bit_requested_but_no_bnb_modules(): + model = _unquantized_model() + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = True, + load_in_8bit = False, + full_finetuning = False, + ) + msgs = [str(w.message) for w in caught] + assert any("load_in_4bit=True was requested" in m for m in msgs), msgs + assert any("issues/5344" in m for m in msgs), msgs + + +def test_silent_when_4bit_succeeded(): + model = _quantized_model() + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = True, + load_in_8bit = False, + full_finetuning = False, + ) + msgs = [str(w.message) for w in caught] + assert not any("load_in_4bit" in m for m in msgs), msgs + + +def test_silent_for_full_finetuning(): + model = _unquantized_model() + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = False, + load_in_8bit = False, + full_finetuning = True, + ) + msgs = [str(w.message) for w in caught] + assert not any("load_in_4bit" in m or "load_in_8bit" in m for m in msgs), msgs + + +def test_silent_when_no_quantization_requested(): + model = _unquantized_model() + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = False, + load_in_8bit = False, + full_finetuning = False, + ) + msgs = [str(w.message) for w in caught] + assert not any("load_in_4bit" in m or "load_in_8bit" in m for m in msgs), msgs + + +def test_fires_for_8bit_silent_bypass(): + model = _unquantized_model() + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = False, + load_in_8bit = True, + full_finetuning = False, + ) + msgs = [str(w.message) for w in caught] + assert any("load_in_8bit=True was requested" in m for m in msgs), msgs + + +class _MoEFusedExpertWrapper(nn.Module): + """Mimics Gemma4TextExperts: fused 3D weights stored as nn.Parameter, not + as separate nn.Linear instances. bnb's replace_with_bnb_linear skips this.""" + + def __init__(self, num_experts = 128, intermediate = 1408, hidden = 2816): + super().__init__() + self.gate_up_proj = nn.Parameter( + torch.zeros((num_experts, intermediate, hidden), dtype = torch.bfloat16), + requires_grad = False, + ) + + +def _partial_quant_model(): + return nn.Sequential(_PretendLinear4bit(), _MoEFusedExpertWrapper()) + + +def test_fires_on_partial_quant_moe_experts(): + model = _partial_quant_model() + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = True, + load_in_8bit = False, + full_finetuning = False, + ) + msgs = [str(w.message) for w in caught] + assert any("partially applied" in m for m in msgs), msgs + assert any("gate_up_proj" in m for m in msgs), msgs + + +class _NormParam(nn.Module): + """An RMSNorm-like module: large BF16 weight whose name is in the skip list.""" + + def __init__(self, dim = 8 * 1024 * 1024 + 10): + super().__init__() + self.norm_weight = nn.Parameter( + torch.zeros(dim, dtype = torch.bfloat16), + requires_grad = False, + ) + + +def test_silent_when_only_skip_list_tensors_unquantized(): + model = nn.Sequential(_PretendLinear4bit(), _NormParam()) + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = True, + load_in_8bit = False, + full_finetuning = False, + ) + msgs = [str(w.message) for w in caught] + assert not any("partially applied" in m for m in msgs), msgs + + +if __name__ == "__main__": + test_fires_when_4bit_requested_but_no_bnb_modules() + test_silent_when_4bit_succeeded() + test_silent_for_full_finetuning() + test_silent_when_no_quantization_requested() + test_fires_for_8bit_silent_bypass() + test_fires_on_partial_quant_moe_experts() + test_silent_when_only_skip_list_tensors_unquantized() + print("All 7 guardrail tests passed.") diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d39f2588ef..dc1cf17482 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2478,6 +2478,16 @@ class FastLlamaModel: and not _head.weight.is_floating_point() ): _head.to(dtype) + # Guardrail: warn before dispatch hooks if quantization was silently dropped. + from unsloth.models.vision import _warn_if_quantization_silently_dropped + + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = load_in_4bit, + load_in_8bit = kwargs.get("load_in_8bit", False), + full_finetuning = kwargs.get("full_finetuning", False), + quantization_config = kwargs.get("quantization_config"), + ) # Attach dispatch hooks for bnb multi-device loads. from unsloth.models.vision import _attach_bnb_multidevice_hooks @@ -2500,9 +2510,19 @@ class FastLlamaModel: attn_implementation = preferred_attn_impl, **kwargs, ) - # Attach dispatch hooks for bnb multi-device loads. - from unsloth.models.vision import _attach_bnb_multidevice_hooks + # Guardrail (#5344) + multi-device dispatch hooks share an import. + from unsloth.models.vision import ( + _warn_if_quantization_silently_dropped, + _attach_bnb_multidevice_hooks, + ) + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = load_in_4bit, + load_in_8bit = kwargs.get("load_in_8bit", False), + full_finetuning = kwargs.get("full_finetuning", False), + quantization_config = kwargs.get("quantization_config"), + ) _attach_bnb_multidevice_hooks( model, load_in_4bit = load_in_4bit, diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index df371e00c8..c7aeb338f3 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -68,7 +68,6 @@ from unsloth_zoo.patching_utils import patch_model_and_tokenizer from unsloth_zoo.training_utils import prepare_model_for_training from unsloth_zoo.utils import Version -from transformers import __version__ as transformers_version import types import functools @@ -99,6 +98,140 @@ __all__ = [ ] +# bnb-quantized Linear class names (see unslothai/unsloth#5344 guardrail). +_BNB_QUANT_CLASS_NAMES = ("Linear4bit", "Linear8bitLt", "LinearNF4", "LinearFP4") + +# Substrings the guardrail treats as intentionally-not-quantized: embeddings, +# norms, biases, routers/gates that need fp16/fp32 precision, vision/audio +# towers, classification heads, rotary tables. +_GUARDRAIL_SKIP_PATTERNS = ( + "embed", + "embedding", + "norm", + "ln_", + "rms", + ".bias", + "lm_head", + "multi_modal_projector", + "merger", + "modality_projection", + "router", + "block_sparse_moe.gate", + "mamba", + "audio_tower", + "vision_tower", + "score", + "classifier", + "qa_outputs", + "rotary", +) + +# A floating Parameter larger than this, found outside the skip list, counts +# as bulk weight that should have been 4-bit. Tuned so head dims and small +# projections do not false-fire but MoE fused expert tensors do. +_GUARDRAIL_BULK_WEIGHT_NUMEL = 8 * 1024 * 1024 + + +def _warn_if_quantization_silently_dropped( + model, + load_in_4bit, + load_in_8bit, + full_finetuning, + quantization_config = None, +): + """Guardrail for unslothai/unsloth#5344. + + Two failure modes covered: + + 1. TOTAL bypass: load_in_4bit was requested but the model contains zero + bnb Linear4bit / Linear8bitLt modules. transformers / bnb / a backend + incompatibility dropped kwargs.quantization_config. + + 2. PARTIAL bypass: bnb quantized the nn.Linear modules but a large fraction + of weight bytes live in non-nn.Linear Parameters (e.g. Gemma-4 MoE fused + 3D expert tensors, custom Linear-like wrappers). bnb only swaps nn.Linear + instances; fused expert weights stay in BF16, defeating QLoRA savings + even though some Linear4bit modules exist. + + Warns rather than raises so non-bnb backends (CPU / MLX / AMD-without-bnb) + with legitimately no Linear4bit are not broken. + """ + if full_finetuning: + return + if quantization_config is not None: + if isinstance(quantization_config, dict): + load_in_4bit = load_in_4bit or bool(quantization_config.get("load_in_4bit")) + load_in_8bit = load_in_8bit or bool(quantization_config.get("load_in_8bit")) + else: + load_in_4bit = load_in_4bit or bool( + getattr(quantization_config, "load_in_4bit", False) + ) + load_in_8bit = load_in_8bit or bool( + getattr(quantization_config, "load_in_8bit", False) + ) + if not (load_in_4bit or load_in_8bit): + return + + has_bnb = any(type(m).__name__ in _BNB_QUANT_CLASS_NAMES for m in model.modules()) + + # Failure mode 1: total bypass. + if not has_bnb: + kind = "4bit" if load_in_4bit else "8bit" + bnb_class_name = "Linear4bit" if load_in_4bit else "Linear8bitLt" + warnings.warn( + f"Unsloth: load_in_{kind}=True was requested but no bitsandbytes " + f"{bnb_class_name} modules were produced. The runtime quantization " + f"config was silently dropped and the model is in full precision. " + f"See https://github.com/unslothai/unsloth/issues/5344 for known " + f"triggers (transformers/bnb version mismatch, MoE checkpoints " + f"without a -bnb-4bit sibling, multi-GPU dispatch). Workaround: " + f'pass device_map="cuda:0" and pin transformers/bitsandbytes to ' + f"a version known to work.", + stacklevel = 3, + ) + return + + # Failure mode 2: partial bypass. Walk named_parameters and find large + # floating tensors outside the skip list. If they aggregate to >= 2x the + # quantized payload, partial quant is essentially negating 4-bit savings. + quantized_bytes = 0 + suspect_bytes = 0 + suspect_samples = [] + for name, p in model.named_parameters(): + if p is None: + continue + nbytes = p.numel() * p.element_size() + if p.dtype in (torch.uint8, torch.int8): + # bnb stores 4-bit payloads as uint8 and 8-bit payloads (Int8Params) as int8. + quantized_bytes += nbytes + continue + if p.dtype not in (torch.bfloat16, torch.float16, torch.float32): + continue + if p.numel() < _GUARDRAIL_BULK_WEIGHT_NUMEL: + continue + lname = name.lower() + if any(pat in lname for pat in _GUARDRAIL_SKIP_PATTERNS): + continue + suspect_bytes += nbytes + if len(suspect_samples) < 3: + suspect_samples.append((name, str(p.dtype), tuple(p.shape))) + + if quantized_bytes > 0 and suspect_bytes >= 2 * quantized_bytes: + kind = "4bit" if load_in_4bit else "8bit" + suspect_human = ", ".join(f"{n} ({d}, {s})" for n, d, s in suspect_samples) + warnings.warn( + f"Unsloth: load_in_{kind}=True is partially applied. " + f"bitsandbytes quantized ~{quantized_bytes/1024**3:.2f} GB of " + f"nn.Linear weights, but ~{suspect_bytes/1024**3:.2f} GB of " + f"non-nn.Linear floating Parameters were left unquantized (e.g. " + f"fused MoE expert tensors, custom Linear-like wrappers). " + f"Examples: {suspect_human}. The model's effective VRAM " + f"footprint is close to its full-precision size. See " + f"https://github.com/unslothai/unsloth/issues/5344.", + stacklevel = 3, + ) + + def _infer_device_map_from_loaded_model(model): """Build a compact device_map by inspecting actual parameter placements.""" device_map = {} @@ -935,6 +1068,14 @@ class FastBaseModel: # attn_implementation = attn_implementation, **kwargs, ) + # Guardrail: see _warn_if_quantization_silently_dropped + #5344. + _warn_if_quantization_silently_dropped( + model, + load_in_4bit = load_in_4bit, + load_in_8bit = load_in_8bit, + full_finetuning = full_finetuning, + quantization_config = kwargs.get("quantization_config"), + ) # Attach dispatch hooks for bnb multi-device loads. _attach_bnb_multidevice_hooks( model,