Update ASFT streaming config to use mode-based API

Replaces `enabled`/`ref_strategy` with unified `mode` parameter in ASFTStreamingConfig. Adds "auto", "seq", "batch", "hybrid", and "off" modes with automatic fallback logic. Implements seq_kv_cache streaming with KV cache reuse and batch microbatching support. Updates notebook defaults to use `mode="auto"` for optimal VRAM reduction. Adds comprehensive tests for mode routing, fallback behavior, and backward compatibility.

Llama3_1_(8B)_Alpaca-ASFT.ipynb

@@ -541,7 +541,7 @@
     "#### Recommended ASFT defaults (optimized)\n",
     "- `asft_mode=\"asft\"` + `kl_weight=0.05`: a solid starting point (stable, not overly restrictive).\n",
     "- `reference_policy=\"disable_adapter\"`: avoids keeping a separate frozen reference copy in most PEFT setups.\n",
-    "- `ASFTStreamingConfig(enabled=True, ref_strategy=\"batch_micro\")`: micro-batches the reference forward pass to reduce peak VRAM.\n",
+    "- `ASFTStreamingConfig(mode=\"auto\")`: tries seq_kv_cache first, falls back to batch micro to reduce peak VRAM.\n",
     "\n",
     "For quick comparisons, try: `asft_mode=\"sft\"` or `asft_mode=\"dft\"` (KL off)."
    ]
@@ -611,9 +611,9 @@
     "\n",
     "# Reference forward streaming to reduce VRAM peak\n",
     "asft_streaming = ASFTStreamingConfig(\n",
-    "    enabled = True,\n",
-    "    ref_strategy = \"batch_micro\",\n",
+    "    mode = \"auto\",\n",
     "    # ref_microbatch_size = 1,  # Set manually if desired; None picks automatically\n",
+    "    # seq_chunk_size = 256,    # Adjust for long sequences if needed\n",
     "    force_fp32_kl = True,\n",
     ")\n",
     "\n",
@@ -660,7 +660,7 @@
     "print(\"Reference policy:\", getattr(trainer, \"reference_policy\", None))\n",
     "streaming = getattr(trainer, \"asft_streaming\", None)\n",
     "if streaming is not None:\n",
-    "    print(\"Streaming enabled:\", getattr(streaming, \"enabled\", None))\n",
+    "    print(\"Streaming mode:\", getattr(streaming, \"mode\", None))\n",
     "    print(\"Streaming strategy:\", getattr(streaming, \"ref_strategy\", None))"
    ]
   },

tests/test_asft.py

@@ -35,6 +35,7 @@ from unsloth.losses.asft import (
     compute_asft_loss,
     _compute_kl_divergence,
     _compute_dft_weights,
+    _compute_kl_seq_kv_cache,
 )
 
 
@@ -194,6 +195,20 @@ class TestFastCrossEntropyLossPerToken:
         # Should be close
         assert torch.allclose(losses, pytorch_losses, atol = 1e-4)
 
+    def test_respects_custom_ignore_index(self):
+        """Test that custom ignore_index is honored by the kernel wrapper."""
+        torch.manual_seed(0)
+        logits = torch.randn(1, 4, 8)
+        labels = torch.tensor([[1, 2, 1, 3]], dtype = torch.long)
+
+        losses, valid_mask = fast_cross_entropy_loss_per_token(
+            logits, labels, ignore_index = 1
+        )
+
+        assert losses.shape == (4,)
+        assert torch.equal(valid_mask, torch.tensor([False, True, False, True]))
+        assert torch.all(losses[~valid_mask] == 0)
+
 
 # -----------------------------------------------------------------------------
 # A3) Test build_shift_labels
@@ -293,6 +308,17 @@ class TestGetReferenceForwardCallable:
         # Should still work (uses frozen copy fallback)
         assert result is not None
 
+    def test_return_outputs_true(self, simple_model):
+        """Test returning full outputs when requested."""
+        ref_forward = get_reference_forward_callable(
+            simple_model, reference_policy = "frozen_copy", return_outputs = True
+        )
+
+        input_ids = torch.tensor([[1, 2, 3, 4]])
+        result = ref_forward(input_ids = input_ids)
+
+        assert hasattr(result, "logits")
+
 
 # -----------------------------------------------------------------------------
 # Test KL divergence computation
@@ -361,6 +387,24 @@ class TestDFTWeights:
 
         assert not weights.requires_grad
 
+    def test_dft_weights_match_exp_neg_ce(self):
+        """Test exp(-CE) matches softmax-gather for DFT weights."""
+        torch.manual_seed(123)
+        logits = torch.randn(2, 3, 7)
+        labels = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype = torch.long)
+
+        ce_losses, valid_mask = fast_cross_entropy_loss_per_token(logits, labels)
+
+        weights_from_ce = _compute_dft_weights(
+            logits,
+            labels,
+            ce_losses = ce_losses,
+            valid_mask = valid_mask,
+        )
+        weights_from_softmax = _compute_dft_weights(logits, labels)
+
+        assert torch.allclose(weights_from_ce, weights_from_softmax, atol = 1e-4)
+
 
 # -----------------------------------------------------------------------------
 # A5) Test compute_asft_loss
@@ -496,6 +540,7 @@ class TestASFTStreamingConfig:
         """Test default configuration values."""
         config = ASFTStreamingConfig()
 
+        assert config.mode is None
         assert config.enabled is False
         assert config.ref_strategy == "none"
         assert config.ref_microbatch_size is None
@@ -506,17 +551,371 @@ class TestASFTStreamingConfig:
     def test_custom_values(self):
         """Test custom configuration values."""
         config = ASFTStreamingConfig(
+            mode = "batch",
             enabled = True,
             ref_strategy = "batch_micro",
             ref_microbatch_size = 4,
             seq_chunk_size = 256,
         )
 
