From 78963ca19cfef544e23a07417f87226f20b04f1f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Feb 2026 03:14:12 -0800 Subject: [PATCH] Fix Nemotron-H and Nemotron-VL model support (#4105) * Fix Nemotron-H and Nemotron-VL model support - Add Mamba kernel precision settings for Nemotron-H hybrid models - Fix VL model auto_model selection for models that only register AutoModelForCausalLM in their auto_map - Skip quantization of out_proj for Nemotron-H Mamba layers * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Simplify VLM auto_model selection logic Reduce three branches to two since the first and third both assign AutoModelForVision2Seq. The simplified condition checks whether the auto_map exclusively registers AutoModelForCausalLM without the VLM class, and defaults to AutoModelForVision2Seq otherwise. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/loader.py | 25 ++++++++++++++++++++++++- unsloth/models/vision.py | 11 +++++++++-- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 4054f1b7f5..04c15bb2c3 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -1084,6 +1084,15 @@ class FastModel(FastBaseModel): ";" "os.environ['TRITON_F32_DEFAULT'] = 'ieee'" ) + elif "nemotron_h" in model_types_all: + # NemotronH (hybrid Mamba-2 + Transformer) uses same Mamba kernels as Falcon-H1 + # Mamba kernels need float32 Triton precision + os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = ( + "float16;torch.float32;torch.float16;" + "if name.endswith(('q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj', 'head')): module.to(torch.float16)" + ";" + "os.environ['TRITON_F32_DEFAULT'] = 'ieee'" + ) elif "gpt_oss" in model_types_all: os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" if not load_in_4bit: @@ -1281,7 +1290,21 @@ class FastModel(FastBaseModel): is_vlm = any(x.endswith("ForConditionalGeneration") for x in architectures) is_vlm = is_vlm or hasattr(model_config, "vision_config") if auto_model is None: - auto_model = AutoModelForVision2Seq if is_vlm else AutoModelForCausalLM + if is_vlm: + # Check if the model's auto_map supports the VLM auto class. + # Some VL models (e.g. Nemotron-VL) only register AutoModelForCausalLM + # in their auto_map, not AutoModelForImageTextToText/AutoModelForVision2Seq. + _auto_map = getattr(model_config, "auto_map", {}) or {} + _vlm_class_name = AutoModelForVision2Seq.__name__ + if ( + "AutoModelForCausalLM" in _auto_map + and _vlm_class_name not in _auto_map + ): + auto_model = AutoModelForCausalLM + else: + auto_model = AutoModelForVision2Seq + else: + auto_model = AutoModelForCausalLM load_in_4bit_kwargs = load_in_4bit load_in_8bit_kwargs = load_in_8bit diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 56cd615a49..24e72e1535 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -654,18 +654,25 @@ class FastBaseModel: raise RuntimeError( "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!" ) + _skip_modules = SKIP_QUANTIZATION_MODULES.copy() + # Nemotron-H uses 'mixer' (not 'mamba') for Mamba layers. + # Mamba fused kernels pass out_proj.weight directly to F.linear, + # which fails with quantized Params4bit. Skip out_proj from quantization. + if any(mt == "nemotron_h" for mt in (model_types or [])): + _skip_modules.append("out_proj") + if load_in_4bit: bnb_config = BitsAndBytesConfig( load_in_4bit = True, bnb_4bit_use_double_quant = True, bnb_4bit_quant_type = "nf4", bnb_4bit_compute_dtype = bnb_compute_dtype, - llm_int8_skip_modules = SKIP_QUANTIZATION_MODULES.copy(), + llm_int8_skip_modules = _skip_modules, ) elif load_in_8bit: bnb_config = BitsAndBytesConfig( load_in_8bit = True, - llm_int8_skip_modules = SKIP_QUANTIZATION_MODULES.copy(), + llm_int8_skip_modules = _skip_modules, ) elif load_in_16bit: bnb_config = None