Auto-enable padding-free SFT (#3672)

* implement (sdpa, xformers, fa2) sample packing * attention dispatching * ddp working OOTB with CLI * packed SWA and softcap support * enable batch flattening * LGPL license headers * mask packed sequence boundaries * auto-enable sample packing * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add explicit toggle for sample packing * Add explicit toggle for sample packing * Update __init__.py * Update unsloth/kernels/rope_embedding.py * Update unsloth/kernels/rope_embedding.py * remove grad output clones; restore deleted FastLanguageModel arg * fix * restore rope embedding clones * xformers mask cache * implement (sdpa, xformers, fa2) sample packing * attention dispatching * ddp working OOTB with CLI * packed SWA and softcap support * enable batch flattening * LGPL license headers * mask packed sequence boundaries * auto-enable sample packing * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add explicit toggle for sample packing * Add explicit toggle for sample packing * Update __init__.py * Update unsloth/kernels/rope_embedding.py * Update unsloth/kernels/rope_embedding.py * remove grad output clones; restore deleted FastLanguageModel arg * fix * restore rope embedding clones * xformers mask cache * add back accidental deletion * Update unsloth/kernels/rope_embedding.py Co-authored-by: Daniel Han <danielhanchen@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix merge conflicts * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add **kwargs * add back clobbered * Update rope_embedding.py * Update rope_embedding.py * simplify trl warnings filter * docstring * nit * bugfix * add padding-free seqlen metadata * auto-enable padding free * gemma2 disable * Apply suggestion from @danielhanchen * Update trainer.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update trainer.py * [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> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2025-12-10 06:07:29 -05:00 · 2025-12-10 06:07:29 -05:00 · 89b042f23b
commit 89b042f23b
parent 538d2c0719
4 changed files with 281 additions and 20 deletions
--- a/tests/utils/test_packing.py
+++ b/tests/utils/test_packing.py
@ -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)