diff --git a/tests/utils/test_packing.py b/tests/utils/test_packing.py index 98c29d9f0f..1b8bb65058 100644 --- a/tests/utils/test_packing.py +++ b/tests/utils/test_packing.py @@ -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. + 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. 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 diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index b0709f7376..4ef3af6add 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -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) diff --git a/unsloth/trainer.py b/unsloth/trainer.py index 61d41aad21..1c30192301 100644 --- a/unsloth/trainer.py +++ b/unsloth/trainer.py @@ -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)}" diff --git a/unsloth/utils/packing.py b/unsloth/utils/packing.py index dd0a1bfb62..f8d539fb93 100644 --- a/unsloth/utils/packing.py +++ b/unsloth/utils/packing.py @@ -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.) 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]]: