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

Merge branch 'main' into nightly

Browse source
This commit is contained in:
Daniel Han 2025-12-10 04:12:19 -08:00
commit 92d84ad6d8
7 changed files with 368 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

52
tests/utils/test_packing.py
View file

@ -16,7 +16,9 @@
from unsloth import FastLanguageModel
from unsloth.utils import attention_dispatch as attention_dispatch_utils
from unsloth.utils.packing import (
configure_padding_free,
configure_sample_packing,
enable_padding_free_metadata,
enable_sample_packing,
mask_packed_sequence_boundaries,
)
@ -150,6 +152,14 @@ def test_configure_sample_packing():
assert config.remove_unused_columns is False
def test_configure_padding_free():
config = SimpleNamespace(remove_unused_columns = True)
configure_padding_free(config)
assert config.padding_free is True
assert config.remove_unused_columns is False
class _DummyChild(torch.nn.Module):
def __init__(self):
super().__init__()
@ -177,6 +187,20 @@ class _DummyTrainer:
)
class _PaddingFreeCollator:
def __init__(self):
self.padding_free = True
self.return_position_ids = False
self.calls = 0
def torch_call(self, examples):
self.calls += 1
return {
"input_ids": torch.tensor([[0]], dtype = torch.long),
"examples_seen": self.calls,
}
def test_enable_sample_packing():
model = _DummyModel()
trainer = _DummyTrainer()
@ -251,6 +275,34 @@ def test_enable_sample_packing_trl_collator(tmp_path):
trainer.accelerator.free_memory()
def test_enable_padding_free_metadata():
model = _DummyModel()
trainer = SimpleNamespace(
args = SimpleNamespace(remove_unused_columns = True),
data_collator = _PaddingFreeCollator(),
)
enable_padding_free_metadata(model, trainer)
assert getattr(model, "_unsloth_allow_packed_overlength") is True
assert getattr(model.child, "_unsloth_allow_packed_overlength") is True
collator = trainer.data_collator
assert collator.return_position_ids is True
assert getattr(collator, "_unsloth_padding_free_lengths_wrapped") is True
examples = [
{"input_ids": [0, 1, 2]},
{"input_ids": [3, 4]},
]
batch = collator.torch_call(examples)
assert torch.equal(
batch["packed_seq_lengths"],
torch.tensor([3, 2], dtype = torch.int32),
)
assert trainer.args.remove_unused_columns is False
def test_packing_sdpa(tmp_path):
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model, batch, trainer, llama_mod = _build_packed_training_setup(tmp_path, device)

10
unsloth/__init__.py
View file

@ -18,10 +18,18 @@ import os, re, subprocess, inspect, functools
import numpy as np
# Fix some issues before importing other packages
from .import_fixes import fix_message_factory_issue
from .import_fixes import (
fix_message_factory_issue,
check_fbgemm_gpu_version,
torchvision_compatibility_check,
)
fix_message_factory_issue()
check_fbgemm_gpu_version()
torchvision_compatibility_check()
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"]

78
unsloth/import_fixes.py
View file

@ -16,11 +16,30 @@ import os
import importlib.util
from pathlib import Path
from importlib.metadata import version as importlib_version
from packaging.version import Version
from packaging.version import Version as TrueVersion
import re
import logging
UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1"
def Version(version):
try:
new_version = str(version)
new_version = re.match(r"[0-9\.]{1,}", new_version)
if new_version is None:
raise Exception(str(e))
new_version = new_version.group(0).rstrip(".")
if new_version != version:
new_version += ".1" # Add .1 for dev / alpha / beta / rc
return TrueVersion(new_version)
except:
from inspect import getframeinfo, stack
caller = getframeinfo(stack()[1][0])
raise RuntimeError(
f"Unsloth: Could not get version for `{version}`\n"\
f"File name = [{caller.filename}] Line number = [{caller.lineno}]"
)
# Ignore logging messages
class HideLoggingMessage(logging.Filter):
@ -239,3 +258,60 @@ def patch_datasets():
f"#### Unsloth: Using `datasets = {str(datasets_version)}` will cause recursion errors.\n"
"Please downgrade datasets to `datasets==4.3.0"
)
def check_fbgemm_gpu_version():
if importlib.util.find_spec("fbgemm_gpu") is None:
return
try:
fbgemm_gpu_version = importlib_version("fbgemm_gpu_genai")
except:
return
# We noticed some SegFault or bad alloc errors on lower versions of fbgemm_gpu.
if Version(fbgemm_gpu_version) < Version("1.4.0"):
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.")
def torchvision_compatibility_check():
if importlib.util.find_spec("torch") is None:
raise ImportError("Unsloth: torch not found. Please install torch first.")
if importlib.util.find_spec("torchvision") is None:
return
torch_version = importlib_version("torch")
torchvision_version = importlib_version("torchvision")
# Torch version -> minimum required torchvision version
# See https://pytorch.org/get-started/previous-versions/
TORCH_TORCHVISION_COMPAT = [
("2.9.0", "0.24.0"),
("2.8.0", "0.23.0"),
("2.7.0", "0.22.0"),
("2.6.0", "0.21.0"),
("2.5.0", "0.20.0"),
("2.4.0", "0.19.0"),
]
required_torchvision = None
for min_torch, min_torchvision in TORCH_TORCHVISION_COMPAT:
if Version(torch_version) >= Version(min_torch):
required_torchvision = min_torchvision
break
if required_torchvision is None:
# Torch version not in compatibility table, skip check
return
if Version(torchvision_version) < Version(required_torchvision):
raise ImportError(
f"Unsloth: torch=={torch_version} requires torchvision>={required_torchvision}, "
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."
)

2
unsloth/models/_utils.py
View file

@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
__version__ = "2025.12.1"
__version__ = "2025.12.2"
__all__ = [
"SUPPORTS_BFLOAT16",

92
unsloth/trainer.py
View file

@ -23,13 +23,19 @@ import trl
import inspect
from trl import SFTTrainer
from . import is_bfloat16_supported
from unsloth.utils import configure_sample_packing, enable_sample_packing
from unsloth.utils import (
configure_padding_free,
configure_sample_packing,
enable_padding_free_metadata,
enable_sample_packing,
)
from unsloth_zoo.training_utils import (
unsloth_train as _unsloth_train,
)
from unsloth_zoo.vision_utils import (
UnslothVisionDataCollator,
)
from unsloth_zoo.hf_utils import get_transformers_model_type
from packaging.version import Version
import dataclasses
@ -47,6 +53,19 @@ _AUTO_PACKING_ENV_DISABLED = os.environ.get(
"UNSLOTH_DISABLE_AUTO_PACKING", ""
).strip().lower() in {"1", "true", "yes", "on"}
_AUTO_PADDING_FREE_ENV_DISABLED = os.environ.get(
"UNSLOTH_DISABLE_AUTO_PADDING_FREE", ""
).strip().lower() in {"1", "true", "yes", "on"}
# [TODO]
# Below cannot work with padding-free
_PADDING_FREE_BLOCK_LIST = {
"gemma2", # - gemma2: Uses slow_attention_softcapping which has torch.compile issues
"gpt_oss", # - gpt_oss: Uses Flex Attention which doesn't handle padding_free correctly
"mistral", # - mistral: Unfortunately I think sliding window attention doesn't work correctly?
}
def _should_auto_pack(config) -> bool:
if config is None or _AUTO_PACKING_ENV_DISABLED:
@ -56,6 +75,14 @@ def _should_auto_pack(config) -> bool:
return not getattr(config, "_unsloth_disable_auto_packing", False)
def _should_auto_padding_free(config) -> bool:
if config is None or _AUTO_PADDING_FREE_ENV_DISABLED:
return False
if getattr(config, "packing", False):
return False
return not getattr(config, "padding_free", False)
def _disable_sample_packing(config):
if config is None:
return
@ -269,11 +296,36 @@ def _patch_sft_trainer_auto_packing(trl_module):
else:
config_arg = kwargs.get("args")
# Check if model type is unsupported for padding_free
model = kwargs.get("model")
is_unsupported_model = False
is_vlm = False
if model is not None:
model_config = getattr(model, "config", None)
if model_config is not None:
model_types = get_transformers_model_type(model_config)
# Blocklist: models that don't work correctly with padding_free
is_unsupported_model = any(
x in PADDING_FREE_BLOCKLIST for x in model_types
)
# Check if VLM
architectures = getattr(model_config, "architectures", None)
if architectures is None:
architectures = []
is_vlm = any(
x.endswith("ForConditionalGeneration") for x in architectures
)
is_vlm = is_vlm or hasattr(model_config, "vision_config")
processing_class = kwargs.get("processing_class") or kwargs.get("tokenizer")
data_collator = kwargs.get("data_collator")
blocked = data_collator is not None or isinstance(
processing_class, ProcessorMixin
# We also disable vision language models for padding free collators
blocked = (
data_collator is not None
or isinstance(processing_class, ProcessorMixin)
or is_vlm
)
if blocked and _should_auto_pack(config_arg):
reason = (
@ -292,6 +344,22 @@ def _patch_sft_trainer_auto_packing(trl_module):
auto_pack_active = True
logger.info("Unsloth: Sample packing auto-enabled for SFTTrainer instance.")
auto_padding_free_active = False
padding_free_requested = getattr(config_arg, "padding_free", None) is True
if not blocked:
if padding_free_requested:
configure_padding_free(config_arg)
elif not is_unsupported_gemma and _should_auto_padding_free(config_arg):
configure_padding_free(config_arg)
auto_padding_free_active = True
logger.info(
"Unsloth: Padding-free batching auto-enabled for SFTTrainer instance."
)
elif is_unsupported_gemma and _should_auto_padding_free(config_arg):
logger.info(
"Unsloth: Padding-free batching auto-disabled for Gemma 2 (requires flash attention)."
)
try:
original_init(self, *args, **kwargs)
except ValueError as exc:
@ -307,11 +375,24 @@ def _patch_sft_trainer_auto_packing(trl_module):
raise
trainer_args = getattr(self, "args", None)
if auto_pack_active and _should_auto_pack(trainer_args):
trainer_packing = bool(trainer_args and getattr(trainer_args, "packing", False))
trainer_padding_free = bool(
trainer_args and getattr(trainer_args, "padding_free", False)
)
if trainer_packing and (auto_pack_active or _should_auto_pack(trainer_args)):
enable_sample_packing(self.model, self)
print(
"🦥 Unsloth: Packing enabled - training is >2x faster and uses less VRAM!"
)
elif trainer_padding_free:
enable_padding_free_metadata(self.model, self)
message = (
"🦥 Unsloth: Padding-free auto-enabled, enabling faster training."
if auto_padding_free_active
else "🦥 Unsloth: Padding-free enabled, enabling faster training."
)
print(message)
sft_trainer.__init__ = new_init
sft_trainer._unsloth_auto_packing_wrapped = True
@ -343,7 +424,6 @@ def _patch_trl_trainer():
except:
continue
if not _AUTO_PACKING_ENV_DISABLED:
_patch_sft_trainer_auto_packing(trl)
_patch_sft_trainer_auto_packing(trl)
trl.__UNSLOTH_BACKWARDS_COMPATIBLE__ = True

11
unsloth/utils/__init__.py
View file

@ -13,7 +13,13 @@
# You should have received a copy of the GNU Lesser General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from .packing import configure_sample_packing, enable_sample_packing
from .packing import (
configure_padding_free,
configure_sample_packing,
enable_padding_free_metadata,
enable_sample_packing,
mark_allow_overlength,
)
from .attention_dispatch import (
AttentionConfig,
AttentionContext,
@ -27,7 +33,10 @@ from .attention_dispatch import (
__all__ = [
"configure_sample_packing",
"configure_padding_free",
"enable_sample_packing",
"enable_padding_free_metadata",
"mark_allow_overlength",
"AttentionConfig",
"AttentionContext",
"FLASH_VARLEN",

146
unsloth/utils/packing.py
View file

@ -84,6 +84,19 @@ def _ensure_trl_warning_filter():
_TRL_FILTER_INSTALLED = True
def mark_allow_overlength(module):
"""Mark a module hierarchy so padding-free batches can exceed max_seq_length."""
if module is None:
return
if hasattr(module, "max_seq_length"):
setattr(module, "_unsloth_allow_packed_overlength", True)
children = getattr(module, "children", None)
if children is None:
return
for child in children():
mark_allow_overlength(child)
def configure_sample_packing(config):
"""Mutate an ``SFTConfig`` so TRL prepares packed batches."""
_ensure_trl_warning_filter()
@ -92,25 +105,37 @@ def configure_sample_packing(config):
setattr(config, "remove_unused_columns", False)
def enable_sample_packing(model, trainer):
def configure_padding_free(config):
"""Mutate an ``SFTConfig`` so TRL enables padding-free batching without packing."""
_ensure_trl_warning_filter()
setattr(config, "padding_free", True)
if hasattr(config, "remove_unused_columns"):
setattr(config, "remove_unused_columns", False)
def enable_sample_packing(
model,
trainer,
*,
sequence_lengths_key: str = "seq_lengths",
) -> None:
"""Enable runtime support for packed batches on an existing trainer."""
if model is None or trainer is None:
raise ValueError("model and trainer must not be None")
def _mark_allow_overlength(module):
if hasattr(module, "max_seq_length"):
setattr(module, "_unsloth_allow_packed_overlength", True)
for child in module.children():
_mark_allow_overlength(child)
mark_allow_overlength(model)
_mark_allow_overlength(model)
if hasattr(trainer, "args") and hasattr(trainer.args, "remove_unused_columns"):
trainer.args.remove_unused_columns = False
collator = getattr(trainer, "data_collator", None)
if (
collator is None
or not hasattr(collator, "torch_call")
or getattr(collator, "_unsloth_packing_wrapped", False)
):
if collator is None or not hasattr(collator, "torch_call"):
return
if getattr(collator, "_unsloth_packing_wrapped", False):
return
if hasattr(collator, "padding_free"):
collator.padding_free = True
if hasattr(collator, "return_position_ids"):
collator.return_position_ids = True
@ -120,18 +145,107 @@ def enable_sample_packing(model, trainer):
batch = original_torch_call(examples)
if examples and isinstance(examples[0], dict):
seq_lengths: list[int] = []
per_example_counts: list[int] = []
for example in examples:
seq_lengths.extend(example["seq_lengths"])
lengths = example.get(sequence_lengths_key)
if isinstance(lengths, Iterable):
numeric_lengths = [int(length) for length in lengths]
seq_lengths.extend(numeric_lengths)
per_example_counts.append(len(numeric_lengths))
else:
per_example_counts.append(0)
if seq_lengths:
batch["packed_seq_lengths"] = torch.tensor(
seq_lengths, dtype = torch.int32
)
position_ids = batch.get("position_ids")
input_ids = batch.get("input_ids")
if position_ids is None and input_ids is not None:
position_ids = torch.zeros_like(
input_ids, dtype = torch.long, device = input_ids.device
)
if position_ids is not None and input_ids is not None:
seq_index = 0
for row_idx, count in enumerate(per_example_counts):
cursor = 0
for _ in range(count):
length = seq_lengths[seq_index]
if length > 0:
position_ids[row_idx, cursor : cursor + length] = (
torch.arange(
length,
dtype = torch.long,
device = position_ids.device,
)
)
cursor += length
seq_index += 1
batch["position_ids"] = position_ids
if "attention_mask" in batch and getattr(
collator, "return_position_ids", False
):
batch.pop("attention_mask")
return batch
collator.torch_call = torch_call_with_lengths
collator._unsloth_packing_wrapped = True
def enable_padding_free_metadata(model, trainer):
"""Inject seq-length metadata when padding-free batching is enabled without packing."""
trainer_args = getattr(trainer, "args", None)
if (
trainer_args is not None
and hasattr(trainer_args, "remove_unused_columns")
and trainer_args.remove_unused_columns
):
trainer_args.remove_unused_columns = False
_ensure_trl_warning_filter()
collator = getattr(trainer, "data_collator", None)
if (
collator is None
or getattr(collator, "_unsloth_padding_free_lengths_wrapped", False)
or not getattr(collator, "padding_free", False)
):
# Nothing to do if there's no collator, we've already wrapped it, or padding-free is off.
return
mark_allow_overlength(model)
if hasattr(collator, "return_position_ids"):
collator.return_position_ids = True
original_torch_call = collator.torch_call
def torch_call_with_padding_free_metadata(examples: Sequence[dict]):
seq_lengths: list[int] = []
if examples and isinstance(examples[0], dict):
for example in examples:
lengths = example.get("seq_lengths")
if lengths is None:
ids = example.get("input_ids")
if ids is None:
continue
lengths = [len(ids)]
example["seq_lengths"] = lengths
seq_lengths.extend(lengths)
batch = original_torch_call(examples)
if seq_lengths:
batch["packed_seq_lengths"] = torch.tensor(
seq_lengths,
dtype = torch.int32,
)
return batch
collator.torch_call = torch_call_with_padding_free_metadata
collator._unsloth_padding_free_lengths_wrapped = True
def get_packed_info_from_kwargs(
kwargs: dict,
device: torch.device,
@ -261,6 +375,12 @@ def mask_packed_sequence_boundaries(
__all__ = [
"configure_sample_packing",
"configure_padding_free",
"enable_sample_packing",
"enable_padding_free_metadata",
"mark_allow_overlength",
"get_packed_info_from_kwargs",
"build_xformers_block_causal_mask",
"build_sdpa_packed_attention_mask",
"mask_packed_sequence_boundaries",
]