Merge branch 'unslothai:main' into FST
This commit is contained in:
electroglyph
GitHub
commit
c8be8fff88
9 changed files with 306 additions and 80 deletions
|
|
@ -51,8 +51,8 @@ Use our official [Unsloth Docker image](https://hub.docker.com/r/unsloth/unsloth
|
|||
For RTX 50x, B200, 6000 GPUs: `pip install unsloth`. Read our [Blackwell Guide](https://docs.unsloth.ai/basics/training-llms-with-blackwell-rtx-50-series-and-unsloth) and [DGX Spark Guide](https://docs.unsloth.ai/new/fine-tuning-llms-with-nvidia-dgx-spark-and-unsloth) for more details.
|
||||
|
||||
## 🦥 Unsloth News
|
||||
- New RoPE & MLP **Triton Kernels** & **Auto Packing**: 3x faster training & 30% less VRAM. [Blog](https://docs.unsloth.ai/new/3x-faster-training-packing)
|
||||
- **Ministral 3** by Mistral: Run Ministral 3 or fine-tune with our vision or RL sodoku notebook. [Guide](https://docs.unsloth.ai/new/ministral-3) • [Notebooks](https://docs.unsloth.ai/new/ministral-3#fine-tuningb)
|
||||
- New RoPE & MLP **Triton Kernels** & **Padding Free + Packing**: 3x faster training & 30% less VRAM. [Blog](https://docs.unsloth.ai/new/3x-faster-training-packing)
|
||||
- **Ministral 3** by Mistral: Run Ministral 3 or fine-tune with vision/RL sodoku notebooks. [Guide](https://docs.unsloth.ai/new/ministral-3) • [Notebooks](https://docs.unsloth.ai/new/ministral-3#fine-tuningb)
|
||||
- **500K Context**: Training a 20B model with >500K context is now possible on an 80GB GPU. [Blog](https://docs.unsloth.ai/new/500k-context-length-fine-tuning)
|
||||
- **FP8 Reinforcement Learning**: You can now do FP8 GRPO on consumer GPUs. [Blog](https://docs.unsloth.ai/new/fp8-reinforcement-learning) • [Notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_8B_FP8_GRPO.ipynb)
|
||||
- **DeepSeek-OCR**: Fine-tune to improve language understanding by 89%. [Guide](https://docs.unsloth.ai/new/deepseek-ocr-run-and-fine-tune) • [Notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Deepseek_OCR_(3B).ipynb)
|
||||
|
|
|
|||
|
|
@ -60,7 +60,7 @@ huggingfacenotorch = [
|
|||
]
|
||||
huggingface = [
|
||||
"unsloth[huggingfacenotorch]",
|
||||
"unsloth_zoo>=2025.12.3",
|
||||
"unsloth_zoo>=2025.12.4",
|
||||
"torchvision",
|
||||
"unsloth[triton]",
|
||||
]
|
||||
|
|
@ -523,7 +523,7 @@ colab-ampere-torch220 = [
|
|||
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
|
||||
]
|
||||
colab-new = [
|
||||
"unsloth_zoo>=2025.12.3",
|
||||
"unsloth_zoo>=2025.12.4",
|
||||
"packaging",
|
||||
"tyro",
|
||||
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,<=4.57.3",
|
||||
|
|
|
|||
|
|
@ -17,6 +17,13 @@ from packaging.version import Version
|
|||
import os, re, subprocess, inspect, functools
|
||||
import numpy as np
|
||||
|
||||
# Log Unsloth is being used
|
||||
os.environ["UNSLOTH_IS_PRESENT"] = "1"
|
||||
|
||||
# Check if modules that need patching are already imported
|
||||
critical_modules = ["trl", "transformers", "peft"]
|
||||
already_imported = [mod for mod in critical_modules if mod in sys.modules]
|
||||
|
||||
# Fix some issues before importing other packages
|
||||
from .import_fixes import (
|
||||
fix_message_factory_issue,
|
||||
|
|
@ -31,10 +38,6 @@ del fix_message_factory_issue
|
|||
del check_fbgemm_gpu_version
|
||||
del torchvision_compatibility_check
|
||||
|
||||
# Check if modules that need patching are already imported
|
||||
critical_modules = ["trl", "transformers", "peft"]
|
||||
already_imported = [mod for mod in critical_modules if mod in sys.modules]
|
||||
|
||||
# This check is critical because Unsloth optimizes these libraries by modifying
|
||||
# their code at import time. If they're imported first, the original (slower,
|
||||
# more memory-intensive) implementations will be used instead of Unsloth's
|
||||
|
|
@ -43,7 +46,7 @@ if already_imported:
|
|||
# stacklevel=2 makes warning point to user's import line rather than this library code,
|
||||
# showing them exactly where to fix the import order in their script
|
||||
warnings.warn(
|
||||
f"WARNING: Unsloth should be imported before {', '.join(already_imported)} "
|
||||
f"WARNING: Unsloth should be imported before [{', '.join(already_imported)}] "
|
||||
f"to ensure all optimizations are applied. Your code may run slower or encounter "
|
||||
f"memory issues without these optimizations.\n\n"
|
||||
f"Please restructure your imports with 'import unsloth' at the top of your file.",
|
||||
|
|
@ -63,8 +66,6 @@ os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
|
|||
# "pinned_use_cuda_host_register:True,"\
|
||||
# "pinned_num_register_threads:8"
|
||||
|
||||
# Log Unsloth is being used
|
||||
os.environ["UNSLOTH_IS_PRESENT"] = "1"
|
||||
|
||||
from importlib.metadata import version as importlib_version
|
||||
from importlib.metadata import PackageNotFoundError
|
||||
|
|
@ -72,7 +73,7 @@ from importlib.metadata import PackageNotFoundError
|
|||
# Check for unsloth_zoo
|
||||
try:
|
||||
unsloth_zoo_version = importlib_version("unsloth_zoo")
|
||||
if Version(unsloth_zoo_version) < Version("2025.12.3"):
|
||||
if Version(unsloth_zoo_version) < Version("2025.12.4"):
|
||||
print(
|
||||
"Unsloth: Please update Unsloth and Unsloth-Zoo to the latest version!\n"
|
||||
"Do this via `pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth unsloth_zoo`"
|
||||
|
|
@ -123,6 +124,8 @@ from .import_fixes import (
|
|||
patch_ipykernel_hf_xet,
|
||||
patch_trackio,
|
||||
patch_datasets,
|
||||
patch_enable_input_require_grads,
|
||||
fix_openenv_no_vllm,
|
||||
)
|
||||
|
||||
fix_xformers_performance_issue()
|
||||
|
|
@ -132,6 +135,8 @@ ignore_logger_messages()
|
|||
patch_ipykernel_hf_xet()
|
||||
patch_trackio()
|
||||
patch_datasets()
|
||||
patch_enable_input_require_grads()
|
||||
fix_openenv_no_vllm()
|
||||
|
||||
del fix_xformers_performance_issue
|
||||
del fix_vllm_aimv2_issue
|
||||
|
|
@ -140,6 +145,8 @@ del ignore_logger_messages
|
|||
del patch_ipykernel_hf_xet
|
||||
del patch_trackio
|
||||
del patch_datasets
|
||||
del patch_enable_input_require_grads
|
||||
del fix_openenv_no_vllm
|
||||
|
||||
# Torch 2.4 has including_emulation
|
||||
if DEVICE_TYPE == "cuda":
|
||||
|
|
|
|||
|
|
@ -19,8 +19,8 @@ from importlib.metadata import version as importlib_version
|
|||
from packaging.version import Version as TrueVersion
|
||||
import re
|
||||
import logging
|
||||
|
||||
UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1"
|
||||
# Cannot import logger here since it'll import transformers
|
||||
# from unsloth_zoo.log import logger
|
||||
|
||||
|
||||
def Version(version):
|
||||
|
|
@ -70,9 +70,10 @@ def fix_message_factory_issue():
|
|||
def GetPrototype(self, *args, **kwargs):
|
||||
return
|
||||
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
if not hasattr(google.protobuf.message_factory, "MessageFactory"):
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
|
||||
logger.info("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
|
||||
google.protobuf.message_factory.MessageFactory = MessageFactory
|
||||
elif (
|
||||
hasattr(google.protobuf.message_factory, "MessageFactory")
|
||||
|
|
@ -82,8 +83,7 @@ def fix_message_factory_issue():
|
|||
and not hasattr(google.protobuf.message_factory, "GetMessageClass")
|
||||
):
|
||||
google.protobuf.message_factory.MessageFactory = MessageFactory
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
|
||||
logger.info("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
|
||||
elif (
|
||||
hasattr(google.protobuf.message_factory, "MessageFactory")
|
||||
and not hasattr(
|
||||
|
|
@ -97,8 +97,7 @@ def fix_message_factory_issue():
|
|||
return GetMessageClass(descriptor)
|
||||
|
||||
google.protobuf.message_factory.MessageFactory.GetPrototype = GetPrototype
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print("Unsloth: Patching protobuf.MessageFactory.GetPrototype")
|
||||
logger.info("Unsloth: Patching protobuf.MessageFactory.GetPrototype")
|
||||
pass
|
||||
except:
|
||||
pass
|
||||
|
|
@ -110,6 +109,8 @@ def fix_xformers_performance_issue():
|
|||
return
|
||||
xformers_version = importlib_version("xformers")
|
||||
if Version(xformers_version) < Version("0.0.29"):
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
xformers_location = importlib.util.find_spec("xformers").origin
|
||||
xformers_location = os.path.split(xformers_location)[0]
|
||||
cutlass = Path(xformers_location) / "ops" / "fmha" / "cutlass.py"
|
||||
|
|
@ -126,13 +127,11 @@ def fix_xformers_performance_issue():
|
|||
f.seek(0)
|
||||
f.write(text)
|
||||
f.truncate()
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print(
|
||||
"Unsloth: Patching Xformers to fix some performance issues."
|
||||
)
|
||||
logger.info(
|
||||
"Unsloth: Patching Xformers to fix some performance issues."
|
||||
)
|
||||
except Exception as e:
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print(f"Unsloth: Failed patching Xformers with error = {str(e)}")
|
||||
logger.info(f"Unsloth: Failed patching Xformers with error = {str(e)}")
|
||||
|
||||
|
||||
# ValueError: 'aimv2' is already used by a Transformers config, pick another name.
|
||||
|
|
@ -141,6 +140,8 @@ def fix_vllm_aimv2_issue():
|
|||
return
|
||||
vllm_version = importlib_version("vllm")
|
||||
if Version(vllm_version) < Version("0.10.1"):
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
vllm_version = importlib.util.find_spec("vllm").origin
|
||||
vllm_version = os.path.split(vllm_version)[0]
|
||||
ovis_config = Path(vllm_version) / "transformers_utils" / "configs" / "ovis.py"
|
||||
|
|
@ -167,13 +168,11 @@ def fix_vllm_aimv2_issue():
|
|||
f.seek(0)
|
||||
f.write(text)
|
||||
f.truncate()
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print(
|
||||
"Unsloth: Patching vLLM to fix `'aimv2' is already used by a Transformers config, pick another name.`"
|
||||
)
|
||||
logger.info(
|
||||
"Unsloth: Patching vLLM to fix `'aimv2' is already used by a Transformers config, pick another name.`"
|
||||
)
|
||||
except Exception as e:
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print(f"Unsloth: Failed patching vLLM with error = {str(e)}")
|
||||
logger.info(f"Unsloth: Failed patching vLLM with error = {str(e)}")
|
||||
|
||||
|
||||
def fix_vllm_guided_decoding_params():
|
||||
|
|
@ -274,8 +273,74 @@ def check_fbgemm_gpu_version():
|
|||
raise ImportError(
|
||||
f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} detected. It might cause unexpected issues like segmentation faults. Please uninstall the current one by doing `pip uninstall fbgemm-gpu` && `pip install fbgemm-gpu` to install fbgemm-gpu 1.4.0 or newer!"
|
||||
)
|
||||
elif UNSLOTH_ENABLE_LOGGING:
|
||||
print(f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} detected.")
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
logger.info(f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} detected.")
|
||||
|
||||
|
||||
def patch_enable_input_require_grads():
|
||||
"""
|
||||
Patch transformers PreTrainedModel.enable_input_require_grads to handle vision models
|
||||
that raise NotImplementedError from get_input_embeddings().
|
||||
|
||||
"""
|
||||
import inspect
|
||||
from transformers import PreTrainedModel
|
||||
|
||||
# Check if the original function iterates over self.modules() instead of just returning the enable_input_require_grads
|
||||
# Ref: https://github.com/huggingface/transformers/pull/41993/files#diff-6b72b98c4c2dcfc6cc606843917733f5d858374fbc22a735ff483bbc0c1e63eaL1979-R1996
|
||||
try:
|
||||
original_source = inspect.getsource(PreTrainedModel.enable_input_require_grads)
|
||||
except:
|
||||
return
|
||||
|
||||
# Only patch if the new pattern exists (iterating over self.modules())
|
||||
if "for module in self.modules()" not in original_source:
|
||||
return
|
||||
|
||||
def _patched_enable_input_require_grads(self):
|
||||
def make_inputs_require_grads(module, input, output):
|
||||
output.requires_grad_(True)
|
||||
|
||||
hooks = []
|
||||
seen_modules = set()
|
||||
|
||||
for module in self.modules():
|
||||
if not (
|
||||
isinstance(module, PreTrainedModel)
|
||||
and hasattr(module, "get_input_embeddings")
|
||||
):
|
||||
continue
|
||||
|
||||
try:
|
||||
input_embeddings = module.get_input_embeddings()
|
||||
except NotImplementedError:
|
||||
# Vision models may not implement get_input_embeddings - skip them
|
||||
# For GLM V4.6 for example, this skips only `self.visual`
|
||||
continue
|
||||
|
||||
if input_embeddings is None:
|
||||
continue
|
||||
|
||||
embedding_id = id(input_embeddings)
|
||||
if embedding_id in seen_modules:
|
||||
continue
|
||||
|
||||
seen_modules.add(embedding_id)
|
||||
hooks.append(
|
||||
input_embeddings.register_forward_hook(make_inputs_require_grads)
|
||||
)
|
||||
|
||||
self._require_grads_hooks = hooks
|
||||
if hooks:
|
||||
self._require_grads_hook = hooks[0]
|
||||
|
||||
PreTrainedModel.enable_input_require_grads = _patched_enable_input_require_grads
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
logger.info(
|
||||
"Unsloth: Patched enable_input_require_grads for vision model compatibility"
|
||||
)
|
||||
|
||||
|
||||
def torchvision_compatibility_check():
|
||||
|
|
@ -313,7 +378,49 @@ def torchvision_compatibility_check():
|
|||
f"but found torchvision=={torchvision_version}. "
|
||||
f"Please refer to https://pytorch.org/get-started/previous-versions/ for more information."
|
||||
)
|
||||
elif UNSLOTH_ENABLE_LOGGING:
|
||||
print(
|
||||
f"Unsloth: torch=={torch_version} and torchvision=={torchvision_version} are compatible."
|
||||
)
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
logger.info(
|
||||
f"Unsloth: torch=={torch_version} and torchvision=={torchvision_version} are compatible."
|
||||
)
|
||||
|
||||
|
||||
# Fix TRL OpenEnv 0.26 NameError: name 'SamplingParams' is not defined
|
||||
def fix_openenv_no_vllm():
|
||||
if importlib.util.find_spec("trl") is None:
|
||||
return
|
||||
trl_location = importlib.util.find_spec("trl").origin
|
||||
trl_location = os.path.split(trl_location)[0]
|
||||
openenv = Path(trl_location) / "experimental" / "openenv" / "utils.py"
|
||||
if not openenv.exists():
|
||||
return
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
try:
|
||||
with open(openenv, "r+", encoding = "utf-8") as f:
|
||||
text = f.read()
|
||||
bad = (
|
||||
"if is_vllm_available():\n"
|
||||
" from vllm import SamplingParams\n"
|
||||
" from vllm.sampling_params import GuidedDecodingParams\n"
|
||||
)
|
||||
if bad + "\n" + "\n" in text:
|
||||
text = text.replace(
|
||||
bad + "\n" + "\n",
|
||||
bad
|
||||
+ (
|
||||
"else:\n"
|
||||
" from typing import Any\n"
|
||||
" SamplingParams = Any\n"
|
||||
" GuidedDecodingParams = Any\n"
|
||||
"\n"
|
||||
),
|
||||
)
|
||||
f.seek(0)
|
||||
f.write(text)
|
||||
f.truncate()
|
||||
logger.info(
|
||||
"Unsloth: Patching TRL OpenEnv to fix SamplingParams not defined"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.info(f"Unsloth: Failed patching TRL OpenEnv with error = {str(e)}")
|
||||
|
|
|
|||
|
|
@ -20,13 +20,6 @@ from ..device_type import DEVICE_COUNT
|
|||
from .utils import calculate_settings, torch_gpu_device, torch_device_stream
|
||||
|
||||
|
||||
@triton.heuristics(
|
||||
{
|
||||
"BACKWARD_PASS": lambda args: bool(args["BACKWARD_PASS"]),
|
||||
"HAS_ROPE_INDICES": lambda args: bool(args["HAS_ROPE_INDICES"]),
|
||||
}
|
||||
)
|
||||
@triton.jit
|
||||
def _rope_embedding_QK(
|
||||
Q,
|
||||
Q_batch_stride,
|
||||
|
|
@ -104,6 +97,15 @@ def _rope_embedding_QK(
|
|||
tl.store(k_ptr + half_head_dim + col_offsets, k1 * cos1 + k0 * sin1, mask = mask)
|
||||
|
||||
|
||||
_rope_embedding_QK = triton.jit(_rope_embedding_QK)
|
||||
_rope_embedding_QK = triton.heuristics(
|
||||
{
|
||||
"BACKWARD_PASS": lambda args: bool(args["BACKWARD_PASS"]),
|
||||
"HAS_ROPE_INDICES": lambda args: bool(args["HAS_ROPE_INDICES"]),
|
||||
}
|
||||
)(_rope_embedding_QK)
|
||||
|
||||
|
||||
ROPE_GROUP_SIZE: int = 4
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
__version__ = "2025.12.4"
|
||||
__version__ = "2025.12.5"
|
||||
|
||||
__all__ = [
|
||||
"SUPPORTS_BFLOAT16",
|
||||
|
|
@ -413,6 +413,16 @@ try:
|
|||
except:
|
||||
pass
|
||||
|
||||
# Flax classes are deprecated and will be removed in Diffusers v1.0.0.
|
||||
try:
|
||||
from diffusers.utils import logger as diffusers_logger
|
||||
|
||||
diffusers_logger.addFilter(HideLoggingMessage("are deprecated"))
|
||||
del diffusers_logger
|
||||
except:
|
||||
pass
|
||||
|
||||
|
||||
# Errors out on
|
||||
# Some weights of Gemma3nForConditionalGeneration were not initialized from the model checkpoint
|
||||
from transformers.modeling_utils import logger as transformers_logger
|
||||
|
|
|
|||
|
|
@ -26,7 +26,10 @@ import torch
|
|||
import inspect
|
||||
from collections import defaultdict
|
||||
from unsloth_zoo.rl_replacements import RL_REPLACEMENTS, left_pack_padding
|
||||
from unsloth_zoo.utils import Version
|
||||
from importlib.metadata import version as importlib_version
|
||||
from unsloth_zoo.log import logger
|
||||
import importlib.util
|
||||
from ..device_type import (
|
||||
is_hip,
|
||||
get_device_type,
|
||||
|
|
@ -942,11 +945,15 @@ def openenv_vllm_reload_weights():
|
|||
#
|
||||
# The fix: Use wake_up() with no tags, which wakes everything. Unsloth's patched
|
||||
# CuMemAllocator.wake_up skips weights anyway, so this is safe.
|
||||
if importlib.util.find_spec("trl") is None:
|
||||
return
|
||||
if Version(importlib_version("trl")) < Version("0.26.0"):
|
||||
return
|
||||
try:
|
||||
import trl.experimental.openenv.utils as openenv_utils
|
||||
import trl.experimental.openenv as openenv
|
||||
except ImportError as e:
|
||||
logger.warning(f"Unsloth: Failed to import trl openenv: {e}")
|
||||
logger.info(f"Unsloth: Failed to import trl openenv: {e}")
|
||||
return
|
||||
|
||||
src = inspect.getsource(openenv_utils.generate_rollout_completions)
|
||||
|
|
|
|||
152
unsloth/save.py
152
unsloth/save.py
|
|
@ -42,6 +42,7 @@ import re
|
|||
from transformers.models.llama.modeling_llama import logger
|
||||
from .tokenizer_utils import fix_sentencepiece_gguf
|
||||
from .models.loader_utils import get_model_name
|
||||
from .models._utils import _convert_torchao_model
|
||||
from .ollama_template_mappers import OLLAMA_TEMPLATES, MODEL_TO_OLLAMA_TEMPLATE_MAPPER
|
||||
from transformers import ProcessorMixin
|
||||
from huggingface_hub import HfApi
|
||||
|
|
@ -2734,11 +2735,35 @@ def unsloth_generic_push_to_hub_merged(
|
|||
gc.collect()
|
||||
|
||||
|
||||
def unsloth_save_pretrained_torchao(
|
||||
self,
|
||||
def _unsloth_save_torchao_with_attached_config(
|
||||
model,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
tokenizer = None,
|
||||
torchao_config = None,
|
||||
tokenizer,
|
||||
push_to_hub: bool = False,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
):
|
||||
"""Save a QAT-trained model by converting fake-quantized weights to real quantized weights."""
|
||||
# Convert QAT fake-quantized weights to real quantized weights
|
||||
_convert_torchao_model(model)
|
||||
|
||||
# TorchAO does not support safe_serialization reliably
|
||||
safe_serialization = False
|
||||
|
||||
if push_to_hub:
|
||||
model.push_to_hub(
|
||||
save_directory, safe_serialization = safe_serialization, token = token
|
||||
)
|
||||
tokenizer.push_to_hub(save_directory, token = token)
|
||||
else:
|
||||
model.save_pretrained(save_directory, safe_serialization = safe_serialization)
|
||||
tokenizer.save_pretrained(save_directory)
|
||||
|
||||
|
||||
def _unsloth_save_torchao_with_given_config(
|
||||
model,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
tokenizer,
|
||||
torchao_config,
|
||||
push_to_hub: bool = False,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
):
|
||||
|
|
@ -2749,23 +2774,26 @@ def unsloth_save_pretrained_torchao(
|
|||
`torchao_config` (TorchAOBaseConfig): configuration for torchao quantization, full list: https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize
|
||||
`push_to_hub` (bool): whether to push the checkpoint to huggingface hub or save locally
|
||||
"""
|
||||
|
||||
if push_to_hub:
|
||||
assert token is not None, "Unsloth: Please specify a token for uploading!"
|
||||
|
||||
assert (
|
||||
torchao_config is not None
|
||||
), "Unsloth: Please specify a torchao_config for post-training quantization!"
|
||||
|
||||
# first merge the lora weights
|
||||
arguments = dict(locals())
|
||||
arguments["model"] = self
|
||||
arguments["tokenizer"] = tokenizer
|
||||
arguments["push_to_hub"] = False # We save ourselves
|
||||
arguments["save_method"] = "merged_16bit" # Must be 16bit
|
||||
del arguments["self"]
|
||||
del arguments["torchao_config"]
|
||||
|
||||
if token is None and push_to_hub:
|
||||
token = get_token()
|
||||
|
||||
if not isinstance(self, PeftModelForCausalLM) and not isinstance(self, PeftModel):
|
||||
self.save_pretrained(save_directory)
|
||||
if not isinstance(model, PeftModelForCausalLM) and not isinstance(model, PeftModel):
|
||||
model.save_pretrained(save_directory)
|
||||
tokenizer.save_pretrained(save_directory)
|
||||
else:
|
||||
unsloth_generic_save(**arguments)
|
||||
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
|
||||
|
|
@ -2778,26 +2806,20 @@ def unsloth_save_pretrained_torchao(
|
|||
)
|
||||
from torchao import quantize_
|
||||
|
||||
if torchao_config is None:
|
||||
from torchao.quantization import Int8DynamicActivationInt8WeightConfig
|
||||
|
||||
print(
|
||||
"Unsloth: You did not specify a `torchao_config`, so defaulting to `Int8DynamicActivationInt8WeightConfig`"
|
||||
)
|
||||
torchao_config = Int8DynamicActivationInt8WeightConfig()
|
||||
quantization_config = TorchAoConfig(quant_type = torchao_config)
|
||||
|
||||
# Determine if this is a VLM
|
||||
is_vlm = False
|
||||
if hasattr(self, "config") and hasattr(self.config, "architectures"):
|
||||
if hasattr(model, "config") and hasattr(model.config, "architectures"):
|
||||
is_vlm = any(
|
||||
x.endswith(("ForConditionalGeneration", "ForVisionText2Text"))
|
||||
for x in self.config.architectures
|
||||
for x in model.config.architectures
|
||||
)
|
||||
is_vlm = is_vlm or hasattr(self.config, "vision_config")
|
||||
is_vlm = is_vlm or hasattr(model.config, "vision_config")
|
||||
auto_model = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM
|
||||
auto_processor = AutoProcessor if is_vlm else AutoTokenizer
|
||||
|
||||
tokenizer = auto_processor.from_pretrained(arguments["save_directory"])
|
||||
tokenizer = auto_processor.from_pretrained(save_directory)
|
||||
|
||||
# TorchAO must only use bfloat16 for loading (float16 fails)
|
||||
if HAS_TORCH_DTYPE:
|
||||
|
|
@ -2805,8 +2827,9 @@ def unsloth_save_pretrained_torchao(
|
|||
else:
|
||||
kwargs = {"dtype": torch.bfloat16}
|
||||
|
||||
model = auto_model.from_pretrained(
|
||||
arguments["save_directory"],
|
||||
# Reload with quantization applied
|
||||
quantized_model = auto_model.from_pretrained(
|
||||
save_directory,
|
||||
device_map = "auto",
|
||||
quantization_config = quantization_config,
|
||||
**kwargs,
|
||||
|
|
@ -2817,25 +2840,92 @@ def unsloth_save_pretrained_torchao(
|
|||
# TorchAO does not support safe_serialization right now 0.14.0 seems broken!
|
||||
safe_serialization = Version(importlib_version("torchao")) > Version("0.14.0")
|
||||
safe_serialization = False
|
||||
|
||||
if push_to_hub:
|
||||
if token is None and push_to_hub:
|
||||
token = get_token()
|
||||
model.push_to_hub(
|
||||
quantized_model.push_to_hub(
|
||||
torchao_save_directory, safe_serialization = safe_serialization, token = token
|
||||
)
|
||||
tokenizer.push_to_hub(torchao_save_directory, token = token)
|
||||
else:
|
||||
model.save_pretrained(
|
||||
quantized_model.save_pretrained(
|
||||
torchao_save_directory, safe_serialization = safe_serialization
|
||||
)
|
||||
tokenizer.save_pretrained(torchao_save_directory)
|
||||
|
||||
# Clean up the intermediate unquantized model
|
||||
if os.path.exists(save_directory):
|
||||
try:
|
||||
import shutil
|
||||
|
||||
shutil.rmtree(save_directory)
|
||||
except:
|
||||
pass
|
||||
|
||||
|
||||
def unsloth_save_pretrained_torchao(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
tokenizer = None,
|
||||
torchao_config = None,
|
||||
push_to_hub: bool = False,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
):
|
||||
"""Saves a torchao quantized model checkpoint.
|
||||
|
||||
This function handles two mutually exclusive workflows:
|
||||
|
||||
1. **QAT (Quantization-Aware Training)**: If the model was trained with `qat_scheme`
|
||||
parameter, do NOT pass `torchao_config`. The function will convert the QAT
|
||||
fake-quantized weights to real quantized weights and save directly.
|
||||
|
||||
2. **PTQ (Post-Training Quantization)**: If you want to apply quantization to a
|
||||
regular model, pass a `torchao_config`. The model must NOT have been trained
|
||||
with `qat_scheme`.
|
||||
|
||||
Args:
|
||||
`save_directory`: local folder path or huggingface hub ID when `push_to_hub` is True
|
||||
`tokenizer`: the tokenizer to save alongside the model
|
||||
`torchao_config` (TorchAOBaseConfig): configuration for torchao quantization.
|
||||
Required for PTQ, must be None for QAT models.
|
||||
Options: https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize
|
||||
`push_to_hub` (bool): whether to push to huggingface hub or save locally
|
||||
`token`: HuggingFace token for pushing to hub
|
||||
"""
|
||||
if token is None and push_to_hub:
|
||||
token = get_token()
|
||||
|
||||
has_qat_config = (
|
||||
hasattr(self, "_torchao_config") and self._torchao_config is not None
|
||||
)
|
||||
|
||||
if torchao_config is not None:
|
||||
# PTQ path: user provided a config, model must NOT have QAT config
|
||||
assert not has_qat_config, (
|
||||
"Unsloth: You passed `torchao_config` but this model was trained with `qat_scheme`. "
|
||||
"For QAT models, do not pass `torchao_config` - the quantization config is already "
|
||||
"attached to the model from training."
|
||||
)
|
||||
_unsloth_save_torchao_with_given_config(
|
||||
model = self,
|
||||
save_directory = save_directory,
|
||||
tokenizer = tokenizer,
|
||||
torchao_config = torchao_config,
|
||||
push_to_hub = push_to_hub,
|
||||
token = token,
|
||||
)
|
||||
else:
|
||||
# QAT path: no config provided, model must have QAT config
|
||||
assert has_qat_config, (
|
||||
"Unsloth: No `torchao_config` provided and model was not trained with `qat_scheme`. "
|
||||
"Either train with `qat_scheme` parameter, or provide a `torchao_config` for "
|
||||
"post-training quantization."
|
||||
)
|
||||
_unsloth_save_torchao_with_attached_config(
|
||||
model = self,
|
||||
save_directory = save_directory,
|
||||
tokenizer = tokenizer,
|
||||
push_to_hub = push_to_hub,
|
||||
token = token,
|
||||
)
|
||||
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ from unsloth_zoo.vision_utils import (
|
|||
UnslothVisionDataCollator,
|
||||
)
|
||||
from unsloth_zoo.hf_utils import get_transformers_model_type
|
||||
from packaging.version import Version
|
||||
from unsloth_zoo.utils import Version
|
||||
import dataclasses
|
||||
|
||||
__all__ = [
|
||||
|
|
@ -315,10 +315,13 @@ def _patch_sft_trainer_auto_packing(trl_module):
|
|||
|
||||
# We also disable vision language models for padding free collators
|
||||
blocked = (
|
||||
data_collator is not None
|
||||
(data_collator is not None)
|
||||
or isinstance(processing_class, ProcessorMixin)
|
||||
or is_vlm
|
||||
or is_unsupported_model
|
||||
or (
|
||||
os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1"
|
||||
) # Disable padding free on forced logits
|
||||
)
|
||||
requested_pack = bool(getattr(config_arg, "packing", False))
|
||||
if blocked:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue