diff --git a/Llama3_1_(8B)_Alpaca-ASFT.ipynb b/Llama3_1_(8B)_Alpaca-ASFT.ipynb index 723b29f4f3..987db4d3c0 100644 --- a/Llama3_1_(8B)_Alpaca-ASFT.ipynb +++ b/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))" ] }, diff --git a/tests/test_asft.py b/tests/test_asft.py index 15bcc207cc..fd9244e7b9 100644 --- a/tests/test_asft.py +++ b/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"]) diff --git a/tests/test_unsloth_cli.py b/tests/test_unsloth_cli.py new file mode 100644 index 0000000000..97d998816f --- /dev/null +++ b/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" diff --git a/unsloth-cli.py b/unsloth-cli.py index 392aeff224..1e895b852f 100644 --- a/unsloth-cli.py +++ b/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) diff --git a/unsloth/losses/asft.py b/unsloth/losses/asft.py index 664f4cf34c..e3c8c8b67d 100644 --- a/unsloth/losses/asft.py +++ b/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