Bug fixes

This commit is contained in:
Daniel Han 2025-03-07 01:43:39 -08:00
commit b4590c9c56
3 changed files with 127 additions and 88 deletions

View file

@ -91,7 +91,7 @@ def original_apply_o(self, X):
pass
from math import sqrt as math_sqrt
KV_CACHE_INCREMENT = 256 # KV Cache update size
KV_CACHE_INCREMENT = 512 # KV Cache update size
torch_nn_functional_softmax = torch.nn.functional.softmax
# SDPA has GQA internally
SDPA_HAS_GQA = "enable_gqa" in scaled_dot_product_attention.__doc__
@ -1656,6 +1656,13 @@ class FastLlamaModel:
"Are you certain you want to do remote code execution?"
)
pass
if fast_inference:
import platform
if platform.system().lower() == 'windows':
print("Unsloth: vLLM does not work in Windows! Will use Unsloth inference!")
fast_inference = False
pass
if token is None: token = get_token()
if model_patcher is None: model_patcher = FastLlamaModel
SUPPORTS_BFLOAT16 = is_bfloat16_supported()
@ -1966,12 +1973,17 @@ class FastLlamaModel:
for layer in model.model.layers:
layer.self_attn.rotary_emb = rotary_emb
pass
# Add for_inference and for_training
model.for_training = functools.partial(FastLlamaModel.for_training, model)
model.for_inference = functools.partial(FastLlamaModel.for_inference, model)
# Patch generate
if model.generate.__name__ != "unsloth_fast_generate":
model._old_generate = model.generate
unsloth_fast_generate.__doc__ = model._old_generate.__doc__
model.generate = types.MethodType(unsloth_fast_generate, model)
pass
return model, tokenizer
pass
@ -2404,7 +2416,7 @@ class FastLlamaModel:
# Add for_inference and for_training
model.for_training = functools.partial(FastLlamaModel.for_training, model)
model.for_inference = functools.partial(FastLlamaModel.for_inference, model)
# Patch generate
if model.generate.__name__ != "unsloth_fast_generate":
model._old_generate = model.generate

View file

@ -611,6 +611,21 @@ __INT_TO_FLOAT_MAPPER = \
"open-thoughts/OpenThinker-7B",
"unsloth/OpenThinker-7B-bnb-4bit",
),
"unsloth/granite-3.2-2b-instruct-unsloth-bnb-4bit" : (
"unsloth/granite-3.2-2b-instruct",
"ibm-granite/granite-3.2-2b-instruct",
"unsloth/granite-3.2-2b-instruct-bnb-4bit",
),
"unsloth/granite-3.2-8b-instruct-unsloth-bnb-4bit" : (
"unsloth/granite-3.2-8b-instruct",
"ibm-granite/granite-3.2-8b-instruct",
"unsloth/granite-3.2-8b-instruct-bnb-4bit",
),
"unsloth/QwQ-32B-unsloth-bnb-4bit" : (
"unsloth/QwQ-32B",
"Qwen/QwQ-32B",
"unsloth/QwQ-32B-bnb-4bit",
),
}
INT_TO_FLOAT_MAPPER = {}

View file

@ -31,40 +31,47 @@ from unsloth_zoo.peft_utils import (
requires_grad_for_gradient_checkpointing,
)
from triton import __version__ as triton_version
from unsloth_zoo.utils import _get_dtype
import types
import functools
__all__ = [
"FastBaseVisionModel",
]
def _wrap_fast_inference(generate, device_type, dtype, model):
# Wraps inference with bfloat16 / float16
@torch.inference_mode
def _fast_generate(*args, **kwargs):
# For num_logits_to_keep
# kwargs["num_logits_to_keep"] = 1
# Remove token_type_ids
kwargs.pop("token_type_ids", None)
def unsloth_vision_fast_generate(
self,
*args,
**kwargs,
):
FastBaseVisionModel.for_inference(self)
# Check pad_token
model_eos_token_id = getattr(model.config, "eos_token_id", None)
if model_eos_token_id is not None and hasattr(model_eos_token_id, "__iter__"):
model_eos_token_id = model_eos_token_id[0]
dtype = _get_dtype(self.config.torch_dtype)
kwargs["pad_token_id"] = kwargs.pop("pad_token_id", model_eos_token_id)
# Remove token_type_ids
kwargs.pop("token_type_ids", None)
try:
kwargs["pixel_values"] = kwargs["pixel_values"].to(model.dtype)
except:
pass
# Check pad_token
model_eos_token_id = getattr(model.config, "eos_token_id", None)
if model_eos_token_id is not None and hasattr(model_eos_token_id, "__iter__"):
model_eos_token_id = model_eos_token_id[0]
# Autocasted
with torch.autocast(device_type = device_type, dtype = dtype):
output = generate(*args, **kwargs)
kwargs["pad_token_id"] = kwargs.pop("pad_token_id", model_eos_token_id)
try:
kwargs["pixel_values"] = kwargs["pixel_values"].to(dtype)
except:
pass
return output
# Mixed precision autocast
with torch.inference_mode(), torch.autocast(device_type = "cuda", dtype = dtype):
output = self._old_generate(*args, **kwargs)
pass
return _fast_generate
FastBaseVisionModel.for_training(self)
return output
pass
@ -94,12 +101,16 @@ class FastBaseVisionModel:
gpu_stats = torch.cuda.get_device_properties(0)
max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
from importlib.metadata import version as importlib_version
try: vllm_version = f" vLLM: {importlib_version('vllm')}."
except: vllm_version = ""
statistics = \
f"==((====))== Unsloth {__version__}: Fast {model_types[0].title()} vision patching. Transformers: {transformers_version}.\n"\
f" {chr(92)}{chr(92)} /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB. Platform: {platform_system}.\n"\
f"==((====))== Unsloth {__version__}: Fast {model_types[0].title()} patching. Transformers: {transformers_version}.{vllm_version}\n"\
f" {chr(92)}{chr(92)} /| {gpu_stats.name}. Num GPUs = {torch.cuda.device_count()}. Max memory: {max_memory} GB. Platform: {platform_system}.\n"\
f"O^O/ {chr(92)}_/ {chr(92)} Torch: {torch.__version__}. CUDA: {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit: {torch.version.cuda}. Triton: {triton_version}\n"\
f"{chr(92)} / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. FA [Xformers = {xformers_version}. FA2 = {HAS_FLASH_ATTENTION}]\n"\
f' "-____-" Free Apache license: http://github.com/unslothai/unsloth'
f' "-____-" Free license: http://github.com/unslothai/unsloth'
print(statistics)
# Warn about fast transfers
@ -136,7 +147,7 @@ class FastBaseVisionModel:
# Cannot be None, since HF now checks for the config
if load_in_4bit: kwargs["quantization_config"] = bnb_config
model = AutoModelForVision2Seq.from_pretrained(
model_name,
device_map = device_map,
@ -190,10 +201,20 @@ class FastBaseVisionModel:
internal_model = model
while hasattr(internal_model, "model"):
internal_model._saved_temp_tokenizer = tokenizer
# Also set is_loaded_in_8bit to disable incorrect DDP
internal_model.is_loaded_in_8bit = True
internal_model = internal_model.model
pass
internal_model._saved_temp_tokenizer = tokenizer
# Also set is_loaded_in_8bit to disable incorrect DDP
internal_model.is_loaded_in_8bit = True
# Patch generate
if model.generate.__name__ != "unsloth_vision_fast_generate":
model._old_generate = model.generate
unsloth_vision_fast_generate.__doc__ = model._old_generate.__doc__
model.generate = types.MethodType(unsloth_vision_fast_generate, model)
return model, tokenizer
pass
@ -281,6 +302,9 @@ class FastBaseVisionModel:
pass
patch_saving_functions(model, vision = True)
# Add for_inference and for_training
model.for_training = functools.partial(FastBaseVisionModel.for_training, model)
model.for_inference = functools.partial(FastBaseVisionModel.for_inference, model)
return model
pass
@ -319,57 +343,52 @@ class FastBaseVisionModel:
if hasattr(internal_model, "_saved_temp_tokenizer"):
internal_model._saved_temp_tokenizer.tokenizer.padding_side = "right"
pass
# Also set is_loaded_in_8bit to disable incorrect DDP
internal_model.is_loaded_in_8bit = True
internal_model = internal_model.model
pass
if hasattr(internal_model, "_saved_temp_tokenizer"):
internal_model._saved_temp_tokenizer.tokenizer.padding_side = "right"
pass
# Also set is_loaded_in_8bit to disable incorrect DDP
internal_model.is_loaded_in_8bit = True
# Clear deleted GPU items
for _ in range(3):
gc.collect()
torch.cuda.empty_cache()
pass
# Add for_inference and for_training
model.for_training = functools.partial(FastBaseVisionModel.for_training, model)
model.for_inference = functools.partial(FastBaseVisionModel.for_inference, model)
# Patch generate
if model.generate.__name__ != "unsloth_vision_fast_generate":
model._old_generate = model.generate
unsloth_vision_fast_generate.__doc__ = model._old_generate.__doc__
model.generate = types.MethodType(unsloth_vision_fast_generate, model)
return model
pass
@staticmethod
def for_inference(model):
model.gradient_checkpointing = False
model.training = False
if not hasattr(model, "parameters"):
raise TypeError("Unsloth: I think you're passing a tokenizer, not the model to for_inference!")
for name, module in model.named_modules():
if hasattr(module, "gradient_checkpointing"):
module.gradient_checkpointing = False
if hasattr(module, "training"):
module.training = False
pass
dtype = model.config.torch_dtype
if type(dtype) is str:
if dtype == "float16": dtype = torch.float16
elif dtype == "bfloat16": dtype = torch.bfloat16
pass
device_type = model.device.type
# Wrap model.generate
if model.generate.__name__ != "_fast_generate":
model._unwrapped_old_generate = model.generate
model.generate = _wrap_fast_inference(model.generate, device_type, dtype, model)
pass
# Patch tokenizer to pad to the left
internal_model = model
while hasattr(internal_model, "model"):
if hasattr(internal_model, "_saved_temp_tokenizer"):
internal_model._saved_temp_tokenizer.tokenizer.padding_side = "left"
pass
internal_model = internal_model.model
pass
if hasattr(internal_model, "_saved_temp_tokenizer"):
internal_model._saved_temp_tokenizer.tokenizer.padding_side = "left"
def _for_inference(m):
if hasattr(m, "gradient_checkpointing"): m.gradient_checkpointing = False
if hasattr(m, "training"): m.training = False
# Pad tokenizer to the left
if hasattr(m, "_saved_temp_tokenizer"): m._saved_temp_tokenizer.padding_side = "left"
# Set a flag for generation!
m._flag_for_generation = True
pass
m = model
while hasattr(m, "model"):
_for_inference(m)
m = m.model
_for_inference(m)
# Also disable training for embeddings for NEFTune
if hasattr(model, "get_input_embeddings"):
@ -380,40 +399,34 @@ class FastBaseVisionModel:
embeddings = model.get_output_embeddings()
if hasattr(embeddings, "training"): embeddings.training = False
pass
return model
pass
@staticmethod
def for_training(model, use_gradient_checkpointing = True):
model.gradient_checkpointing = use_gradient_checkpointing
model.training = True
if not hasattr(model, "parameters"):
raise TypeError("Unsloth: I think you're passing a tokenizer, not the model to for_training!")
for name, module in model.named_modules():
if hasattr(module, "gradient_checkpointing"):
module.gradient_checkpointing = use_gradient_checkpointing
if hasattr(module, "training"):
module.training = True
# Delete all fast inference loras
for param in model.parameters():
if hasattr(param, "_fast_lora"):
del param._fast_lora
pass
# Also revert model.generate
if hasattr(model, "_unwrapped_old_generate"):
model.generate = model._unwrapped_old_generate
del model._unwrapped_old_generate
pass
# Patch tokenizer to pad to the right
internal_model = model
while hasattr(internal_model, "model"):
if hasattr(internal_model, "_saved_temp_tokenizer"):
internal_model._saved_temp_tokenizer.tokenizer.padding_side = "right"
pass
internal_model = internal_model.model
pass
if hasattr(internal_model, "_saved_temp_tokenizer"):
internal_model._saved_temp_tokenizer.tokenizer.padding_side = "right"
def _for_training(m):
if hasattr(m, "gradient_checkpointing"): m.gradient_checkpointing = use_gradient_checkpointing
if hasattr(m, "training"): m.training = True
# Pad tokenizer to the left
if hasattr(m, "_saved_temp_tokenizer"): m._saved_temp_tokenizer.padding_side = "right"
# Set a flag for generation!
if hasattr(m, "_flag_for_generation"): del m._flag_for_generation
pass
m = model
while hasattr(m, "model"):
_for_training(m)
m = m.model
_for_training(m)
# Also re-enable training for embeddings for NEFTune
if hasattr(model, "get_input_embeddings"):
@ -424,7 +437,6 @@ class FastBaseVisionModel:
embeddings = model.get_output_embeddings()
if hasattr(embeddings, "training"): embeddings.training = True
pass
return model
pass
pass