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

Merge branch 'main' into pr/3719

Browse source
This commit is contained in:
Daniel Han 2026-01-04 05:50:05 -08:00
commit df569dab89
4 changed files with 156 additions and 50 deletions
Download patch file Download diff file Expand all files Collapse all files

164
unsloth/models/loader.py
View file

@ -151,8 +151,41 @@ class FastLanguageModel(FastLlamaModel):
*args,
**kwargs,
):
# Respect user-provided quantization_config (e.g. BitsAndBytesConfig)
quantization_config = kwargs.get("quantization_config", None)
if quantization_config is not None:
if isinstance(quantization_config, dict):
q_load_in_4bit = quantization_config.get("load_in_4bit", False)
q_load_in_8bit = quantization_config.get("load_in_8bit", False)
else:
q_load_in_4bit = getattr(quantization_config, "load_in_4bit", False)
q_load_in_8bit = getattr(quantization_config, "load_in_8bit", False)
if q_load_in_4bit:
load_in_4bit = True
load_in_8bit = False
if q_load_in_8bit:
load_in_8bit = True
load_in_4bit = False
# Login to allow private models
token = hf_login(token)
# Align dtype with bnb_4bit_compute_dtype if provided and dtype is unset.
if dtype is None and quantization_config is not None:
bnb_compute_dtype = None
if isinstance(quantization_config, dict):
if quantization_config.get("load_in_4bit", False):
bnb_compute_dtype = quantization_config.get(
"bnb_4bit_compute_dtype", None
)
else:
if getattr(quantization_config, "load_in_4bit", False):
bnb_compute_dtype = getattr(
quantization_config, "bnb_4bit_compute_dtype", None
)
if isinstance(bnb_compute_dtype, str):
bnb_compute_dtype = getattr(torch, bnb_compute_dtype, None)
if isinstance(bnb_compute_dtype, torch.dtype):
dtype = bnb_compute_dtype
if load_in_8bit or full_finetuning or qat_scheme is not None:
return FastModel.from_pretrained(
model_name = model_name,
@ -542,11 +575,17 @@ class FastLanguageModel(FastLlamaModel):
if fast_inference:
fast_inference, model_name = fast_inference_setup(model_name, model_config)
load_in_4bit_kwargs = load_in_4bit
load_in_8bit_kwargs = load_in_8bit
if quantization_config is not None and not fast_inference:
load_in_4bit_kwargs = False
load_in_8bit_kwargs = False
model, tokenizer = dispatch_model.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = _get_dtype(dtype),
load_in_4bit = load_in_4bit,
load_in_4bit = load_in_4bit_kwargs,
token = token,
device_map = device_map,
rope_scaling = rope_scaling,
@ -583,22 +622,30 @@ class FastLanguageModel(FastLlamaModel):
)
if load_in_4bit:
# Fix up bitsandbytes config
compute_dtype = dtype_from_config(model.config)
quantization_config = {
# Sometimes compute_dtype is not a string!!
"bnb_4bit_compute_dtype": compute_dtype,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_use_double_quant": True,
"llm_int8_enable_fp32_cpu_offload": False,
"llm_int8_has_fp16_weight": False,
"llm_int8_skip_modules": None,
"llm_int8_threshold": 6.0,
"load_in_4bit": True,
"load_in_8bit": False,
"quant_method": "bitsandbytes",
}
model.config.update({"quantization_config": quantization_config})
# Fix up bitsandbytes config, but respect user-provided quantization_config
if quantization_config is None:
compute_dtype = dtype_from_config(model.config)
quantization_config = {
# Sometimes compute_dtype is not a string!!
"bnb_4bit_compute_dtype": compute_dtype,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_use_double_quant": True,
"llm_int8_enable_fp32_cpu_offload": False,
"llm_int8_has_fp16_weight": False,
"llm_int8_skip_modules": None,
"llm_int8_threshold": 6.0,
"load_in_4bit": True,
"load_in_8bit": False,
"quant_method": "bitsandbytes",
}
model.config.update({"quantization_config": quantization_config})
else:
if hasattr(quantization_config, "to_dict"):
model.config.update(
{"quantization_config": quantization_config.to_dict()}
)
elif isinstance(quantization_config, dict):
model.config.update({"quantization_config": quantization_config})
if load_in_fp8 != False:
_tag_model_with_fp8_torchao_config(model, fp8_mode)
@ -690,12 +737,45 @@ class FastModel(FastBaseModel):
*args,
**kwargs,
):
# Respect user-provided quantization_config (e.g. BitsAndBytesConfig)
quantization_config = kwargs.get("quantization_config", None)
if quantization_config is not None:
if isinstance(quantization_config, dict):
q_load_in_4bit = quantization_config.get("load_in_4bit", False)
q_load_in_8bit = quantization_config.get("load_in_8bit", False)
else:
q_load_in_4bit = getattr(quantization_config, "load_in_4bit", False)
q_load_in_8bit = getattr(quantization_config, "load_in_8bit", False)
if q_load_in_4bit:
load_in_4bit = True
load_in_8bit = False
if q_load_in_8bit:
load_in_8bit = True
load_in_4bit = False
# Login to allow private models
token = hf_login(token)
if whisper_language is not None:
assert type(whisper_language) is str
if whisper_task is not None:
assert type(whisper_task) is str
# Align dtype with bnb_4bit_compute_dtype if provided and dtype is unset.
if dtype is None and quantization_config is not None:
bnb_compute_dtype = None
if isinstance(quantization_config, dict):
if quantization_config.get("load_in_4bit", False):
bnb_compute_dtype = quantization_config.get(
"bnb_4bit_compute_dtype", None
)
else:
if getattr(quantization_config, "load_in_4bit", False):
bnb_compute_dtype = getattr(
quantization_config, "bnb_4bit_compute_dtype", None
)
if isinstance(bnb_compute_dtype, str):
bnb_compute_dtype = getattr(torch, bnb_compute_dtype, None)
if isinstance(bnb_compute_dtype, torch.dtype):
dtype = bnb_compute_dtype
SUPPORTS_BFLOAT16 = is_bfloat16_supported()
if dtype is None:
dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16
@ -1169,12 +1249,18 @@ class FastModel(FastBaseModel):
if auto_model is None:
auto_model = AutoModelForVision2Seq if is_vlm else AutoModelForCausalLM
load_in_4bit_kwargs = load_in_4bit
load_in_8bit_kwargs = load_in_8bit
if quantization_config is not None and not fast_inference:
load_in_4bit_kwargs = False
load_in_8bit_kwargs = False
model, tokenizer = FastBaseModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = _get_dtype(dtype),
load_in_4bit = load_in_4bit,
load_in_8bit = load_in_8bit,
load_in_4bit = load_in_4bit_kwargs,
load_in_8bit = load_in_8bit_kwargs,
load_in_16bit = load_in_16bit,
full_finetuning = full_finetuning,
token = token,
@ -1220,22 +1306,30 @@ class FastModel(FastBaseModel):
)
if load_in_4bit:
# Fix up bitsandbytes config
compute_dtype = dtype_from_config(model.config)
quantization_config = {
# Sometimes compute_dtype is not a string!!
"bnb_4bit_compute_dtype": compute_dtype,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_use_double_quant": True,
"llm_int8_enable_fp32_cpu_offload": False,
"llm_int8_has_fp16_weight": False,
"llm_int8_skip_modules": None,
"llm_int8_threshold": 6.0,
"load_in_4bit": True,
"load_in_8bit": False,
"quant_method": "bitsandbytes",
}
model.config.update({"quantization_config": quantization_config})
# Fix up bitsandbytes config, but respect user-provided quantization_config
if quantization_config is None:
compute_dtype = dtype_from_config(model.config)
quantization_config = {
# Sometimes compute_dtype is not a string!!
"bnb_4bit_compute_dtype": compute_dtype,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_use_double_quant": True,
"llm_int8_enable_fp32_cpu_offload": False,
"llm_int8_has_fp16_weight": False,
"llm_int8_skip_modules": None,
"llm_int8_threshold": 6.0,
"load_in_4bit": True,
"load_in_8bit": False,
"quant_method": "bitsandbytes",
}
model.config.update({"quantization_config": quantization_config})
else:
if hasattr(quantization_config, "to_dict"):
model.config.update(
{"quantization_config": quantization_config.to_dict()}
)
elif isinstance(quantization_config, dict):
model.config.update({"quantization_config": quantization_config})
if load_in_fp8 != False:
_tag_model_with_fp8_torchao_config(model, fp8_mode)

21
unsloth/models/rl.py
View file

@ -199,15 +199,15 @@ def PatchRL(FastLanguageModel):
unwrap = "unwrap_model_for_generation"
for trainer in trainers:
try:
current_trainer = eval(f"trl.trainer.{trainer}")
current_trainer = getattr(trl.trainer, trainer)
except:
continue
if hasattr(current_trainer, unwrap):
try:
exec(f"trl.trainer.{trainer}.{unwrap} = unsloth_{unwrap}")
setattr(current_trainer, unwrap, unsloth_unwrap_model_for_generation)
except:
continue
exec(f"Trainer.prediction_step=unsloth_prediction_step")
Trainer.prediction_step = unsloth_prediction_step
selective_log_softmax = RL_REPLACEMENTS["selective_log_softmax"]
@ -234,6 +234,10 @@ from transformers.training_args import ParallelMode
# Also patches W&B since multiple runs must use wandb.finish()
import functools
from types import MethodType
try:
from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers
except:
def reset_unsloth_gradient_checkpointing_buffers(): pass
def prepare_for_training_mode(f):
@functools.wraps(f)
def wrapper(self, *args, **kwargs):
@ -244,6 +248,11 @@ def prepare_for_training_mode(f):
# Return inference mode
if hasattr(self, 'model') and hasattr(self.model, "for_inference"):
self.model.for_inference()
# Reset gradient checkpointing buffers to free memory while staying ready for next run
try:
reset_unsloth_gradient_checkpointing_buffers()
except:
pass
# Patch W&B to enable logging on future runs, otherwise it'll overwrite the first run
try:
import wandb
@ -817,7 +826,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
num_proc_check = (
"if dataset_num_proc is None:\n"
" import psutil\n"
" dataset_num_proc = min(max(psutil.cpu_count()+4, 2), 64)\n"
" dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64)\n"
" memory_gb_left = psutil.virtual_memory().available / (1024**3)\n"
" if memory_gb_left <= 4: dataset_num_proc = 1 # Too risky, so set to 1\n"
" elif memory_gb_left <= 6: dataset_num_proc = min(2, dataset_num_proc)\n"
@ -994,10 +1003,10 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
# Temporary patch _is_vlm to False
# as of 0.22 it only exists in sfttrainer
oriignal_is_vlm_text = "self._is_vlm = True"
original_is_vlm_text = "self._is_vlm = True"
new_is_vlm_text = "self._is_vlm = False"
RLTrainer_source = RLTrainer_source.replace(
oriignal_is_vlm_text, new_is_vlm_text
original_is_vlm_text, new_is_vlm_text
)
# Remove multiple doc strings

7
unsloth/models/vision.py
View file

@ -529,6 +529,7 @@ class FastBaseModel:
del kwargs["attn_implementation"]
bnb_config = None
user_quantization_config = kwargs.get("quantization_config", None)
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."
@ -596,7 +597,8 @@ class FastBaseModel:
):
pass
else:
kwargs["quantization_config"] = bnb_config
if user_quantization_config is None:
kwargs["quantization_config"] = bnb_config
else:
if auto_config is None:
auto_config = AutoConfig.from_pretrained(
@ -641,7 +643,8 @@ class FastBaseModel:
)
except:
pass
kwargs["quantization_config"] = quantization_config
if user_quantization_config is None:
kwargs["quantization_config"] = quantization_config
# Check if using forced float32 - we load it in bfloat16, then cast to float16!
torch_dtype = dtype

