Experimental: correct varlen sample packing for hybrid linear-attention models (#7249)

* Fix text-only VLM CPT packing truncation

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Handle streaming vision datasets in packing

* Harden multimodal packing detection

* Preserve safe packing boundaries

* Scope stream packing checks to VLMs

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Narrow VLM packing detection

* Align packing mode and eval safety

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add qwen3_5/qwen3_next to PADDING_FREE_BLOCKLIST to avoid packed-sequence contamination

* Detect hybrid linear-attention models structurally instead of by name for packing guard

* Add experimental varlen packing for hybrid linear-attention models

Feed seq_idx to the causal conv and cu_seqlens to the gated-delta scan so
sample packing / padding-free reset state at sequence boundaries for hybrid
linear-attention models (Qwen3.5, Qwen3-Next). Gated behind
UNSLOTH_EXPERIMENTAL_HYBRID_PACKING and fail-closed: when the flag is off or
the accelerated kernels (causal_conv1d + fla) are unavailable, the guard keeps
these models on the padded path.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden hybrid linear-attention varlen packing shim

Make patch_hybrid_linear_attention_varlen robust across transformers 4.57.6
through 5.x and TRL 0.22.2 through 1.x, following the import_fixes.py style:

- Read UNSLOTH_EXPERIMENTAL_HYBRID_PACKING at call time so the flag takes effect
  when set after importing unsloth.
- Idempotent: repeat calls on a patched model return True without re-validating
  the wrappers or double-wrapping; signatures are checked on captured originals.
- Prefer the authoritative packed_seq_lengths (via get_packed_info_from_kwargs)
  over position_ids resets, handling pad_to_multiple_of trailing tokens.
- Suppress injection for cached forwards (use_cache / past_key_values) so
  generation and eval are left on the untouched decode path.
- Validate every gated-delta module before mutating any (transactional).
- Bind position_ids / use_cache from both positional and keyword args.
- Verify dispatch at runtime (Unsloth wraps each module forward, so the mixer
  source is not statically inspectable) and warn once if the shim is never hit.
- Emit one deduped diagnostic on each fail-closed path.

Add CPU unit tests covering the hybrid guard detection, the boundary builders,
and the shim (fail-closed, active, idempotent, cached no-op, runtime handshake).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Abort hybrid packing when the varlen shim is not fully dispatched

The runtime handshake used a single per-module hit flag written by both the conv
and scan wrappers, so a partial dispatch (only one kernel routed through
self.<kernel>) passed the any() check and trained on contaminated data, and a
missing dispatch only logged a warning. Track conv and scan dispatch separately,
require both on every gated-delta module on the first packed forward, and raise
before loss/backward when either is missing (the batch is already flattened, so
there is no padded recovery at that point). Also skip an empty packed_seq_lengths
before it reaches max(), and document the position_ids fallback's left-pad
assumption.

Add tests for no-dispatch and partial (conv-only / scan-only) abort, the
packed_seq_lengths preference over a competing position_ids, MRoPE 3D position
ids, and the pad_to_multiple_of trailing-segment path through the metadata builder.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Import the hybrid packing patch from its submodule to satisfy the import-hoist lint

* Fail closed for hybrid packing on encoder-decoder, chunked-loss, and string-name models

The varlen shim only helps decoder-only hybrid models that run their mixer
through self.<kernel> on a live nn.Module forward. Three cases slipped past
the guard:

- Encoder-decoder configs (is_encoder_decoder) reached the packing path even
  though flattening a cross-attention batch is unsound. Block them explicitly.
- TRL's chunked_nll loss (the 1.x default) calls the backbone directly and
  bypasses model.forward, so the per-instance forward wrapper that refreshes
  the varlen stash never runs. Detect that path and keep the model padded.
- A string model_name reaches the trainer before the module exists, so the
  instance shim has nothing to patch. Resolve the config up front and keep
  string hybrids on the padded path.

Adds encoder-decoder / decoder-only / chunked-loss / string-model tests.

* Harden the SFT source-injection replacements and forward auth args for string models

The wrapped-packing injection rewrote the sourced unsloth_zoo sft_prepare_dataset
with str.replace anchored on the exact 'All Unsloth Zoo code licensed under
LGPLv3' comment. str.replace never raises on a missing anchor, so a supported
newer unsloth_zoo (the dependency is only lower-bounded) that moved that header
would silently drop the setup while the truncation and pack_dataset edits still
referenced _unsloth_wrapped_packing / _inspect, raising NameError on every SFT
dataset preparation.

- Install the setup at the sft_prepare_dataset signature via re.subn (a structural
  anchor that always exists) and raise if even that is missing.
- Route the remaining edits through a _require_replace helper that fails loudly on a
  missing required anchor (or warns once for an optional one), formalizing the
  verify-then-replace idiom the DPO patchers in this file already use.
- Reuse the guarded _unsloth_pack_has_strategy at the pack_dataset call instead of
  re-calling inspect.signature(pack_dataset) unguarded, so a non-introspectable
  pack_dataset cannot crash there after the setup already handled it.
- _resolve_string_model_config now forwards token / use_auth_token / cache_dir /
  code_revision, so a private hybrid resolves its config instead of falling through
  as non-hybrid and enabling packing without the varlen shim.

Adds regression tests for the drift-resistant injection, the helper, and the
string-model auth forwarding.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Honor top-level SFTConfig.trust_remote_code when resolving a string model

TRL merges the top-level args.trust_remote_code into the load via
model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) before
create_model_from_path, so a remote-code hybrid is commonly set with
SFTConfig(trust_remote_code=True) rather than inside model_init_kwargs. The config
probe only read model_init_kwargs, so AutoConfig could fail for such a model, leave
model_config None, and let the guard treat it as non-hybrid, enabling packing
without the varlen shim. Mirror TRL's setdefault (model_init_kwargs wins).

* Tighten hybrid-packing comments for concision

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: alkinun <alkinunl@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Etherl <61019402+Etherll@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-20 00:57:02 -07:00 committed by GitHub
commit 07272b9278
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 984 additions and 80 deletions

View file

