From 6f36e8403a0830b311059006bcce3aad95f77229 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Mon, 8 Jun 2026 23:23:53 +0530 Subject: [PATCH] Merge nvfp4_load CI fixes Merged latest main, resolved KTO test conflicts, fixed nvfp4 test to use synthetic configs, fixed TRL/GRPO KTO drift --- tests/test_nvfp4_quant_load.py | 143 +++++++++++++++++++++++++++++++++ unsloth/models/llama.py | 16 +++- unsloth/models/loader_utils.py | 43 +++++++++- unsloth/models/vision.py | 8 ++ 4 files changed, 207 insertions(+), 3 deletions(-) create mode 100644 tests/test_nvfp4_quant_load.py diff --git a/tests/test_nvfp4_quant_load.py b/tests/test_nvfp4_quant_load.py new file mode 100644 index 0000000000..988fae5c97 --- /dev/null +++ b/tests/test_nvfp4_quant_load.py @@ -0,0 +1,143 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Test cases for NVFP4 / compressed-tensors model loading. +Ensures that models with non-bitsandbytes quantization configs +don't conflict with Unsloth's default load_in_4bit=True behavior. + +Uses synthetic config objects (no network access) so this suite +runs offline in CI where tests/security/conftest.py blocks socket +connections. +""" + +from types import SimpleNamespace + +# Import unsloth first to set UNSLOTH_IS_PRESENT env var +import unsloth +from unsloth_zoo.utils import get_quant_type +from unsloth.models.loader_utils import check_and_disable_bitsandbytes_loading + + +def _make_config(quantization_config = None, model_type = "llama"): + return SimpleNamespace( + quantization_config = quantization_config, + model_type = model_type, + ) + + +_NVFP4_QCFG_DICT = { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "quantization_config": {"num_bits": 4}, +} + +_BNB_QCFG_DICT = { + "quant_method": "bitsandbytes", + "load_in_4bit": True, + "bnb_4bit_compute_dtype": "float16", + "llm_int8_skip_modules": [], +} + + +def test_nvfp4_config_has_compressed_tensors(): + config = _make_config(quantization_config = _NVFP4_QCFG_DICT) + qcfg = config.quantization_config + assert qcfg is not None + assert qcfg.get("quant_method") == "compressed-tensors" + assert qcfg.get("format") == "nvfp4-pack-quantized" + + +def test_regular_bnb_config_has_bitsandbytes(): + config = _make_config(quantization_config = _BNB_QCFG_DICT) + qcfg = config.quantization_config + assert qcfg is not None + assert qcfg.get("quant_method") == "bitsandbytes" + + +def test_nvfp4_disables_load_in_4bit(): + config = _make_config(quantization_config = _NVFP4_QCFG_DICT) + quant_method = get_quant_type(config) + assert quant_method == "compressed-tensors" + + load_in_4bit, load_in_8bit, _ = check_and_disable_bitsandbytes_loading( + config, load_in_4bit = True, load_in_8bit = False, verbose = False + ) + assert load_in_4bit is False + assert load_in_8bit is False + + +def test_bnb_does_not_disable_load_in_4bit(): + config = _make_config(quantization_config = _BNB_QCFG_DICT) + quant_method = get_quant_type(config) + assert quant_method == "bitsandbytes" + + load_in_4bit, load_in_8bit, _ = check_and_disable_bitsandbytes_loading( + config, load_in_4bit = True, load_in_8bit = False, verbose = False + ) + assert load_in_4bit is True + assert load_in_8bit is False + + +def test_no_quantization_config_leaves_settings_unchanged(): + config = _make_config(quantization_config = None) + quant_method = get_quant_type(config) + assert quant_method is None + + load_in_4bit, load_in_8bit, _ = check_and_disable_bitsandbytes_loading( + config, load_in_4bit = True, load_in_8bit = False, verbose = False + ) + assert load_in_4bit is True + assert load_in_8bit is False + + +def test_nvfp4_disables_both_4bit_and_8bit(): + config = _make_config(quantization_config = _NVFP4_QCFG_DICT) + + load_in_4bit, load_in_8bit, _ = check_and_disable_bitsandbytes_loading( + config, load_in_4bit = True, load_in_8bit = True, verbose = False + ) + assert load_in_4bit is False + assert load_in_8bit is False + + +def test_verbose_flag_does_not_raise(): + config = _make_config(quantization_config = _NVFP4_QCFG_DICT) + load_in_4bit, load_in_8bit, _ = check_and_disable_bitsandbytes_loading( + config, load_in_4bit = True, load_in_8bit = False, verbose = True + ) + assert load_in_4bit is False + assert load_in_8bit is False + + +def test_empty_quantization_config_is_not_quantized(): + config = _make_config(quantization_config = {}) + assert get_quant_type(config) is None + + load_in_4bit, load_in_8bit, _ = check_and_disable_bitsandbytes_loading( + config, load_in_4bit = True, load_in_8bit = False, verbose = False + ) + assert load_in_4bit is True + + +if __name__ == "__main__": + test_nvfp4_config_has_compressed_tensors() + test_regular_bnb_config_has_bitsandbytes() + test_nvfp4_disables_load_in_4bit() + test_bnb_does_not_disable_load_in_4bit() + test_no_quantization_config_leaves_settings_unchanged() + test_nvfp4_disables_both_4bit_and_8bit() + test_verbose_flag_does_not_raise() + test_empty_quantization_config_is_not_quantized() + print("All tests passed!") diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 945d827573..17d485f9f9 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2352,7 +2352,20 @@ class FastLlamaModel: # Add to kwargs kwargs["rope_scaling"] = rope_scaling + from .loader_utils import check_and_disable_bitsandbytes_loading + from unsloth_zoo.utils import get_quant_type + + # Extract load_in_8bit from kwargs if provided + load_in_8bit = kwargs.get("load_in_8bit", False) + + # Check and disable bitsandbytes loading if model has non-bitsandbytes quantization + load_in_4bit, load_in_8bit, _ckpt_quant_method = check_and_disable_bitsandbytes_loading( + model_config, load_in_4bit = load_in_4bit, load_in_8bit = load_in_8bit + ) + bnb_config = None + _ckpt_qcfg = getattr(model_config, "quantization_config", None) + if load_in_4bit: llm_int8_skip_modules = SKIP_QUANTIZATION_MODULES.copy() if IS_FALCON_H1: @@ -2371,8 +2384,7 @@ class FastLlamaModel: # we pass via kwargs. Merge our skip list into that bundled config # so task heads like `score` (for *ForSequenceClassification) stay # in the compute dtype. See unslothai/unsloth#5027. - _ckpt_qcfg = getattr(model_config, "quantization_config", None) - if _ckpt_qcfg is not None: + if _ckpt_quant_method == "bitsandbytes" and _ckpt_qcfg is not None: if isinstance(_ckpt_qcfg, dict): _ckpt_skip = list(_ckpt_qcfg.get("llm_int8_skip_modules") or []) for _m in llm_int8_skip_modules: diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index ea18dc125a..d304bb454e 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -30,7 +30,7 @@ from .mapper import ( # https://github.com/huggingface/transformers/pull/26037 allows 4 bit loading! from transformers import __version__ as transformers_version from unsloth.models._utils import TorchAOConfig -from unsloth_zoo.utils import Version +from unsloth_zoo.utils import Version, get_quant_type import gc transformers_version = Version(transformers_version) @@ -335,6 +335,47 @@ def _tag_model_with_fp8_torchao_config(model: torch.nn.Module, fp8_mode: str): pass +def check_and_disable_bitsandbytes_loading( + model_config, + load_in_4bit = True, + load_in_8bit = False, + verbose = True, +): + """ + Check if we should disable bitsandbytes loading (load_in_4bit/load_in_8bit) + because the model already has a non-bitsandbytes quantization config. + If so, disable BOTH 4bit and 8bit loading and print a warning message. + + Args: + model_config: The AutoConfig object from the model + load_in_4bit: Whether load_in_4bit is currently enabled + load_in_8bit: Whether load_in_8bit is currently enabled + verbose: Whether to print warning messages + + Returns: + tuple: (load_in_4bit, load_in_8bit, quant_method) + load_in_4bit/load_in_8bit will be False if they were disabled + quant_method is the detected quantization method or None + """ + quant_method = get_quant_type(model_config) + + if quant_method is None or quant_method == "bitsandbytes": + return load_in_4bit, load_in_8bit, quant_method + + # Model has a non-bitsandbytes quantization config (e.g., compressed-tensors, gptq, awq) + # We should disable BOTH bitsandbytes loading to avoid config conflicts + if load_in_4bit or load_in_8bit: + if verbose: + print( + f"Unsloth: Model already quantized with {quant_method}. " + f"Disabling `load_in_4bit` and `load_in_8bit` to avoid quantization config conflict." + ) + load_in_4bit = False + load_in_8bit = False + + return load_in_4bit, load_in_8bit, quant_method + + def _get_fp8_mode_and_check_settings( load_in_fp8: Union[bool, str], fast_inference: bool, diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index b308a6f1da..3e779a46b1 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -770,6 +770,14 @@ class FastBaseModel: bnb_config = None user_quantization_config = kwargs.get("quantization_config", None) + + # Check if model already has a non-bitsandbytes quantization config (e.g. compressed-tensors/NVFP4) + from .loader_utils import check_and_disable_bitsandbytes_loading + + load_in_4bit, load_in_8bit, _ = check_and_disable_bitsandbytes_loading( + auto_config, load_in_4bit = load_in_4bit, load_in_8bit = load_in_8bit + ) + if full_finetuning and (load_in_4bit or load_in_8bit): print( "Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA."