ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix Nemotron-H and Nemotron-VL model support (#4105)

Browse source 
* 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>
This commit is contained in:
Daniel Han 2026-02-25 03:14:12 -08:00 • committed by GitHub
commit 78963ca19c
2 changed files with 33 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

25
unsloth/models/loader.py
View file

@ -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

11
unsloth/models/vision.py
View file

@ -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