14
unsloth/save.py
View file

@ -879,12 +879,12 @@ def install_llama_cpp_make_non_blocking():
IS_CMAKE = False
if check == 0:
# Uses old MAKE
n_jobs = max(int(psutil.cpu_count() * 1.5), 1)
n_jobs = max(int((psutil.cpu_count() or 1) * 1.5), 1)
full_command = ["make", "all", "-j" + str(n_jobs), "-C", "llama.cpp"]
IS_CMAKE = False
else:
# Uses new CMAKE
n_jobs = max(int(psutil.cpu_count()), 1) # Use less CPUs since 1.5x faster
n_jobs = max(int(psutil.cpu_count() or 1), 1) # Use less CPUs since 1.5x faster
check = os.system(
f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}"
)
@ -994,13 +994,13 @@ def install_llama_cpp_old(version = -10):
# Try using MAKE
commands = [
"make clean -C llama.cpp",
f"make all -j{psutil.cpu_count()*2} -C llama.cpp",
f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp",
]
if try_execute(commands) == "CMAKE":
# Instead use CMAKE
commands = [
f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}",
f"cmake --build llama.cpp/build --config Release -j{psutil.cpu_count()*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}",
f"cmake --build llama.cpp/build --config Release -j{(psutil.cpu_count() or 1)*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}",
"cp llama.cpp/build/bin/llama-* llama.cpp",
"rm -rf llama.cpp/build",
]
@ -1040,14 +1040,14 @@ def install_llama_cpp_blocking(use_cuda = False):
"make clean -C llama.cpp",
# https://github.com/ggerganov/llama.cpp/issues/7062
# Weirdly GPU conversion for GGUF breaks??
# f"{use_cuda} make all -j{psutil.cpu_count()*2} -C llama.cpp",
f"make all -j{psutil.cpu_count()*2} -C llama.cpp",
# f"{use_cuda} make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp",
f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp",
]
if try_execute(commands) == "CMAKE":
# Instead use CMAKE
commands = [
f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}",
f"cmake --build llama.cpp/build --config Release -j{psutil.cpu_count()*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}",
f"cmake --build llama.cpp/build --config Release -j{(psutil.cpu_count() or 1)*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}",
"cp llama.cpp/build/bin/llama-* llama.cpp",
"rm -rf llama.cpp/build",
]