+        assert config.mode == "batch"
         assert config.enabled is True
         assert config.ref_strategy == "batch_micro"
         assert config.ref_microbatch_size == 4
         assert config.seq_chunk_size == 256
 
+
+class TestStreamingModeMapping:
+    """Tests for streaming mode routing in compute_asft_loss."""
+
+    def test_mode_batch_uses_batch_micro(self, simple_model):
+        """Test that mode=batch routes to batch micro."""
+        inputs = {
+            "input_ids": torch.tensor([[1, 2, 3, 4], [2, 3, 4, 5]]),
+            "labels": torch.tensor([[1, 2, 3, 4], [2, 3, 4, 5]]),
+        }
+        config = ASFTStreamingConfig(
+            mode = "batch",
+            ref_microbatch_size = 1,
+            enabled = False,
+            ref_strategy = "seq_kv_cache",
+        )
+
+        def batch_side_effect(
+            model,
+            cur_logits,
+            shift_labels,
+            valid_mask,
+            ref_forward,
+            forward_inputs,
+            microbatch_size,
+            logit_softcapping = 0,
+            logit_scaling = 0,
+            force_fp32 = True,
+        ):
+            batch, seq_len = shift_labels.shape
+            return torch.zeros(batch, seq_len, device = shift_labels.device)
+
+        with patch(
+            "unsloth.losses.asft._compute_kl_batch_micro",
+            side_effect = batch_side_effect,
+        ) as batch_mock, patch(
+            "unsloth.losses.asft._compute_kl_seq_kv_cache",
+            side_effect = AssertionError("seq_kv_cache should not be used"),
+        ):
+            loss = compute_asft_loss(
+                simple_model,
+                inputs,
+                asft_mode = "sft+kl",
+                kl_weight = 0.1,
+                reference_policy = "frozen_copy",
+                streaming_config = config,
+            )
+
+        assert batch_mock.called
+        assert batch_mock.call_args[0][6] == 1
+        assert loss.dim() == 0
+
+    @pytest.mark.parametrize("mode", ["seq", "auto"])
+    def test_mode_seq_and_auto_use_seq_kv_cache(self, mode, simple_model):
+        """Test that mode=seq/auto routes to seq_kv_cache."""
+        inputs = {
+            "input_ids": torch.tensor([[1, 2, 3, 4]]),
+            "labels": torch.tensor([[1, 2, 3, 4]]),
+        }
+        config = ASFTStreamingConfig(
+            mode = mode,
+            seq_chunk_size = 2,
+            enabled = False,
+            ref_strategy = "batch_micro",
+        )
+
+        def seq_side_effect(
+            model,
+            cur_logits,
+            shift_labels,
+            valid_mask,
+            ref_forward,
+            forward_inputs,
+            seq_chunk_size,
+            **kwargs,
+        ):
+            batch, seq_len = shift_labels.shape
+            return torch.zeros(batch, seq_len, device = shift_labels.device)
+
+        with patch(
+            "unsloth.losses.asft._compute_kl_seq_kv_cache",
+            side_effect = seq_side_effect,
+        ) as seq_mock, patch(
+            "unsloth.losses.asft._compute_kl_batch_micro",
+            side_effect = AssertionError("batch_micro should not be used"),
+        ):
+            loss = compute_asft_loss(
+                simple_model,
+                inputs,
+                asft_mode = "sft+kl",
+                kl_weight = 0.1,
+                reference_policy = "frozen_copy",
+                streaming_config = config,
+            )
+
+        assert seq_mock.called
+        assert seq_mock.call_args[0][6] == 2
+        assert seq_mock.call_args.kwargs["microbatch_size"] is None
+        assert loss.dim() == 0
+
+    def test_mode_hybrid_defaults_microbatch(self, simple_model):
+        """Test that hybrid mode sets a default microbatch size."""
+        inputs = {
+            "input_ids": torch.tensor([[1, 2, 3, 4], [2, 3, 4, 5]]),
+            "labels": torch.tensor([[1, 2, 3, 4], [2, 3, 4, 5]]),
+        }
+        config = ASFTStreamingConfig(
+            mode = "hybrid",
+            seq_chunk_size = 2,
+            ref_microbatch_size = None,
+        )
+
+        def seq_side_effect(
+            model,
+            cur_logits,
+            shift_labels,
+            valid_mask,
+            ref_forward,
+            forward_inputs,
+            seq_chunk_size,
+            **kwargs,
+        ):
+            batch, seq_len = shift_labels.shape
+            return torch.zeros(batch, seq_len, device = shift_labels.device)
+
+        with patch(
+            "unsloth.losses.asft._compute_kl_seq_kv_cache",
+            side_effect = seq_side_effect,
+        ) as seq_mock:
+            loss = compute_asft_loss(
+                simple_model,
+                inputs,
+                asft_mode = "sft+kl",
+                kl_weight = 0.1,
+                reference_policy = "frozen_copy",
+                streaming_config = config,
+            )
+
+        assert seq_mock.called
+        assert seq_mock.call_args.kwargs["microbatch_size"] == 1
+        assert config.ref_microbatch_size is None
+        assert loss.dim() == 0
+
+    def test_mode_off_uses_full_forward(self, simple_model):
+        """Test that mode=off bypasses streaming helpers."""
+        inputs = {
+            "input_ids": torch.tensor([[1, 2, 3, 4]]),
+            "labels": torch.tensor([[1, 2, 3, 4]]),
+        }
+        config = ASFTStreamingConfig(
+            mode = "off",
+            enabled = True,
+            ref_strategy = "seq_kv_cache",
+        )
+
+        def kl_side_effect(
+            cur_logits,
+            ref_logits,
+            model = None,
+            logit_softcapping = 0,
+            logit_scaling = 0,
+            force_fp32 = True,
+        ):
+            batch, seq_len = ref_logits.shape[:2]
+            return torch.zeros(batch * seq_len, device = ref_logits.device)
+
+        with patch(
+            "unsloth.losses.asft._compute_kl_divergence",
+            side_effect = kl_side_effect,
+        ) as kl_mock, patch(
+            "unsloth.losses.asft._compute_kl_seq_kv_cache",
+            side_effect = AssertionError("seq_kv_cache should not be used"),
+        ), patch(
+            "unsloth.losses.asft._compute_kl_batch_micro",
+            side_effect = AssertionError("batch_micro should not be used"),
+        ):
+            loss = compute_asft_loss(
+                simple_model,
+                inputs,
+                asft_mode = "sft+kl",
+                kl_weight = 0.1,
+                reference_policy = "frozen_copy",
+                streaming_config = config,
+            )
+
+        assert kl_mock.called
+        assert loss.dim() == 0
+
+    def test_invalid_mode_raises(self, simple_model):
+        """Test that invalid streaming mode raises a ValueError."""
+        inputs = {
+            "input_ids": torch.tensor([[1, 2, 3, 4]]),
+            "labels": torch.tensor([[1, 2, 3, 4]]),
+        }
+        config = ASFTStreamingConfig(mode = "invalid")
+
+        with pytest.raises(ValueError):
+            compute_asft_loss(
+                simple_model,
+                inputs,
+                asft_mode = "sft+kl",
+                kl_weight = 0.1,
+                reference_policy = "frozen_copy",
+                streaming_config = config,
+            )
+
+
+class TestSeqKVCacheStreaming:
+    """Tests for seq_kv_cache streaming behavior."""
+
+    def test_seq_kv_cache_runs_when_use_cache_false(self):
+        """Test that seq_kv_cache attempts chunking even if config.use_cache=False."""
+        batch_size, seq_len, vocab_size = 1, 6, 5
+        cur_logits = torch.randn(batch_size, seq_len, vocab_size)
+        shift_labels = torch.zeros(batch_size, seq_len, dtype = torch.long)
+        valid_mask = shift_labels != -100
+        input_ids = torch.arange(seq_len).view(1, -1)
+
+        class DummyModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.config = SimpleNamespace(
+                    use_cache = False,
+                    final_logit_softcapping = 0,
+                    logit_scale = 0,
+                )
+
+        model = DummyModel()
+        call_state = {"saw_past": False}
+
+        def ref_forward(**kwargs):
+            input_ids_local = kwargs["input_ids"]
+            if input_ids_local.shape[1] == seq_len:
+                raise AssertionError("full forward not expected")
+            if "past_key_values" in kwargs:
+                call_state["saw_past"] = True
+            batch, chunk_len = input_ids_local.shape
+            logits = torch.zeros(
+                batch, chunk_len, vocab_size, device = input_ids_local.device
+            )
+            return (logits, ("cache",))
+
+        forward_inputs = {"input_ids": input_ids}
+
+        kl = _compute_kl_seq_kv_cache(
+            model,
+            cur_logits,
+            shift_labels,
+            valid_mask,
+            ref_forward,
+            forward_inputs,
+            seq_chunk_size = 4,
+        )
+
+        assert kl.shape == (batch_size, seq_len)
+        assert call_state["saw_past"] is True
+
+    def test_seq_kv_cache_supports_microbatching(self):
+        """Test that seq_kv_cache can be microbatched by batch dimension."""
+        batch_size, seq_len, vocab_size = 2, 4, 3
+        cur_logits = torch.randn(batch_size, seq_len, vocab_size)
+        shift_labels = torch.zeros(batch_size, seq_len, dtype = torch.long)
+        valid_mask = shift_labels != -100
+        input_ids = torch.arange(seq_len).repeat(batch_size, 1)
+
+        class DummyModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.config = SimpleNamespace(
+                    use_cache = True,
+                    final_logit_softcapping = 0,
+                    logit_scale = 0,
+                )
+
+        model = DummyModel()
+        call_state = {"max_batch": 0}
+
+        def ref_forward(**kwargs):
+            input_ids_local = kwargs["input_ids"]
+            call_state["max_batch"] = max(
+                call_state["max_batch"], input_ids_local.shape[0]
+            )
+            if input_ids_local.shape[0] > 1:
+                raise AssertionError("expected microbatching")
+            batch, chunk_len = input_ids_local.shape
+            logits = torch.zeros(
+                batch, chunk_len, vocab_size, device = input_ids_local.device
+            )
+            return (logits, ("cache",))
+
+        forward_inputs = {"input_ids": input_ids}
+
+        kl = _compute_kl_seq_kv_cache(
+            model,
+            cur_logits,
+            shift_labels,
+            valid_mask,
+            ref_forward,
+            forward_inputs,
+            seq_chunk_size = 2,
+            microbatch_size = 1,
+        )
+
+        assert kl.shape == (batch_size, seq_len)
+        assert call_state["max_batch"] == 1
+
+    def test_seq_kv_cache_falls_back_to_batch_micro(self):
+        """Test that seq_kv_cache falls back to batch micro on cache failure."""
+        batch_size, seq_len, vocab_size = 2, 6, 5
+        cur_logits = torch.randn(batch_size, seq_len, vocab_size)
+        shift_labels = torch.zeros(batch_size, seq_len, dtype = torch.long)
+        valid_mask = shift_labels != -100
+        input_ids = torch.arange(seq_len).repeat(batch_size, 1)
+
+        class DummyModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.config = SimpleNamespace(
+                    use_cache = True,
+                    final_logit_softcapping = 0,
+                    logit_scale = 0,
+                )
+
+        model = DummyModel()
+
+        def ref_forward(**kwargs):
+            input_ids_local = kwargs["input_ids"]
+            if (
+                input_ids_local.shape[0] == batch_size
+                and input_ids_local.shape[1] == seq_len
+            ):
+                raise AssertionError("full forward not expected on fallback")
+            batch, chunk_len = input_ids_local.shape
+            logits = torch.zeros(
+                batch, chunk_len, vocab_size, device = input_ids_local.device
+            )
+            return (logits, None)
+
+        forward_inputs = {"input_ids": input_ids}
+
+        kl = _compute_kl_seq_kv_cache(
+            model,
+            cur_logits,
+            shift_labels,
+            valid_mask,
+            ref_forward,
+            forward_inputs,
+            seq_chunk_size = 2,
+        )
+
+        assert kl.shape == (batch_size, seq_len)
+
     def test_config_immutability_when_none_values(self, simple_model):
         """Test that streaming_config is not mutated when values are None."""
         config = ASFTStreamingConfig(
@@ -604,6 +1003,113 @@ class TestBackwardCompatibility:
         # Should be very close
         assert torch.allclose(full_loss, streaming_loss, atol = 1e-4)
 
+    def test_seq_kv_cache_equivalence(self):
+        """Test that seq_kv_cache matches full forward for KL loss."""
+        torch.manual_seed(123)
+
+        class CacheModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.config = SimpleNamespace(
+                    use_cache = True,
+                    final_logit_softcapping = 0,
+                    logit_scale = 0,
+                )
+                self.embedding = nn.Embedding(32, 8)
+                self.linear = nn.Linear(8, 32)
+
+            def forward(
+                self, input_ids = None, past_key_values = None, use_cache = None, **kwargs
+            ):
+                embeddings = self.embedding(input_ids)
+                logits = self.linear(embeddings)
+                past = ("cache",) if (use_cache or past_key_values is not None) else None
+                return SimpleNamespace(logits = logits, past_key_values = past)
+
+        model = CacheModel()
+        inputs = {
+            "input_ids": torch.tensor([[1, 2, 3, 4, 5, 6], [6, 5, 4, 3, 2, 1]]),
+            "labels": torch.tensor([[1, 2, 3, 4, 5, 6], [6, 5, 4, 3, 2, 1]]),
+        }
+
+        full_loss = compute_asft_loss(
+            model,
+            inputs,
+            asft_mode = "sft+kl",
+            kl_weight = 0.1,
+            reference_policy = "frozen_copy",
+            streaming_config = ASFTStreamingConfig(enabled = False),
+        )
+
+        seq_loss = compute_asft_loss(
+            model,
+            inputs,
+            asft_mode = "sft+kl",
+            kl_weight = 0.1,
+            reference_policy = "frozen_copy",
+            streaming_config = ASFTStreamingConfig(
+                enabled = True,
+                ref_strategy = "seq_kv_cache",
+                seq_chunk_size = 2,
+            ),
+        )
+
+        assert torch.allclose(full_loss, seq_loss, atol = 1e-4)
+
+    def test_seq_kv_cache_microbatch_equivalence(self):
+        """Test that seq_kv_cache + microbatching matches full forward."""
+        torch.manual_seed(456)
+
+        class CacheModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.config = SimpleNamespace(
+                    use_cache = True,
+                    final_logit_softcapping = 0,
+                    logit_scale = 0,
+                )
+                self.embedding = nn.Embedding(32, 8)
+                self.linear = nn.Linear(8, 32)
+
+            def forward(
+                self, input_ids = None, past_key_values = None, use_cache = None, **kwargs
+            ):
+                embeddings = self.embedding(input_ids)
+                logits = self.linear(embeddings)
+                past = ("cache",) if (use_cache or past_key_values is not None) else None
+                return SimpleNamespace(logits = logits, past_key_values = past)
+
+        model = CacheModel()
+        inputs = {
+            "input_ids": torch.tensor([[1, 2, 3, 4, 5, 6], [6, 5, 4, 3, 2, 1]]),
+            "labels": torch.tensor([[1, 2, 3, 4, 5, 6], [6, 5, 4, 3, 2, 1]]),
+        }
+
+        full_loss = compute_asft_loss(
+            model,
+            inputs,
+            asft_mode = "sft+kl",
+            kl_weight = 0.1,
+            reference_policy = "frozen_copy",
+            streaming_config = ASFTStreamingConfig(enabled = False),
+        )
+
+        combined_loss = compute_asft_loss(
+            model,
+            inputs,
+            asft_mode = "sft+kl",
+            kl_weight = 0.1,
+            reference_policy = "frozen_copy",
+            streaming_config = ASFTStreamingConfig(
+                enabled = True,
+                ref_strategy = "seq_kv_cache",
+                seq_chunk_size = 2,
+                ref_microbatch_size = 1,
+            ),
+        )
+
+        assert torch.allclose(full_loss, combined_loss, atol = 1e-4)
+
 
 # -----------------------------------------------------------------------------
 # Integration Tests
@@ -627,5 +1133,123 @@ class TestASFTTrainerIntegration:
         assert issubclass(ASFTTrainer, UnslothTrainer)
 
 
+class TestASFTTrainerComputeLoss:
+    """Tests for ASFTTrainer.compute_loss behavior."""
+
+    def test_compute_loss_calls_asft_loss(self):
+        """Test ASFTTrainer compute_loss calls compute_asft_loss."""
+        from unsloth.trainer import ASFTTrainer, ASFTStreamingConfig
+
+        trainer = ASFTTrainer.__new__(ASFTTrainer)
+        trainer.asft_enabled = True
+        trainer.asft_mode = "sft"
+        trainer.kl_weight = 0.0
+        trainer.reference_policy = "disable_adapter"
+        trainer.asft_streaming = ASFTStreamingConfig()
+        trainer._asft_original_model = None
+
+        model = nn.Module()
+        inputs = {
+            "input_ids": torch.tensor([[1, 2, 3, 4]]),
+            "labels": torch.tensor([[1, 2, 3, 4]]),
+        }
+        expected = torch.tensor(1.0, device = inputs["input_ids"].device)
+
+        with patch(
+            "unsloth.trainer.compute_asft_loss", return_value = expected
+        ) as loss_mock:
+            result = ASFTTrainer.compute_loss(
+                trainer, model, inputs, return_outputs = False, num_items_in_batch = 7
+            )
+
+        assert result is expected
+        assert inputs["num_items_in_batch"] == 7
+        assert loss_mock.called
+        assert loss_mock.call_args.kwargs["model"] is model
+        assert loss_mock.call_args.kwargs["asft_mode"] == "sft"
+        assert loss_mock.call_args.kwargs["kl_weight"] == 0.0
+        assert loss_mock.call_args.kwargs["reference_policy"] == "disable_adapter"
+        assert loss_mock.call_args.kwargs["streaming_config"] is trainer.asft_streaming
+
+    def test_compute_loss_creates_frozen_copy_once(self):
+        """Test frozen copy is created once when needed."""
+        from unsloth.trainer import ASFTTrainer, ASFTStreamingConfig
+
+        trainer = ASFTTrainer.__new__(ASFTTrainer)
+        trainer.asft_enabled = True
+        trainer.asft_mode = "asft"
+        trainer.kl_weight = 0.1
+        trainer.reference_policy = "frozen_copy"
+        trainer.asft_streaming = ASFTStreamingConfig()
+        trainer._asft_original_model = None
+
+        model = nn.Module()
+        model_copy = MagicMock()
+        inputs = {
+            "input_ids": torch.tensor([[1, 2, 3, 4]]),
+            "labels": torch.tensor([[1, 2, 3, 4]]),
+        }
+
+        with patch(
+            "unsloth.trainer.deepcopy", return_value = model_copy
+        ) as deepcopy_mock, patch(
+            "unsloth.trainer.compute_asft_loss",
+            return_value = torch.tensor(0.5, device = inputs["input_ids"].device),
+        ):
+            ASFTTrainer.compute_loss(trainer, model, inputs)
+            ASFTTrainer.compute_loss(trainer, model, inputs)
+
+        assert deepcopy_mock.call_count == 1
+        assert trainer._asft_original_model is model_copy
+        assert model_copy.eval.called
+        assert model_copy.requires_grad_.called
+
+    def test_compute_loss_skips_copy_with_disable_adapter(self):
+        """Test disable_adapter policy skips frozen copy."""
+        from unsloth.trainer import ASFTTrainer, ASFTStreamingConfig
+
+        trainer = ASFTTrainer.__new__(ASFTTrainer)
+        trainer.asft_enabled = True
+        trainer.asft_mode = "asft"
+        trainer.kl_weight = 0.1
+        trainer.reference_policy = "disable_adapter"
+        trainer.asft_streaming = ASFTStreamingConfig()
+        trainer._asft_original_model = None
+
+        model = MagicMock()
+        model.disable_adapter = MagicMock()
+        inputs = {
+            "input_ids": torch.tensor([[1, 2, 3, 4]]),
+            "labels": torch.tensor([[1, 2, 3, 4]]),
+        }
+
+        with patch(
+            "unsloth.trainer.deepcopy"
+        ) as deepcopy_mock, patch(
+            "unsloth.trainer.compute_asft_loss",
+            return_value = torch.tensor(0.5, device = inputs["input_ids"].device),
+        ) as loss_mock:
+            ASFTTrainer.compute_loss(trainer, model, inputs)
+
+        assert not deepcopy_mock.called
+        assert loss_mock.call_args.kwargs["original_model"] is None
+
+
+class TestUnslothTrainingArguments:
+    """Tests for UnslothTrainingArguments."""
+
+    def test_embedding_learning_rate_is_set(self):
+        """Test embedding_learning_rate is stored on the args object."""
+        from unsloth import trainer as trainer_module
+
+        with patch.object(
+            trainer_module.TrainingArguments, "__init__", return_value = None
+        ) as base_init:
+            args = trainer_module.UnslothTrainingArguments(embedding_learning_rate = 0.01)
+
+        assert args.embedding_learning_rate == 0.01
+        assert base_init.called
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

tests/test_unsloth_cli.py

@@ -0,0 +1,80 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CLI argument parsing tests for unsloth-cli.py."""
+
+from pathlib import Path
+import importlib.util
+
+import pytest
+
+
+def _load_cli_module():
+    root = Path(__file__).resolve().parents[1]
+    cli_path = root / "unsloth-cli.py"
+    spec = importlib.util.spec_from_file_location("unsloth_cli", cli_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_cli_defaults_asft():
+    cli = _load_cli_module()
+    parser = cli.build_parser()
+    args = parser.parse_args([])
+
+    assert args.asft is False
+    assert args.asft_mode == "asft"
+    assert args.kl_weight == 0.0
+    assert args.reference_policy == "disable_adapter"
+    assert args.asft_streaming == "off"
+    assert args.ref_microbatch_size is None
+    assert args.seq_chunk_size is None
+
+
+def test_cli_asft_streaming_flag_defaults_auto():
+    cli = _load_cli_module()
+    parser = cli.build_parser()
+    args = parser.parse_args(["--asft_streaming"])
+
+    assert args.asft_streaming == "auto"
+
+
+def test_cli_asft_streaming_value():
+    cli = _load_cli_module()
+    parser = cli.build_parser()
+    args = parser.parse_args(["--asft_streaming", "batch"])
+
+    assert args.asft_streaming == "batch"
+
+
+def test_cli_asft_options_parsed():
+    cli = _load_cli_module()
+    parser = cli.build_parser()
+    args = parser.parse_args(
+        [
+            "--asft",
+            "--asft_mode",
+            "sft+kl",
+            "--kl_weight",
+            "0.2",
+            "--reference_policy",
+            "frozen_copy",
+        ]
+    )
+
+    assert args.asft is True
+    assert args.asft_mode == "sft+kl"
+    assert args.kl_weight == pytest.approx(0.2)
+    assert args.reference_policy == "frozen_copy"

unsloth-cli.py

@@ -161,8 +161,7 @@ def run(args):
     if asft_enabled:
         # Build ASFT streaming config
         asft_streaming = ASFTStreamingConfig(
-            enabled = getattr(args, "asft_streaming", False),
-            ref_strategy = getattr(args, "ref_strategy", "none"),
+            mode = getattr(args, "asft_streaming", None),
             ref_microbatch_size = getattr(args, "ref_microbatch_size", None),
             seq_chunk_size = getattr(args, "seq_chunk_size", None),
         )
@@ -234,7 +233,7 @@ def run(args):
         print("Warning: The model is not saved!")
 
 
-if __name__ == "__main__":
+def build_parser():
     parser = argparse.ArgumentParser(
         description = "🦥 Fine-tune your llm faster using unsloth!"
     )
@@ -535,25 +534,23 @@ if __name__ == "__main__":
     )
     asft_group.add_argument(
         "--asft_streaming",
-        action = "store_true",
-        help = "Enable streaming for reference model to reduce VRAM usage",
-    )
-    asft_group.add_argument(
-        "--ref_strategy",
-        type = str,
-        default = "none",
-        choices = ["none", "batch_micro", "seq_kv_cache"],
+        nargs = "?",
+        const = "auto",
+        default = "off",
+        choices = ["off", "auto", "batch", "seq", "hybrid"],
         help = (
-            "Streaming strategy for reference forward: 'none' (full forward), "
-            "'batch_micro' (microbatch by batch), 'seq_kv_cache' (sequence chunking with KV cache). "
-            "Default: 'none'"
+            "Streaming mode for reference forward: 'off' (full forward), "
+            "'auto' (seq_kv_cache with batch-micro fallback), "
+            "'batch' (microbatch by batch), 'seq' (sequence chunking with KV cache), "
+            "'hybrid' (batch micro + seq_kv_cache). "
+            "Use flag without value for 'auto'."
         ),
     )
     asft_group.add_argument(
         "--ref_microbatch_size",
         type = int,
         default = None,
-        help = "Microbatch size for batch_micro strategy",
+        help = "Microbatch size for batch_micro or seq_kv_cache strategy",
     )
     asft_group.add_argument(
         "--seq_chunk_size",
@@ -562,5 +559,10 @@ if __name__ == "__main__":
         help = "Sequence chunk size for seq_kv_cache strategy",
     )
 
+    return parser
+
+
+if __name__ == "__main__":
+    parser = build_parser()
     args = parser.parse_args()
     run(args)

unsloth/losses/asft.py

@@ -68,17 +68,24 @@ class ASFTStreamingConfig:
     """Configuration for ASFT streaming strategies to reduce VRAM peak.
 
     Attributes:
+        mode: High-level streaming mode.
+            - "off": Disable streaming.
+            - "auto": Try seq_kv_cache with automatic batch-micro fallback.
+            - "batch": Use batch microbatching only.
+            - "seq": Use seq_kv_cache (with fallback).
+            - "hybrid": Combine batch micro + seq_kv_cache.
         enabled: Whether streaming is enabled.
         ref_strategy: Strategy for reference model forward pass.
             - "none": Full reference forward (no streaming).
             - "batch_micro": Microbatch reference forward by batch dimension.
             - "seq_kv_cache": Sequence chunking via KV cache.
-        ref_microbatch_size: Microbatch size for batch_micro strategy.
+        ref_microbatch_size: Microbatch size for batch_micro or seq_kv_cache.
         seq_chunk_size: Chunk size for seq_kv_cache strategy (e.g., 128-512).
         kl_token_chunk_size: Optional extra chunking of valid tokens for KL.
         force_fp32_kl: Whether to force FP32 for KL computation.
     """
 
+    mode: Optional[Literal["off", "auto", "batch", "seq", "hybrid"]] = None
     enabled: bool = False
     ref_strategy: Literal["none", "batch_micro", "seq_kv_cache"] = "none"
     ref_microbatch_size: Optional[int] = None
@@ -184,11 +191,16 @@ def fast_cross_entropy_loss_per_token(
     # Create valid mask before computing loss
     valid_mask = labels != ignore_index
 
+    labels_for_kernel = labels
+    if ignore_index != -100:
+        labels_for_kernel = labels.clone()
+        labels_for_kernel[labels_for_kernel == ignore_index] = -100
+
     # Compute per-token CE using Unsloth's Triton kernel
     # The kernel already handles ignore_index (-100) internally and returns 0 for those
     losses = Fast_CrossEntropyLoss.apply(
         logits,
-        labels,
+        labels_for_kernel,
         logit_softcapping,
         logit_scaling,
     )
@@ -379,10 +391,14 @@ def _compute_dft_weights(
     logit_softcapping: float = 0,
     logit_scaling: float = 0,
     ignore_index: int = -100,
+    ce_losses: Optional[torch.Tensor] = None,
+    valid_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """Compute DFT weights: probability of target token under current model.
 
     w = p(label) where p = softmax(effective_logits), detached.
+    If ce_losses is provided, compute weights via exp(-ce_losses) to avoid
+    a full softmax over the vocab.
 
     Args:
         logits: Model logits (B*T, V) or (B, T, V).
@@ -391,10 +407,18 @@ def _compute_dft_weights(
         logit_softcapping: Softcapping value.
         logit_scaling: Scaling value.
         ignore_index: Index to ignore.
+        ce_losses: Optional per-token CE losses aligned with labels.
+        valid_mask: Optional mask of valid tokens (same shape as labels).
 
     Returns:
         DFT weights of same shape as labels, detached.
     """
+    if ce_losses is not None:
+        weights = torch.exp(-ce_losses.detach())
+        if valid_mask is not None:
+            weights = weights * valid_mask
+        return weights
+
     # Flatten if 3D
     if logits.dim() == 3:
         batch, seq_len, vocab_size = logits.shape
@@ -436,6 +460,22 @@ def _unwrap_reference_outputs(
     return ref_outputs, None
 
 
+def _slice_batch_inputs(
+    forward_inputs: Dict[str, Any],
+    batch_size: int,
+    b_start: int,
+    b_end: int,
+) -> Dict[str, Any]:
+    """Slice batch-first tensors for microbatch processing."""
+    mb_inputs = {}
+    for key, value in forward_inputs.items():
+        if torch.is_tensor(value) and value.shape[0] == batch_size:
+            mb_inputs[key] = value[b_start:b_end]
+        else:
+            mb_inputs[key] = value
+    return mb_inputs
+
+
 def _compute_kl_batch_micro(
     model: nn.Module,
     cur_logits: torch.Tensor,
@@ -475,12 +515,7 @@ def _compute_kl_batch_micro(
         b_end = min(b_start + microbatch_size, batch_size)
 
         # Slice inputs for microbatch
-        mb_inputs = {}
-        for key, value in forward_inputs.items():
-            if torch.is_tensor(value) and value.shape[0] == batch_size:
-                mb_inputs[key] = value[b_start:b_end]
-            else:
-                mb_inputs[key] = value
+        mb_inputs = _slice_batch_inputs(forward_inputs, batch_size, b_start, b_end)
 
         # Get reference logits for microbatch
         ref_outputs_mb = ref_forward(**mb_inputs)
@@ -518,6 +553,8 @@ def _compute_kl_seq_kv_cache(
     ref_forward: Callable,
     forward_inputs: Dict[str, Any],
     seq_chunk_size: int,
+    microbatch_size: Optional[int] = None,
+    allow_auto_microbatch_fallback: bool = True,
     logit_softcapping: float = 0,
     logit_scaling: float = 0,
     force_fp32: bool = True,
@@ -535,6 +572,8 @@ def _compute_kl_seq_kv_cache(
         ref_forward: Reference forward callable.
         forward_inputs: Forward inputs (without labels).
         seq_chunk_size: Size of each sequence chunk.
+        microbatch_size: Optional microbatch size for batch dimension.
+        allow_auto_microbatch_fallback: Allow automatic microbatch fallback on errors.
         logit_softcapping: Softcapping value.
         logit_scaling: Scaling value.
         force_fp32: Whether to use FP32 for KL.
@@ -544,39 +583,38 @@ def _compute_kl_seq_kv_cache(
     """
     batch_size, seq_len, vocab_size = cur_logits.shape
     device = cur_logits.device
+
+    if microbatch_size is not None:
+        microbatch_size = max(1, microbatch_size)
+    if microbatch_size is not None and microbatch_size < batch_size:
+        kl = torch.zeros(batch_size, seq_len, dtype = torch.float32, device = device)
+        for b_start in range(0, batch_size, microbatch_size):
+            b_end = min(b_start + microbatch_size, batch_size)
+            mb_inputs = _slice_batch_inputs(
+                forward_inputs, batch_size, b_start, b_end
+            )
+            kl_mb = _compute_kl_seq_kv_cache(
+                model,
+                cur_logits[b_start:b_end],
+                shift_labels[b_start:b_end],
+                valid_mask[b_start:b_end],
+                ref_forward,
+                mb_inputs,
+                seq_chunk_size,
+                microbatch_size = None,
+                allow_auto_microbatch_fallback = False,
+                logit_softcapping = logit_softcapping,
+                logit_scaling = logit_scaling,
+                force_fp32 = force_fp32,
+            )
+            if kl_mb.dim() == 1:
+                mb_batch = b_end - b_start
+                kl_mb = kl_mb.view(mb_batch, -1)
+            kl[b_start:b_end] = kl_mb
+        return kl
+
     kl = torch.zeros(batch_size, seq_len, dtype = torch.float32, device = device)
 
-    # Try to get the underlying model for KV cache support
-    underlying_model = model
-    if hasattr(model, "model"):
-        underlying_model = model.model
-    elif hasattr(model, "base_model"):
-        if hasattr(model.base_model, "model"):
-            underlying_model = model.base_model.model
-        else:
-            underlying_model = model.base_model
-
-    # Check if model supports KV cache
-    supports_cache = hasattr(underlying_model, "config") and getattr(
-        underlying_model.config, "use_cache", True
-    )
-
-    if not supports_cache:
-        # Fallback to full forward
-        ref_outputs = ref_forward(**forward_inputs)
-        ref_logits, _ = _unwrap_reference_outputs(ref_outputs)
-        kl_full = _compute_kl_divergence(
-            cur_logits,
-            ref_logits,
-            model,
-            logit_softcapping,
-            logit_scaling,
-            force_fp32,
-        )
-        if kl_full.dim() == 1:
-            kl_full = kl_full.view(batch_size, seq_len)
-        return kl_full
-
     # Process in chunks with KV cache
     past_key_values = None
 
@@ -616,7 +654,30 @@ def _compute_kl_seq_kv_cache(
                 ref_outputs
             )
             if ref_past_key_values is None and s_end < seq_len:
-                # Can't continue without cache; fall back to full forward
+                # Can't continue without cache; fall back to batch micro if allowed
+                fallback_microbatch = None
+                if allow_auto_microbatch_fallback:
+                    fallback_microbatch = (
+                        microbatch_size
+                        if microbatch_size is not None
+                        else max(1, batch_size // _DEFAULT_REF_MICROBATCH_DIVISOR)
+                    )
+                if (
+                    fallback_microbatch is not None
+                    and fallback_microbatch < batch_size
+                ):
+                    return _compute_kl_batch_micro(
+                        model,
+                        cur_logits,
+                        shift_labels,
+                        valid_mask,
+                        ref_forward,
+                        forward_inputs,
+                        fallback_microbatch,
+                        logit_softcapping,
+                        logit_scaling,
+                        force_fp32,
+                    )
                 ref_outputs = ref_forward(**forward_inputs)
                 ref_logits, _ = _unwrap_reference_outputs(ref_outputs)
                 kl_full = _compute_kl_divergence(
@@ -653,9 +714,32 @@ def _compute_kl_seq_kv_cache(
             del ref_logits_chunk
 
         except (RuntimeError, ValueError, KeyError, TypeError) as e:
-            # Fallback to full forward on KV cache errors
+            # Fallback to batch micro or full forward on KV cache errors
             # These exceptions typically indicate the model doesn't support
             # the chunked KV cache approach (e.g., missing past_key_values support)
+            fallback_microbatch = None
+            if allow_auto_microbatch_fallback:
+                fallback_microbatch = (
+                    microbatch_size
+                    if microbatch_size is not None
+                    else max(1, batch_size // _DEFAULT_REF_MICROBATCH_DIVISOR)
+                )
+            if (
+                fallback_microbatch is not None
+                and fallback_microbatch < batch_size
+            ):
+                return _compute_kl_batch_micro(
+                    model,
+                    cur_logits,
+                    shift_labels,
+                    valid_mask,
+                    ref_forward,
+                    forward_inputs,
+                    fallback_microbatch,
+                    logit_softcapping,
+                    logit_scaling,
+                    force_fp32,
+                )
             ref_outputs = ref_forward(**forward_inputs)
             ref_logits, _ = _unwrap_reference_outputs(ref_outputs)
             kl_full = _compute_kl_divergence(
@@ -713,6 +797,24 @@ def compute_asft_loss(
     if streaming_config is None:
         streaming_config = ASFTStreamingConfig()
 
+    # Resolve streaming mode (new API) vs legacy enabled/ref_strategy
+    mode = streaming_config.mode
+    if mode is None:
+        streaming_enabled = streaming_config.enabled
+        ref_strategy = streaming_config.ref_strategy
+    else:
+        if mode == "off":
+            streaming_enabled = False
+            ref_strategy = "none"
+        elif mode == "batch":
+            streaming_enabled = True
+            ref_strategy = "batch_micro"
+        elif mode in ("seq", "auto", "hybrid"):
+            streaming_enabled = True
+            ref_strategy = "seq_kv_cache"
+        else:
+            raise ValueError(f"Unknown streaming mode: {mode}")
+
     # Get model config for softcapping/scaling
     config = getattr(model, "config", None)
     logit_softcapping = 0
@@ -767,7 +869,13 @@ def compute_asft_loss(
     elif asft_mode == "dft":
         # DFT: CE weighted by model confidence
         dft_weights = _compute_dft_weights(
-            logits, shift_labels, model, logit_softcapping, logit_scaling
+            logits,
+            shift_labels,
+            model,
+            logit_softcapping,
+            logit_scaling,
+            ce_losses = ce_losses,
+            valid_mask = valid_mask,
         )
         dft_weights = dft_weights.view(batch_size, seq_len)
         token_loss = ce_losses * dft_weights
@@ -775,7 +883,7 @@ def compute_asft_loss(
     elif asft_mode in ("sft+kl", "asft"):
         # Need KL divergence
         needs_outputs = (
-            streaming_config.enabled and streaming_config.ref_strategy == "seq_kv_cache"
+            streaming_enabled and ref_strategy == "seq_kv_cache"
         )
         ref_forward = get_reference_forward_callable(
             model,
@@ -786,7 +894,7 @@ def compute_asft_loss(
 
         # Compute KL based on streaming strategy
         # Use local variables to avoid mutating the input config
-        if streaming_config.enabled and streaming_config.ref_strategy == "batch_micro":
+        if streaming_enabled and ref_strategy == "batch_micro":
             ref_microbatch_size = streaming_config.ref_microbatch_size
             if ref_microbatch_size is None:
                 ref_microbatch_size = max(
@@ -804,12 +912,16 @@ def compute_asft_loss(
                 logit_scaling,
                 streaming_config.force_fp32_kl,
             )
-        elif (
-            streaming_config.enabled and streaming_config.ref_strategy == "seq_kv_cache"
-        ):
+        elif streaming_enabled and ref_strategy == "seq_kv_cache":
             seq_chunk_size = streaming_config.seq_chunk_size
             if seq_chunk_size is None:
                 seq_chunk_size = _DEFAULT_SEQ_CHUNK_SIZE
+            ref_microbatch_size = streaming_config.ref_microbatch_size
+            if mode == "hybrid" and ref_microbatch_size is None:
+                ref_microbatch_size = max(
+                    1, batch_size // _DEFAULT_REF_MICROBATCH_DIVISOR
+                )
+            allow_auto_microbatch_fallback = True
             kl = _compute_kl_seq_kv_cache(
                 model,
                 logits,
@@ -818,9 +930,11 @@ def compute_asft_loss(
                 ref_forward,
                 forward_inputs,
                 seq_chunk_size,
-                logit_softcapping,
-                logit_scaling,
-                streaming_config.force_fp32_kl,
+                microbatch_size = ref_microbatch_size,
+                allow_auto_microbatch_fallback = allow_auto_microbatch_fallback,
+                logit_softcapping = logit_softcapping,
+                logit_scaling = logit_scaling,
+                force_fp32 = streaming_config.force_fp32_kl,
             )
         else:
             # Full reference forward
@@ -843,7 +957,13 @@ def compute_asft_loss(
         else:
             # Full ASFT: DFT + KL
             dft_weights = _compute_dft_weights(
-                logits, shift_labels, model, logit_softcapping, logit_scaling
+                logits,
+                shift_labels,
+                model,
+                logit_softcapping,
+                logit_scaling,
+                ce_losses = ce_losses,
+                valid_mask = valid_mask,
             )
             dft_weights = dft_weights.view(batch_size, seq_len)
             dft_loss = ce_losses * dft_weights