@ -15,6 +15,7 @@
from unsloth import FastLanguageModel
import unsloth.trainer as trainer_module
import unsloth.utils.packing as packing_module
from unsloth.utils import attention_dispatch as attention_dispatch_utils
from unsloth.utils.packing import (
configure_padding_free,
@ -22,6 +23,7 @@ from unsloth.utils.packing import (
enable_padding_free_metadata,
enable_sample_packing,
mask_packed_sequence_boundaries,
patch_hybrid_linear_attention_varlen,
)
from contextlib import ExitStack
@ -161,6 +163,327 @@ def test_configure_padding_free():
assert config.remove_unused_columns is False
# --- Hybrid linear-attention guard + varlen shim (PR #7211 / #7249) ---------------
def _hybrid_config_model():
# Qwen3.5 / Qwen3-Next style: explicit linear_attention layer schedule.
return SimpleNamespace(
config = SimpleNamespace(layer_types = ["linear_attention", "full_attention"])
)
def _gemma3_model():
# Has layer_types but no linear_attention -> must NOT be flagged as hybrid.
return SimpleNamespace(
config = SimpleNamespace(
model_type = "gemma3", layer_types = ["sliding_attention", "full_attention"]
),
)
def _dense_qwen3_model():
return SimpleNamespace(
config = SimpleNamespace(model_type = "qwen3", architectures = ["Qwen3ForCausalLM"])
)
class _FakeGatedDeltaNet(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1d = torch.nn.Conv1d(4, 4, 3, groups = 4)
self.A_log = torch.nn.Parameter(torch.zeros(4))
def forward(self, hidden_states, **kwargs): # dispatch through self.<kernel>
return self.chunk_gated_delta_rule(self.causal_conv1d_fn(hidden_states))
class _FakeHybridModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.config = SimpleNamespace() # no markers -> forces module-level detection
self.linear_attn = _FakeGatedDeltaNet()
def test_is_hybrid_linear_attention_detects_and_excludes():
is_hybrid = trainer_module._is_hybrid_linear_attention_model
assert is_hybrid(_hybrid_config_model()) is True
assert is_hybrid(_FakeHybridModel()) is True # module-structural evidence
assert is_hybrid(_text_model()) is False # Llama
assert is_hybrid(_gemma3_model()) is False # layer_types without linear_attention
assert is_hybrid(_dense_qwen3_model()) is False # dense Qwen3
assert is_hybrid(None) is False
def test_varlen_from_position_ids():
cu, seq_idx = packing_module._varlen_from_position_ids(torch.tensor([[0, 1, 0, 0, 1, 2]]))
assert cu.tolist() == [0, 2, 3, 6]
assert seq_idx.tolist() == [[0, 0, 1, 2, 2, 2]]
assert (
packing_module._varlen_from_position_ids(torch.tensor([[0, 1, 2, 3]])) is None
) # single sequence
assert packing_module._varlen_from_position_ids(torch.tensor([[1, 2, 3]])) is None # first != 0
assert (
packing_module._varlen_from_position_ids(torch.tensor([[0, 1], [0, 1]])) is None
) # normal 2-row batch
assert packing_module._varlen_from_position_ids(None) is None
def test_seq_idx_from_cu_seqlens_handles_trailing_pad():
cu = torch.tensor([0, 2, 5], dtype = torch.int32)
boundaries, seq_idx = packing_module._seq_idx_from_cu_seqlens(cu, total = 8) # pad_to_multiple_of
assert boundaries.tolist() == [0, 2, 5, 8]
assert seq_idx.tolist() == [[0, 0, 1, 1, 1, 2, 2, 2]]
boundaries2, _ = packing_module._seq_idx_from_cu_seqlens(cu, total = 5) # exact fit
assert boundaries2.tolist() == [0, 2, 5]
assert (
packing_module._seq_idx_from_cu_seqlens(torch.tensor([1, 2], dtype = torch.int32), total = 2)
is None
)
assert packing_module._seq_idx_from_cu_seqlens(cu, total = 3) is None # boundaries exceed total
def test_hybrid_varlen_metadata_prefers_packed_seq_lengths():
# A competing position_ids would segment [0, 3, 6]; packed_seq_lengths must win.
kwargs = {
"input_ids": torch.zeros(1, 6, dtype = torch.long),
"packed_seq_lengths": torch.tensor([2, 1, 3], dtype = torch.int32),
"position_ids": torch.tensor([[0, 1, 2, 0, 1, 2]]),
}
cu, seq_idx = packing_module._hybrid_varlen_metadata(kwargs)
assert cu.tolist() == [0, 2, 3, 6]
assert seq_idx.tolist() == [[0, 0, 1, 2, 2, 2]]
def test_hybrid_varlen_metadata_suppressed_when_cached():
base = {
"input_ids": torch.zeros(1, 6, dtype = torch.long),
"packed_seq_lengths": torch.tensor([2, 1, 3], dtype = torch.int32),
}
assert packing_module._hybrid_varlen_metadata({**base, "use_cache": True}) is None
assert packing_module._hybrid_varlen_metadata({**base, "past_key_values": object()}) is None
def test_hybrid_varlen_metadata_none_for_plain_batch():
kwargs = {
"input_ids": torch.zeros(1, 4, dtype = torch.long),
"position_ids": torch.tensor([[0, 1, 2, 3]]),
}
assert packing_module._hybrid_varlen_metadata(kwargs) is None
def _make_fake_kernels():
def causal_conv1d_fn(
x,
weight = None,
bias = None,
activation = None,
seq_idx = None,
):
causal_conv1d_fn.calls.append(seq_idx)
return x
causal_conv1d_fn.calls = []
def chunk_gated_delta_rule(
q,
k = None,
v = None,
cu_seqlens = None,
**kw,
):
chunk_gated_delta_rule.calls.append(cu_seqlens)
return q
chunk_gated_delta_rule.calls = []
return causal_conv1d_fn, chunk_gated_delta_rule
class _ShimGatedDeltaNet(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1d = torch.nn.Conv1d(4, 4, 3, groups = 4)
self.causal_conv1d_fn, self.chunk_gated_delta_rule = _make_fake_kernels()
def forward(self, hidden_states, **kwargs):
return self.chunk_gated_delta_rule(self.causal_conv1d_fn(hidden_states))
class _ShimHybridModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.config = SimpleNamespace(layer_types = ["linear_attention", "full_attention"])
self.linear_attn = _ShimGatedDeltaNet()
def forward(
self,
input_ids = None,
position_ids = None,
packed_seq_lengths = None,
use_cache = None,
**kwargs,
):
return self.linear_attn(input_ids.float())
def test_patch_hybrid_varlen_flag_off(monkeypatch):
monkeypatch.delenv("UNSLOTH_EXPERIMENTAL_HYBRID_PACKING", raising = False)
model = _ShimHybridModel()
assert patch_hybrid_linear_attention_varlen(model) is False
assert not getattr(model, "_unsloth_varlen_forward_wrapped", False)
def test_patch_hybrid_varlen_active_and_idempotent(monkeypatch):
monkeypatch.setenv("UNSLOTH_EXPERIMENTAL_HYBRID_PACKING", "1")
model = _ShimHybridModel()
conv_orig, scan_orig = (
model.linear_attn.causal_conv1d_fn,
model.linear_attn.chunk_gated_delta_rule,
)
assert patch_hybrid_linear_attention_varlen(model) is True
assert model._unsloth_varlen_forward_wrapped is True
assert model.linear_attn._unsloth_varlen_wrapped is True
assert patch_hybrid_linear_attention_varlen(model) is True # idempotent, no double-wrap
conv_orig.calls.clear()
scan_orig.calls.clear()
packing_module._HYBRID_WARNED.clear()
ids = torch.zeros(1, 6, dtype = torch.long)
model(
input_ids = ids,
packed_seq_lengths = torch.tensor([2, 1, 3], dtype = torch.int32),
use_cache = False,
)
assert conv_orig.calls[-1] is not None # seq_idx injected
assert scan_orig.calls[-1].tolist() == [0, 2, 3, 6] # cu_seqlens injected
assert not packing_module._HYBRID_WARNED # handshake passed, no rejection
conv_orig.calls.clear()
scan_orig.calls.clear()
model(
input_ids = ids, packed_seq_lengths = torch.tensor([2, 1, 3], dtype = torch.int32), use_cache = True
)
assert conv_orig.calls[-1] is None # cached forward -> no injection
assert scan_orig.calls[-1] is None
def test_patch_hybrid_varlen_torch_fallback_fail_closed(monkeypatch):
monkeypatch.setenv("UNSLOTH_EXPERIMENTAL_HYBRID_PACKING", "1")
model = _ShimHybridModel()
def torch_chunk_gated_delta_rule(
q,
cu_seqlens = None,
**kw,
):
return q
model.linear_attn.chunk_gated_delta_rule = torch_chunk_gated_delta_rule
assert patch_hybrid_linear_attention_varlen(model) is False
assert not getattr(model, "_unsloth_varlen_forward_wrapped", False)
def test_patch_hybrid_varlen_bad_signature_fail_closed(monkeypatch):
monkeypatch.setenv("UNSLOTH_EXPERIMENTAL_HYBRID_PACKING", "1")
model = _ShimHybridModel()
def scan_no_cu(q, **kw): # missing cu_seqlens
return q
model.linear_attn.chunk_gated_delta_rule = scan_no_cu
assert patch_hybrid_linear_attention_varlen(model) is False
def _hybrid_model_with_gdn(gdn_forward):
# Build a fake hybrid model whose gated-delta mixer forward is `gdn_forward`.
class _GatedDeltaNet(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1d = torch.nn.Conv1d(4, 4, 3, groups = 4)
self.causal_conv1d_fn, self.chunk_gated_delta_rule = _make_fake_kernels()
forward = gdn_forward
class _Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.config = SimpleNamespace(layer_types = ["linear_attention", "full_attention"])
self.linear_attn = _GatedDeltaNet()
def forward(
self,
input_ids = None,
packed_seq_lengths = None,
use_cache = None,
**kwargs,
):
return self.linear_attn(input_ids.float())
return _Model()
def test_patch_hybrid_varlen_no_dispatch_aborts(monkeypatch):
# Dispatch is verified at runtime, not statically. A mixer that never calls
# self.<kernel> installs the shim, but the first packed forward aborts (both
# boundary kernels are load-bearing).
monkeypatch.setenv("UNSLOTH_EXPERIMENTAL_HYBRID_PACKING", "1")
model = _hybrid_model_with_gdn(lambda self, hidden_states, **kw: hidden_states)
assert patch_hybrid_linear_attention_varlen(model) is True # kernels valid -> installs
with pytest.raises(RuntimeError, match = "both invoked"):
model(
input_ids = torch.zeros(1, 6),
packed_seq_lengths = torch.tensor([2, 1, 3], dtype = torch.int32),
use_cache = False,
)
def test_patch_hybrid_varlen_partial_dispatch_aborts(monkeypatch):
# Only the conv fires; the scan would leak state. Both must be invoked, so abort.
monkeypatch.setenv("UNSLOTH_EXPERIMENTAL_HYBRID_PACKING", "1")
conv_only = _hybrid_model_with_gdn(
lambda self, hidden_states, **kw: self.causal_conv1d_fn(hidden_states)
)
assert patch_hybrid_linear_attention_varlen(conv_only) is True
with pytest.raises(RuntimeError, match = "both invoked"):
conv_only(
input_ids = torch.zeros(1, 6),
packed_seq_lengths = torch.tensor([2, 1, 3], dtype = torch.int32),
use_cache = False,
)
scan_only = _hybrid_model_with_gdn(
lambda self, hidden_states, **kw: self.chunk_gated_delta_rule(hidden_states)
)
assert patch_hybrid_linear_attention_varlen(scan_only) is True
with pytest.raises(RuntimeError, match = "both invoked"):
scan_only(
input_ids = torch.zeros(1, 6),
packed_seq_lengths = torch.tensor([2, 1, 3], dtype = torch.int32),
use_cache = False,
)
def test_varlen_from_position_ids_mrope_3d():
pos = (
torch.tensor([[0, 1, 0, 0, 1, 2]]).unsqueeze(0).expand(3, 1, 6).clone()
) # [3,1,T] text plane
cu, seq_idx = packing_module._varlen_from_position_ids(pos)
assert cu.tolist() == [0, 2, 3, 6]
assert seq_idx.tolist() == [[0, 0, 1, 2, 2, 2]]
def test_hybrid_varlen_metadata_trailing_pad():
# packed_seq_lengths sum to 6 but the flattened input is 8 (pad_to_multiple_of).
kwargs = {
"input_ids": torch.zeros(1, 8, dtype = torch.long),
"packed_seq_lengths": torch.tensor([2, 1, 3], dtype = torch.int32),
}
cu, seq_idx = packing_module._hybrid_varlen_metadata(kwargs)
assert cu.tolist() == [0, 2, 3, 6, 8]
assert seq_idx.tolist() == [[0, 0, 1, 2, 2, 2, 3, 3]]
def _patch_fake_sft_trainer():
class FakeSFTTrainer:
def __init__(self, *args, **kwargs):
@ -245,29 +568,101 @@ def test_vlm_without_processing_class_still_disables_packing():
("t5", "T5ForConditionalGeneration"),
("bart", "BartForConditionalGeneration"),
("whisper", "WhisperForConditionalGeneration"),
("csm", "CsmForConditionalGeneration"),
),
)
def test_nonvision_conditional_generation_keeps_packing(model_type, architecture):
def test_encoder_decoder_disables_packing(model_type, architecture):
# Text-only encoder-decoder models are not VLMs, but their bidirectional encoder
# attends across concatenated samples once padding-free drops attention_mask.
fake_trainer = _patch_fake_sft_trainer()
config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
model = SimpleNamespace(
config = SimpleNamespace(model_type = model_type, architectures = [architecture]),
config = SimpleNamespace(
model_type = model_type,
architectures = [architecture],
is_encoder_decoder = True,
),
max_seq_length = 16,
)
trainer = fake_trainer(
model,
config,
None,
Dataset.from_dict({"text": ["text-only sample"]}),
trainer = fake_trainer(model, config, None, Dataset.from_dict({"text": ["text-only sample"]}))
assert config.packing is False
assert config.padding_free is False
def test_decoder_only_conditional_generation_keeps_packing():
# CSM is decoder-only despite the ForConditionalGeneration name -> packing stays on.
fake_trainer = _patch_fake_sft_trainer()
config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
model = SimpleNamespace(
config = SimpleNamespace(
model_type = "csm",
architectures = ["CsmForConditionalGeneration"],
is_encoder_decoder = False,
),
max_seq_length = 16,
)
trainer = fake_trainer(model, config, None, Dataset.from_dict({"text": ["text-only sample"]}))
assert config.packing is True
assert config.padding_free is True
assert trainer.model._unsloth_allow_packed_overlength is True
def _hybrid_trainer_model():
return SimpleNamespace(
config = SimpleNamespace(
model_type = "qwen3_next",
architectures = ["Qwen3NextForCausalLM"],
layer_types = ["linear_attention", "full_attention"],
),
max_seq_length = 16,
)
def test_hybrid_varlen_active_enables_packing(monkeypatch):
# Baseline: shim active + no forward bypass -> hybrid packing is allowed.
monkeypatch.setattr(trainer_module, "_chunked_loss_bypasses_forward", lambda config: False)
monkeypatch.setattr(trainer_module, "patch_hybrid_linear_attention_varlen", lambda model: True)
fake_trainer = _patch_fake_sft_trainer()
config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
fake_trainer(_hybrid_trainer_model(), config, None, Dataset.from_dict({"text": ["x"]}))
assert config.packing is True
assert config.padding_free is True
def test_hybrid_chunked_loss_stays_on_padded_path(monkeypatch):
# TRL's chunked-loss forward bypass leaves the varlen shim off -> block packing.
monkeypatch.setattr(trainer_module, "_chunked_loss_bypasses_forward", lambda config: True)
monkeypatch.setattr(trainer_module, "patch_hybrid_linear_attention_varlen", lambda model: True)
fake_trainer = _patch_fake_sft_trainer()
config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
fake_trainer(_hybrid_trainer_model(), config, None, Dataset.from_dict({"text": ["x"]}))
assert config.packing is False
assert config.padding_free is False
def test_string_hybrid_model_disables_packing(monkeypatch):
# A string model= is materialized after init; a hybrid string is blocked because the
# shim cannot patch a not-yet-built model.
monkeypatch.setattr(
trainer_module,
"_resolve_string_model_config",
lambda name, cfg: SimpleNamespace(
model_type = "qwen3_next",
architectures = ["Qwen3NextForCausalLM"],
layer_types = ["linear_attention", "full_attention"],
),
)
monkeypatch.setattr(trainer_module, "patch_hybrid_linear_attention_varlen", lambda model: True)
fake_trainer = _patch_fake_sft_trainer()
config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
fake_trainer("Qwen/Qwen3-Next-80B-A3B", config, None, Dataset.from_dict({"text": ["x"]}))
assert config.packing is False
assert config.padding_free is False
def test_vlm_vision_dataset_still_disables_packing():
fake_trainer = _patch_fake_sft_trainer()
config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
@ -486,49 +881,6 @@ def test_wrapped_packing_preserves_overlength_tokens(monkeypatch, legacy_api):
assert all(len(input_ids) <= args.max_length for input_ids in packed_ids)
# Named to match the unsloth_zoo helper: sft_trainer_prepare_dataset sources it by
# name and renames "def sft_prepare_dataset" -> "def _prepare_dataset". This fixture
# deliberately omits the "All Unsloth Zoo code licensed under LGPLv3" header to emulate
# a newer, compatible Zoo whose header moved (the dependency is only lower-bounded).
def sft_prepare_dataset(
self, dataset, processing_class, args, packing, formatting_func, dataset_text_field
):
do_truncation = True
# Mirror the Zoo call so the "truncation = do_truncation," injection anchor
# survives formatting (a bare tuple assignment gets rewritten to a paren form).
dataset = processing_class(
dataset,
truncation = do_truncation,
)
return dataset
def test_wrapped_packing_setup_survives_missing_zoo_header(monkeypatch):
# Regression: the wrapped-packing setup used to anchor on the Zoo license comment,
# so a header change made it a no-op while the truncation reference still landed,
# NameError-ing every SFT dataset preparation. It must now install via the
# signature and always precede the reference.
import ast
import textwrap
import unsloth.models.rl_replacements as rlr
monkeypatch.setitem(rlr.RL_REPLACEMENTS, "sft_prepare_dataset", sft_prepare_dataset)
source = (
"def _prepare_dataset(self, dataset, processing_class, args, packing, "
"formatting_func, dataset_text_field):\n return dataset\n"
)
patched = rlr.sft_trainer_prepare_dataset("_prepare_dataset", source)
assert "_unsloth_wrapped_packing = packing" in patched
assert "import inspect as _inspect" in patched
assert "not _unsloth_wrapped_packing" in patched
assert patched.index("_unsloth_wrapped_packing = packing") < patched.index(
"truncation = do_truncation and not _unsloth_wrapped_packing"
)
ast.parse(textwrap.dedent(patched))
class _DummyChild(torch.nn.Module):
def __init__(self):
super().__init__()
@ -759,3 +1111,128 @@ def test_packing_sdpa(tmp_path):
if hasattr(trainer, "accelerator"):
trainer.accelerator.free_memory()
# --- wrapped-packing source-injection robustness (reviewer.py / fork findings) --------
# fmt: off
# Named to match the unsloth_zoo helper (sourced by name, "def sft_prepare_dataset" ->
# "def _prepare_dataset"). Deliberately OMITS the "licensed under LGPLv3" header to
# emulate a newer Zoo whose header moved (dependency is only lower-bounded). Source only.
def sft_prepare_dataset(
self, dataset, processing_class, args, packing, formatting_func, dataset_text_field
):
do_truncation = True
max_seq_length = 4
used_column_names = ["text"]
map_kwargs = {}
dataset = processing_class(dataset, truncation = do_truncation,)
if do_truncation and max_seq_length > 0:
pass
if packing:
dataset = pack_dataset(
dataset.select_columns(used_column_names),
max_seq_length,
getattr(args, "packing_strategy", "bfd"),
map_kwargs,
)
return dataset
# fmt: on
def test_wrapped_packing_injection_is_drift_resistant(monkeypatch):
# Regression: the setup used to anchor on the Zoo license comment, so a header
# change silently no-op'd it while the truncation/pack edits still referenced its
# variables -> NameError on every SFT prep. It must now install via the signature
# before those references, and the pack edit must reuse the guarded
# _unsloth_pack_has_strategy instead of re-calling _inspect.signature(pack_dataset).
import ast
import textwrap
import unsloth.models.rl_replacements as rlr
monkeypatch.setitem(rlr.RL_REPLACEMENTS, "sft_prepare_dataset", sft_prepare_dataset)
source = (
"def _prepare_dataset(self, dataset, processing_class, args, packing, "
"formatting_func, dataset_text_field):\n return dataset\n"
)
patched = rlr.sft_trainer_prepare_dataset("_prepare_dataset", source)
# setup installed despite the missing header, and before it is referenced
assert "_unsloth_wrapped_packing = packing" in patched
assert "import inspect as _inspect" in patched
assert patched.index("_unsloth_wrapped_packing = packing") < patched.index(
"truncation = do_truncation and not _unsloth_wrapped_packing"
)
# the pack edit reuses the guarded flag (signature inspected exactly once, in setup)
assert "if _unsloth_pack_has_strategy:" in patched
assert patched.count("_inspect.signature(pack_dataset)") == 1
ast.parse(textwrap.dedent(patched))
def test_require_replace_raises_on_missing_anchor():
from unsloth.models.rl_replacements import _require_replace
assert _require_replace("abc", "b", "B") == "aBc"
with pytest.raises(RuntimeError):
_require_replace("abc", "z", "Z", where = "unit test")
# an optional edit warns once and returns the source unchanged (no dangling ref)
assert _require_replace("abc", "z", "Z", required = False, where = "optional") == "abc"
def test_resolve_string_model_config_forwards_token(monkeypatch):
import transformers
captured = {}
class _FakeAutoConfig:
@staticmethod
def from_pretrained(name, **kwargs):
captured.update(kwargs)
return SimpleNamespace(is_encoder_decoder = False)
monkeypatch.setattr(transformers, "AutoConfig", _FakeAutoConfig)
config_arg = SimpleNamespace(
model_init_kwargs = {
"token": "hf_secret",
"trust_remote_code": True,
"cache_dir": "/tmp/cache",
"torch_dtype": "bfloat16", # not a config arg -> must NOT be forwarded
}
)
result = trainer_module._resolve_string_model_config("org/private-hybrid", config_arg)
assert result is not None
assert captured.get("token") == "hf_secret"
assert captured.get("trust_remote_code") is True
assert captured.get("cache_dir") == "/tmp/cache"
assert "torch_dtype" not in captured
def test_resolve_string_model_config_merges_top_level_trust_remote_code(monkeypatch):
import transformers
captured = {}
class _FakeAutoConfig:
@staticmethod
def from_pretrained(name, **kwargs):
captured.update(kwargs)
return SimpleNamespace(is_encoder_decoder = False)
monkeypatch.setattr(transformers, "AutoConfig", _FakeAutoConfig)
# SFTConfig(trust_remote_code=True) with no model_init_kwargs entry is honored
config_arg = SimpleNamespace(model_init_kwargs = {}, trust_remote_code = True)
trainer_module._resolve_string_model_config("org/remote-hybrid", config_arg)
assert captured.get("trust_remote_code") is True
# model_init_kwargs wins over the top-level flag (mirrors TRL's setdefault)
captured.clear()
config_arg = SimpleNamespace(
model_init_kwargs = {"trust_remote_code": False}, trust_remote_code = True
)
trainer_module._resolve_string_model_config("org/remote-hybrid", config_arg)
assert captured.get("trust_remote_code") is False

View file

@ -437,6 +437,52 @@ RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_compute_loss_liger)
RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_data_collator_vision_keys)
_WRAPPED_PACKING_SETUP = (
" import inspect as _inspect\n"
" try:\n"
' _unsloth_pack_has_strategy = "strategy" in _inspect.signature(pack_dataset).parameters\n'
" except Exception:\n"
" _unsloth_pack_has_strategy = True\n"
" _unsloth_wrapped_packing = packing and (\n"
' getattr(args, "packing_strategy", None) == "wrapped"\n'
" or not _unsloth_pack_has_strategy\n"
" )\n"
)
_WARNED_MISSING_ANCHORS = set()
def _require_replace(
function,
old,
new,
*,
count = 1,
required = True,
where = "",
):
"""str.replace that never silently no-ops a load-bearing source edit.
Plain str.replace returns the source unchanged when the anchor is absent, so a
drifted anchor in a newer TRL / unsloth_zoo would skip the edit while later edits
still reference helper variables it should have introduced (NameError at runtime).
Fail loudly for a required edit, warn once and skip for an optional one, so a
drifted source can never corrupt the patched function silently.
"""
if old not in function:
detail = f" ({where})" if where else ""
if required:
raise RuntimeError(
f"Unsloth: source anchor not found{detail}; the patched function is out "
"of sync with this TRL / unsloth_zoo version. Please file a bug report."
)
if where not in _WARNED_MISSING_ANCHORS:
_WARNED_MISSING_ANCHORS.add(where)
logger.warning(f"Unsloth: skipped an optional source edit{detail} (anchor not found).")
return function
return function.replace(old, new, count)
# Fix tokenizer double BOS
def sft_trainer_prepare_dataset(function_name, function):
if function_name != "_prepare_non_packed_dataloader" and function_name != "_prepare_dataset":
@ -454,27 +500,14 @@ def sft_trainer_prepare_dataset(function_name, function):
if matched:
# Use fast version!
function = inspect.getsource(fast_sft_prepare_dataset)
# why: install the wrapped-packing setup (and the `_inspect` import the
# truncation / pack_dataset rewrites below depend on) at the function
# signature, a structural anchor that always exists, rather than the
# unsloth_zoo license-comment line. That header is only lower-bounded, so a
# newer Zoo may move or drop it; anchoring there let the setup silently
# no-op while the references still landed, NameError-ing every SFT dataset
# preparation. Fail loudly if even the signature cannot be located.
_wrapped_packing_setup = (
" import inspect as _inspect\n"
" try:\n"
' _unsloth_pack_has_strategy = "strategy" in _inspect.signature(pack_dataset).parameters\n'
" except Exception:\n"
" _unsloth_pack_has_strategy = True\n"
" _unsloth_wrapped_packing = packing and (\n"
' getattr(args, "packing_strategy", None) == "wrapped"\n'
" or not _unsloth_pack_has_strategy\n"
" )\n"
)
# why: anchor the wrapped-packing setup on the function signature -- a
# structural anchor that always exists -- not the unsloth_zoo license comment,
# which is only lower-bounded and a newer Zoo may move or drop. Anchoring there
# let the setup silently no-op while edits below referenced its variables,
# NameError-ing every SFT dataset prep. Fail loudly if the signature is missing.
function, _n_setup = re.subn(
r"(def sft_prepare_dataset\s*\(.*?\)\s*(?:->[^:\n]*)?:[ \t]*\n)",
lambda match: match.group(1) + _wrapped_packing_setup,
lambda match: match.group(1) + _WRAPPED_PACKING_SETUP,
function,
count = 1,
flags = re.DOTALL,
@ -484,15 +517,25 @@ def sft_trainer_prepare_dataset(function_name, function):
"Unsloth: failed to install wrapped-packing support into "
"sft_prepare_dataset (signature not found); please file a bug report."
)
function = function.replace(
# why: route each edit through _require_replace so a drifted anchor fails
# loudly instead of leaving a dangling reference to the setup variables.
function = _require_replace(
function,
"truncation = do_truncation,",
"truncation = do_truncation and not _unsloth_wrapped_packing,",
where = "sft_prepare_dataset truncation flag",
)
function = function.replace(
function = _require_replace(
function,
"if do_truncation and max_seq_length > 0:",
"if do_truncation and not _unsloth_wrapped_packing and max_seq_length > 0:",
where = "sft_prepare_dataset truncation guard",
)
function = function.replace(
# why: reuse the guarded _unsloth_pack_has_strategy from the setup instead of
# re-calling _inspect.signature(pack_dataset) here -- the setup wraps that call
# in try/except, so a non-introspectable pack_dataset must not crash here.
function = _require_replace(
function,
"""dataset = pack_dataset(
dataset.select_columns(used_column_names),
max_seq_length,
@ -500,13 +543,14 @@ def sft_trainer_prepare_dataset(function_name, function):
map_kwargs,
)""",
"""_pack_kwargs = {"map_kwargs": map_kwargs}
if "strategy" in _inspect.signature(pack_dataset).parameters:
if _unsloth_pack_has_strategy:
_pack_kwargs["strategy"] = getattr(args, "packing_strategy", "bfd")
dataset = pack_dataset(
dataset.select_columns(used_column_names),
max_seq_length,
**_pack_kwargs,
)""",
where = "sft_prepare_dataset pack_dataset call",
)
function = function.split("\n")
function = "\n".join(" " * 4 + x for x in function)

View file

@ -17,6 +17,7 @@ import os
import psutil
import warnings
from dataclasses import dataclass, field
from types import SimpleNamespace
from typing import Optional, List
from functools import wraps
@ -32,6 +33,7 @@ from unsloth.utils import (
enable_padding_free_metadata,
enable_sample_packing,
)
from unsloth.utils.packing import patch_hybrid_linear_attention_varlen
from unsloth_zoo.training_utils import (
unsloth_train as _unsloth_train,
)
@ -101,9 +103,9 @@ PADDING_FREE_BLOCKLIST = {
"gpt_oss", # - gpt_oss: Uses Flex Attention which doesn't handle padding_free correctly
}
# Hybrid linear-attention / state-space models (Qwen3.5, Qwen3-Next, ...) carry a
# recurrent gated-delta state plus a causal conv1d. Sample packing / padding-free
# flattens the batch, so those ops leak state across sequence boundaries. Detected
# structurally by _is_hybrid_linear_attention_model rather than by model name.
# recurrent gated-delta state plus a causal conv1d that leak across sequence
# boundaries once packing flattens the batch. Detected structurally by
# _is_hybrid_linear_attention_model, not by model name.
def _should_pack(config) -> bool:
@ -267,6 +269,57 @@ def _is_hybrid_linear_attention_model(model) -> bool:
return False
def _resolve_string_model_config(model_name, config_arg):
"""TRL materializes a string ``model=`` inside ``__init__``; resolve its config
up front so the packing guards run before the dataset is packed. Best-effort:
returns None if the config cannot be loaded."""
try:
from transformers import AutoConfig
init_kwargs = getattr(config_arg, "model_init_kwargs", None) or {}
# why: forward auth + cache args too. Dropping token/use_auth_token made a
# private hybrid fail to load (resolve as None) -> treated as non-hybrid ->
# packing enabled without the shim even though TRL later loads it with the token.
forward = {
key: init_kwargs[key]
for key in (
"trust_remote_code",
"revision",
"subfolder",
"token",
"use_auth_token",
"cache_dir",
"code_revision",
)
if key in init_kwargs
}
# why: TRL merges top-level args.trust_remote_code into the load via setdefault
# before create_model_from_path, so honor it here (model_init_kwargs wins), else
# a remote-code hybrid with SFTConfig(trust_remote_code=True) resolves as None
# and skips the guard.
top_level_trust_remote_code = getattr(config_arg, "trust_remote_code", None)
if top_level_trust_remote_code is not None:
forward.setdefault("trust_remote_code", top_level_trust_remote_code)
return AutoConfig.from_pretrained(model_name, **forward)
except Exception:
return None
def _chunked_loss_bypasses_forward(config) -> bool:
"""TRL's default ``loss_type="chunked_nll"`` patches the model forward and calls
the backbone directly, so a forward wrapper never runs. Detect it so hybrid
packing stays on the padded path instead of silently skipping the varlen shim."""
try:
import trl.trainer.sft_trainer as _sft_trainer
except Exception:
return False
if not hasattr(_sft_trainer, "_patch_chunked_ce_lm_head"):
return False # TRL has no chunked-CE path -> forward is not bypassed
if getattr(config, "use_liger_kernel", False):
return False # liger forces loss_type="nll" -> normal forward
return getattr(config, "loss_type", None) in (None, "chunked_nll")
# Unsloth gradient accumulation fix:
from transformers import __version__ as transformers_version, ProcessorMixin
@ -632,13 +685,38 @@ def _patch_sft_trainer_auto_packing(trl_module):
is_vlm = False
is_unsupported_model = False
is_hybrid = False
is_encoder_decoder = False
hybrid_varlen_active = False
if model is not None:
model_config = getattr(model, "config", None)
if model_config is None and isinstance(model, str):
# TRL builds a string model inside __init__; resolve its config now.
model_config = _resolve_string_model_config(model, config_arg)
if model_config is not None:
model_types = get_transformers_model_type(model_config)
is_unsupported_model = any(x in PADDING_FREE_BLOCKLIST for x in model_types)
is_vlm = _is_vlm_config(model_config, model_types)
is_hybrid = _is_hybrid_linear_attention_model(model)
is_encoder_decoder = bool(getattr(model_config, "is_encoder_decoder", False))
hybrid_target = (
SimpleNamespace(config = model_config)
if isinstance(model, str) and model_config is not None
else model
)
is_hybrid = _is_hybrid_linear_attention_model(hybrid_target)
# Hybrid models corrupt packed batches unless the gated-delta conv + scan
# reset at sequence boundaries. Enable the experimental varlen shim (flag +
# kernels) so packing stays correct, else keep them blocked. A string model
# (patched only after init) and TRL's chunked-loss forward bypass both leave
# the shim off, so hybrid packing falls back to the padded path.
if (
is_hybrid
and not isinstance(model, str)
and not _chunked_loss_bypasses_forward(config_arg)
):
try:
hybrid_varlen_active = patch_hybrid_linear_attention_varlen(model)
except Exception:
hybrid_varlen_active = False
processing_class = (
args[5] if len(args) >= 6 else kwargs.get("processing_class") or kwargs.get("tokenizer")
@ -664,7 +742,8 @@ def _patch_sft_trainer_auto_packing(trl_module):
or is_auto_processor_vlm
or is_vision_dataset
or is_unsupported_model
or is_hybrid
or is_encoder_decoder
or (is_hybrid and not hybrid_varlen_active)
or (
os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1"
) # Disable padding free on forced logits
@ -684,7 +763,9 @@ def _patch_sft_trainer_auto_packing(trl_module):
reason = "vision-language model with auto processor"
elif is_vision_dataset:
reason = "vision dataset"
elif is_hybrid:
elif is_encoder_decoder:
reason = "encoder-decoder model"
elif is_hybrid and not hybrid_varlen_active:
reason = "hybrid linear-attention model"
elif is_unsupported_model:
reason = f"unsupported model type(s): {', '.join(model_types)}"

View file

@ -17,8 +17,11 @@
from __future__ import annotations
import inspect
import logging
import os
from collections import OrderedDict
from functools import wraps
from typing import Any, Iterable, Optional, Sequence, Tuple
import torch
@ -218,6 +221,305 @@ def enable_padding_free_metadata(model, trainer):
collator._unsloth_padding_free_lengths_wrapped = True
# --- Experimental: correct packing / padding-free for hybrid linear-attention ---
# Qwen3.5 / Qwen3-Next mix a gated-delta recurrence with a causal conv1d. Packing
# flattens the batch, and both ops leak state across sequence boundaries unless we
# pass seq_idx (conv) and cu_seqlens (scan). Only the accelerated kernels accept
# these, so we fail closed on the pure-torch fallbacks. Gated behind an env flag.
#
# Overrides only the per-module prefill kernels (causal_conv1d_fn /
# chunk_gated_delta_rule), leaving decode untouched so generation is unaffected.
# Recompute-safe under gradient checkpointing; never fires for cached forwards.
# Feature-detect (never version-detect), fail closed, idempotent, one deduped
# diagnostic when it declines to activate.
_HYBRID_PACKING_ENV_VAR = "UNSLOTH_EXPERIMENTAL_HYBRID_PACKING"
_HYBRID_LOGGER = logging.getLogger("unsloth.hybrid_packing")
_HYBRID_WARNED: set = set()
def _hybrid_packing_enabled() -> bool:
# Read at call time so setting the flag after `import unsloth` still takes effect.
return os.environ.get(_HYBRID_PACKING_ENV_VAR, "0").strip().lower() in {
"1",
"true",
"yes",
"on",
}
def _hybrid_reject(reason: str) -> bool:
# One deduped diagnostic explaining why hybrid packing stayed on the padded path.
if reason not in _HYBRID_WARNED:
_HYBRID_WARNED.add(reason)
_HYBRID_LOGGER.warning(
"Unsloth: hybrid linear-attention packing disabled (padded path): %s.",
reason,
)
return False
def _iter_gated_delta_modules(model):
modules, seen = [], set()
for module in model.modules():
if id(module) in seen:
continue
seen.add(id(module))
if type(module).__name__.endswith("GatedDeltaNet") and hasattr(module, "conv1d"):
modules.append(module)
return modules
def _hybrid_varlen_kernels_available(gated_delta_modules) -> Optional[str]:
"""None if every module can use the accelerated varlen path, else a short
reason string. All modules are validated before any are mutated; signatures
are read off the captured originals when already wrapped.
Dispatch (the mixer actually calling self.causal_conv1d_fn /
self.chunk_gated_delta_rule) is verified at RUNTIME by the forward-wrapper
handshake, not statically: Unsloth's compile-disable shim hides it from
inspect.getsource, and every supported transformers release dispatches
through the instance attribute."""
if not gated_delta_modules:
return "no gated-delta modules found"
for module in gated_delta_modules:
conv = getattr(module, "_unsloth_varlen_orig_conv", None) or getattr(
module,
"causal_conv1d_fn",
None,
)
scan = getattr(module, "_unsloth_varlen_orig_scan", None) or getattr(
module,
"chunk_gated_delta_rule",
None,
)
if conv is None or scan is None:
return "accelerated kernels missing (install causal_conv1d and fla)"
if getattr(scan, "__name__", "").startswith("torch_") or getattr(
conv,
"__name__",
"",
).startswith("torch_"):
return "pure-torch kernel fallback in use"
try:
if "seq_idx" not in inspect.signature(conv).parameters:
return "conv kernel does not accept seq_idx"
if "cu_seqlens" not in inspect.signature(scan).parameters:
return "scan kernel does not accept cu_seqlens"
except (TypeError, ValueError):
return "kernel signature not introspectable"
return None
def _varlen_from_position_ids(position_ids):
"""(cu_seqlens int32[n+1], seq_idx int32[1,T]) for a flattened padding-free
batch, else None. Padding-free position_ids reset to 0 at each sequence start;
accepts only a validated single-row pack (normal batch or single sequence ->
None). Fallback used only when packed_seq_lengths is absent: it assumes
right-packed reset position_ids and would mis-segment a left-padded row, which
is why packed_seq_lengths is always preferred."""
if position_ids is None:
return None
pos = position_ids
if pos.dim() == 3: # MRoPE [n_planes, 1, T] -> text plane is index 0
pos = pos[0]
if pos.dim() != 2 or pos.shape[0] != 1:
return None
row = pos[0]
total = row.shape[0]
starts = (row == 0).nonzero(as_tuple = False).flatten()
if starts.numel() <= 1 or int(starts[0].item()) != 0:
return None
cu_seqlens = torch.cat(
[
starts.to(torch.int32),
torch.tensor([total], dtype = torch.int32, device = row.device),
]
)
return _seq_idx_from_cu_seqlens(cu_seqlens, total)
def _seq_idx_from_cu_seqlens(cu_seqlens, total):
"""(cu_seqlens int32[n+1], seq_idx int32[1,total]) partitioning [0, total),
else None. Appends a trailing segment for pad_to_multiple_of zero tokens so the
boundaries always cover the full flattened length the kernels see."""
if cu_seqlens is None or cu_seqlens.numel() < 2 or int(cu_seqlens[0].item()) != 0:
return None
boundaries = cu_seqlens.to(torch.int32)
last = int(boundaries[-1].item())
if last > total:
return None
if last < total: # trailing pad tokens -> one final segment
boundaries = torch.cat(
[
boundaries,
torch.tensor([total], dtype = torch.int32, device = boundaries.device),
]
)
lengths = boundaries[1:] - boundaries[:-1]
if not bool((lengths > 0).all()):
return None
seq_idx = torch.repeat_interleave(
torch.arange(lengths.numel(), dtype = torch.int32, device = boundaries.device),
lengths.to(torch.int64),
).unsqueeze(0)
return boundaries, seq_idx
def _hybrid_varlen_metadata(kwargs):
"""Boundary metadata (cu_seqlens, seq_idx) for one flattened packed forward,
else None. Prefers the authoritative packed_seq_lengths, falls back to
reset-style position_ids. Returns None for cached forwards and non-packed
batches so decode / eval / normal batches are a strict no-op."""
if kwargs.get("use_cache"):
return None
if kwargs.get("past_key_values") is not None or kwargs.get("cache_params") is not None:
return None
total, device = None, None
for key in ("input_ids", "inputs_embeds", "position_ids"):
tensor = kwargs.get(key)
if tensor is not None and hasattr(tensor, "shape"):
total = tensor.shape[1] if key == "inputs_embeds" else tensor.shape[-1]
device = tensor.device
break
if total is None:
return None
psl = kwargs.get("packed_seq_lengths")
if psl is not None and getattr(psl, "numel", lambda: 1)() > 0: # skip empty (no max())
info = get_packed_info_from_kwargs(kwargs, device)
if info is not None:
_, cu_seqlens, _ = info
built = _seq_idx_from_cu_seqlens(cu_seqlens, total)
if built is not None:
return built
return _varlen_from_position_ids(kwargs.get("position_ids"))
def patch_hybrid_linear_attention_varlen(model) -> bool:
"""Feed seq_idx / cu_seqlens to the gated-delta conv + scan so packing and
padding-free reset state at sequence boundaries. Gated by
UNSLOTH_EXPERIMENTAL_HYBRID_PACKING and fail-closed. Returns True when the
varlen path is active, so the caller may allow packing for the model.
Idempotent: repeat calls on an already-patched model return True."""
if not _hybrid_packing_enabled():
return False
gated_delta_modules = _iter_gated_delta_modules(model)
# Idempotency: an already fully-patched model stays active without re-validation.
if (
getattr(model, "_unsloth_varlen_forward_wrapped", False)
and gated_delta_modules
and all(getattr(m, "_unsloth_varlen_wrapped", False) for m in gated_delta_modules)
):
return True
reason = _hybrid_varlen_kernels_available(gated_delta_modules)
if reason is not None:
return _hybrid_reject(reason)
# Transactional: every module validated above, now wrap each and stash originals.
for module in gated_delta_modules:
if getattr(module, "_unsloth_varlen_wrapped", False):
continue
conv_orig, scan_orig = module.causal_conv1d_fn, module.chunk_gated_delta_rule
module._unsloth_varlen_orig_conv = conv_orig
module._unsloth_varlen_orig_scan = scan_orig
@wraps(conv_orig)
def conv_fn(
*args,
_orig = conv_orig,
_module = module,
**kwargs,
):
varlen = getattr(_module, "_unsloth_varlen", None)
if varlen is not None:
_module._unsloth_varlen_conv_hit = True # runtime dispatch handshake
if kwargs.get("seq_idx") is None:
kwargs["seq_idx"] = varlen[1]
return _orig(*args, **kwargs)
@wraps(scan_orig)
def scan_fn(
*args,
_orig = scan_orig,
_module = module,
**kwargs,
):
varlen = getattr(_module, "_unsloth_varlen", None)
if varlen is not None:
_module._unsloth_varlen_scan_hit = True
if kwargs.get("cu_seqlens") is None:
kwargs["cu_seqlens"] = varlen[0]
return _orig(*args, **kwargs)
module.causal_conv1d_fn = conv_fn
module.chunk_gated_delta_rule = scan_fn
module._unsloth_varlen = None
module._unsloth_varlen_wrapped = True
# Refresh the boundary stash on the outermost forward (once per step, outside
# gradient-checkpoint recompute, so it stays valid for recomputed inner
# forwards). Read from both positional and keyword args via the bound signature.
if not getattr(model, "_unsloth_varlen_forward_wrapped", False):
forward_orig = model.forward
try:
forward_sig = inspect.signature(forward_orig)
except (TypeError, ValueError):
forward_sig = None
@wraps(forward_orig)
def forward_with_varlen(*args, **kwargs):
try:
bound = dict(kwargs)
if forward_sig is not None and args:
bound.update(forward_sig.bind_partial(*args).arguments)
varlen = _hybrid_varlen_metadata(bound)
except Exception:
varlen = None
first_pack = varlen is not None and not getattr(
model,
"_unsloth_varlen_handshake_done",
False,
)
for module in gated_delta_modules:
module._unsloth_varlen = varlen
if first_pack:
module._unsloth_varlen_conv_hit = False
module._unsloth_varlen_scan_hit = False
out = forward_orig(*args, **kwargs)
# Runtime dispatch handshake: on the first packed forward, confirm BOTH
# boundary kernels ran for EVERY module. seq_idx (conv) and cu_seqlens
# (scan) are both load-bearing, so a partial/absent dispatch (a future
# version no longer routing through self.<kernel>) leaves cross-sequence
# contamination. The batch is already flattened with no padded recovery,
# so abort before loss/backward rather than train on corrupted data.
if first_pack:
model._unsloth_varlen_handshake_done = True
missing = [
type(m).__name__
for m in gated_delta_modules
if not (
getattr(m, "_unsloth_varlen_conv_hit", False)
and getattr(m, "_unsloth_varlen_scan_hit", False)
)
]
if missing:
for m in gated_delta_modules:
m._unsloth_varlen = None
_hybrid_reject("varlen conv/scan not both dispatched (dispatch changed?)")
raise RuntimeError(
"Unsloth: experimental hybrid packing cannot continue because the "
"varlen conv/scan wrappers were not both invoked for "
f"{sorted(set(missing))}. Unset UNSLOTH_EXPERIMENTAL_HYBRID_PACKING "
"to train these models on the padded path."
)
return out
model.forward = forward_with_varlen
model._unsloth_varlen_forward_wrapped = True
return True
def get_packed_info_from_kwargs(
kwargs: dict, device: torch.device
) -> Optional[Tuple[torch.Tensor, torch.Tensor, int]]: