diff --git a/tests/test_asft.py b/tests/test_asft.py index fd9244e7b9..6428604971 100644 --- a/tests/test_asft.py +++ b/tests/test_asft.py @@ -334,7 +334,7 @@ class TestKLDivergence: cur_logits = torch.randn(4, 8) # (B*T, V) ref_logits = torch.randn(4, 8) - kl = _compute_kl_divergence(cur_logits, ref_logits) + kl = _compute_kl_divergence(cur_logits, ref_logits, kl_direction = "forward") # KL should be non-negative assert torch.all(kl >= -1e-6) # Allow small numerical errors @@ -343,7 +343,7 @@ class TestKLDivergence: """Test that KL is zero when distributions are identical.""" logits = torch.randn(4, 8) - kl = _compute_kl_divergence(logits, logits.clone()) + kl = _compute_kl_divergence(logits, logits.clone(), kl_direction = "forward") # Should be close to zero assert torch.allclose(kl, torch.zeros_like(kl), atol = 1e-5) @@ -353,11 +353,27 @@ class TestKLDivergence: cur_logits = torch.randn(2, 4, 8) # (B, T, V) ref_logits = torch.randn(2, 4, 8) - kl = _compute_kl_divergence(cur_logits, ref_logits) + kl = _compute_kl_divergence(cur_logits, ref_logits, kl_direction = "forward") # Should be flattened to (B*T,) assert kl.shape == (8,) + def test_kl_reverse_matches_manual(self): + """Test reverse KL matches manual computation.""" + torch.manual_seed(321) + cur_logits = torch.randn(2, 5) + ref_logits = torch.randn(2, 5) + + kl_reverse = _compute_kl_divergence( + cur_logits, ref_logits, kl_direction = "reverse" + ) + + cur_p = F.softmax(cur_logits, dim = -1) + ref_p = F.softmax(ref_logits, dim = -1) + manual = (cur_p * (cur_p.log() - ref_p.log())).sum(dim = -1) + + assert torch.allclose(kl_reverse, manual, atol = 1e-5) + # ----------------------------------------------------------------------------- # Test DFT weights computation @@ -439,6 +455,37 @@ class TestComputeASFTLoss: assert loss.dim() == 0 assert loss.requires_grad + def test_dft_normalize_by_weights(self, simple_model): + """Test DFT normalization by weight sum.""" + inputs = { + "input_ids": torch.tensor([[1, 2, 3, 4]]), + "labels": torch.tensor([[1, 2, 3, 4]]), + } + + logits = simple_model(input_ids = inputs["input_ids"]).logits + shift_labels = build_shift_labels(inputs["labels"]) + valid_mask = shift_labels != -100 + ce_losses, _ = fast_cross_entropy_loss_per_token(logits, shift_labels) + ce_losses = ce_losses.view(shift_labels.shape) + dft_weights = _compute_dft_weights( + logits, + shift_labels, + ce_losses = ce_losses, + valid_mask = valid_mask, + ).view(shift_labels.shape) + token_loss = ce_losses * dft_weights + expected = token_loss[valid_mask].sum() / dft_weights[valid_mask].sum().clamp_min(1e-8) + + loss = compute_asft_loss( + simple_model, + inputs, + asft_mode = "dft", + kl_weight = 0.0, + normalize_by = "weights", + ) + + assert torch.allclose(loss, expected, atol = 1e-5) + def test_sft_kl_mode(self, simple_model): """Test SFT+KL mode.""" inputs = { @@ -592,6 +639,7 @@ class TestStreamingModeMapping: logit_softcapping = 0, logit_scaling = 0, force_fp32 = True, + kl_direction = "forward", ): batch, seq_len = shift_labels.shape return torch.zeros(batch, seq_len, device = shift_labels.device) @@ -726,6 +774,7 @@ class TestStreamingModeMapping: logit_softcapping = 0, logit_scaling = 0, force_fp32 = True, + kl_direction = "forward", ): batch, seq_len = ref_logits.shape[:2] return torch.zeros(batch * seq_len, device = ref_logits.device) @@ -1144,8 +1193,10 @@ class TestASFTTrainerComputeLoss: trainer.asft_enabled = True trainer.asft_mode = "sft" trainer.kl_weight = 0.0 + trainer.kl_direction = "forward" trainer.reference_policy = "disable_adapter" trainer.asft_streaming = ASFTStreamingConfig() + trainer.normalize_by = "tokens" trainer._asft_original_model = None model = nn.Module() @@ -1168,8 +1219,10 @@ class TestASFTTrainerComputeLoss: 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["kl_direction"] == "forward" assert loss_mock.call_args.kwargs["reference_policy"] == "disable_adapter" assert loss_mock.call_args.kwargs["streaming_config"] is trainer.asft_streaming + assert loss_mock.call_args.kwargs["normalize_by"] == "tokens" def test_compute_loss_creates_frozen_copy_once(self): """Test frozen copy is created once when needed.""" @@ -1179,8 +1232,10 @@ class TestASFTTrainerComputeLoss: trainer.asft_enabled = True trainer.asft_mode = "asft" trainer.kl_weight = 0.1 + trainer.kl_direction = "forward" trainer.reference_policy = "frozen_copy" trainer.asft_streaming = ASFTStreamingConfig() + trainer.normalize_by = "tokens" trainer._asft_original_model = None model = nn.Module() @@ -1190,7 +1245,7 @@ class TestASFTTrainerComputeLoss: "labels": torch.tensor([[1, 2, 3, 4]]), } - with patch( + with pytest.warns(UserWarning), patch( "unsloth.trainer.deepcopy", return_value = model_copy ) as deepcopy_mock, patch( "unsloth.trainer.compute_asft_loss", @@ -1212,8 +1267,10 @@ class TestASFTTrainerComputeLoss: trainer.asft_enabled = True trainer.asft_mode = "asft" trainer.kl_weight = 0.1 + trainer.kl_direction = "forward" trainer.reference_policy = "disable_adapter" trainer.asft_streaming = ASFTStreamingConfig() + trainer.normalize_by = "tokens" trainer._asft_original_model = None model = MagicMock() diff --git a/unsloth/losses/asft.py b/unsloth/losses/asft.py index e3c8c8b67d..4b4b47bda8 100644 --- a/unsloth/losses/asft.py +++ b/unsloth/losses/asft.py @@ -335,10 +335,18 @@ def _compute_kl_divergence( logit_softcapping: float = 0, logit_scaling: float = 0, force_fp32: bool = True, + kl_direction: Literal["forward", "reverse"] = "forward", ) -> torch.Tensor: - """Compute per-token KL divergence: KL(p_ref || p_cur). + """Compute per-token KL divergence (forward KL by default). - KL(p_ref || p_cur) = sum_i p_ref(i) * (log p_ref(i) - log p_cur(i)) + KL naming is frequently confused due to PyTorch's kl_div signature + (target is the weighting distribution). Definitions here follow standard + math/RLHF convention: + - Forward KL: KL(p_ref || p_cur), expectation over p_ref (mass-covering). + - Reverse KL: KL(p_cur || p_ref), expectation over p_cur (mode-seeking). + The original ASFT repo's paper text says "reverse KL" but its code uses + F.kl_div(log(cur), ref), which is forward KL; we match the code behavior. + Reverse KL is available via kl_direction="reverse". Args: cur_logits: Current model logits (B*T, V) or (B, T, V). @@ -347,6 +355,7 @@ def _compute_kl_divergence( logit_softcapping: Softcapping value. logit_scaling: Scaling value. force_fp32: Whether to compute in FP32 for stability. + kl_direction: "forward" for KL(p_ref || p_cur), "reverse" for KL(p_cur || p_ref). Returns: Per-token KL divergence of shape (B*T,) or (B, T). @@ -367,14 +376,19 @@ def _compute_kl_divergence( cur_eff = cur_eff.float() ref_eff = ref_eff.float() - # Compute log probabilities and probabilities - cur_logp = F.log_softmax(cur_eff, dim = -1) - ref_p = F.softmax(ref_eff, dim = -1) - - # KL(p_ref || p_cur) = sum_i p_ref(i) * (log p_ref(i) - log p_cur(i)) - # Using F.kl_div: kl_div(input=log_cur, target=ref) computes the right thing - # with reduction='none', we get per-element, then sum over vocab - kl = F.kl_div(cur_logp, ref_p, reduction = "none").sum(dim = -1) + if kl_direction == "forward": + # Forward KL: KL(p_ref || p_cur) + cur_logp = F.log_softmax(cur_eff, dim = -1) + ref_p = F.softmax(ref_eff, dim = -1) + # Using F.kl_div: kl_div(input=log_cur, target=ref) computes KL(ref || cur) + kl = F.kl_div(cur_logp, ref_p, reduction = "none").sum(dim = -1) + elif kl_direction == "reverse": + # Reverse KL: KL(p_cur || p_ref) + ref_logp = F.log_softmax(ref_eff, dim = -1) + cur_p = F.softmax(cur_eff, dim = -1) + kl = F.kl_div(ref_logp, cur_p, reduction = "none").sum(dim = -1) + else: + raise ValueError(f"Unknown kl_direction: {kl_direction}") return kl @@ -487,6 +501,7 @@ def _compute_kl_batch_micro( logit_softcapping: float = 0, logit_scaling: float = 0, force_fp32: bool = True, + kl_direction: Literal["forward", "reverse"] = "forward", ) -> torch.Tensor: """Compute KL using batch microbatching strategy. @@ -503,6 +518,7 @@ def _compute_kl_batch_micro( logit_softcapping: Softcapping value. logit_scaling: Scaling value. force_fp32: Whether to use FP32 for KL. + kl_direction: "forward" for KL(p_ref || p_cur), "reverse" for KL(p_cur || p_ref). Returns: KL tensor of shape (B, T). @@ -530,6 +546,7 @@ def _compute_kl_batch_micro( logit_softcapping, logit_scaling, force_fp32, + kl_direction, ) # Reshape if needed @@ -558,6 +575,7 @@ def _compute_kl_seq_kv_cache( logit_softcapping: float = 0, logit_scaling: float = 0, force_fp32: bool = True, + kl_direction: Literal["forward", "reverse"] = "forward", ) -> torch.Tensor: """Compute KL using sequence chunking with KV cache strategy. @@ -577,6 +595,7 @@ def _compute_kl_seq_kv_cache( logit_softcapping: Softcapping value. logit_scaling: Scaling value. force_fp32: Whether to use FP32 for KL. + kl_direction: "forward" for KL(p_ref || p_cur), "reverse" for KL(p_cur || p_ref). Returns: KL tensor of shape (B, T). @@ -606,6 +625,7 @@ def _compute_kl_seq_kv_cache( logit_softcapping = logit_softcapping, logit_scaling = logit_scaling, force_fp32 = force_fp32, + kl_direction = kl_direction, ) if kl_mb.dim() == 1: mb_batch = b_end - b_start @@ -677,6 +697,7 @@ def _compute_kl_seq_kv_cache( logit_softcapping, logit_scaling, force_fp32, + kl_direction, ) ref_outputs = ref_forward(**forward_inputs) ref_logits, _ = _unwrap_reference_outputs(ref_outputs) @@ -687,6 +708,7 @@ def _compute_kl_seq_kv_cache( logit_softcapping, logit_scaling, force_fp32, + kl_direction, ) if kl_full.dim() == 1: kl_full = kl_full.view(batch_size, seq_len) @@ -703,6 +725,7 @@ def _compute_kl_seq_kv_cache( logit_softcapping, logit_scaling, force_fp32, + kl_direction, ) if kl_chunk.dim() == 1: @@ -739,6 +762,7 @@ def _compute_kl_seq_kv_cache( logit_softcapping, logit_scaling, force_fp32, + kl_direction, ) ref_outputs = ref_forward(**forward_inputs) ref_logits, _ = _unwrap_reference_outputs(ref_outputs) @@ -749,6 +773,7 @@ def _compute_kl_seq_kv_cache( logit_softcapping, logit_scaling, force_fp32, + kl_direction, ) if kl_full.dim() == 1: kl_full = kl_full.view(batch_size, seq_len) @@ -768,9 +793,11 @@ def compute_asft_loss( *, asft_mode: Literal["sft", "dft", "sft+kl", "asft"] = "asft", kl_weight: float = 0.0, + kl_direction: Literal["forward", "reverse"] = "forward", reference_policy: Literal["disable_adapter", "frozen_copy"] = "disable_adapter", streaming_config: Optional[ASFTStreamingConfig] = None, original_model: Optional[nn.Module] = None, + normalize_by: Literal["tokens", "weights"] = "tokens", return_outputs: bool = False, ) -> Union[torch.Tensor, Tuple[torch.Tensor, Any]]: """Compute ASFT loss. @@ -786,9 +813,11 @@ def compute_asft_loss( - "sft+kl": CE + KL divergence from reference - "asft": DFT + KL divergence (full ASFT) kl_weight: Weight for KL term (only for sft+kl and asft modes). + kl_direction: "forward" for KL(p_ref || p_cur), "reverse" for KL(p_cur || p_ref). reference_policy: How to get reference distribution. streaming_config: Configuration for streaming strategies. original_model: Optional pre-created frozen reference model. + normalize_by: "tokens" (default, matches reference) or "weights" for DFT/ASFT. return_outputs: Whether to return model outputs alongside loss. Returns: @@ -841,10 +870,10 @@ def compute_asft_loss( # Valid mask and normalization valid_mask = shift_labels != -100 - n_items = inputs.get("num_items_in_batch", None) - if n_items is None: - n_items = valid_mask.sum() - n_items = max(n_items, 1) # Avoid division by zero + n_items_tokens = inputs.get("num_items_in_batch", None) + if n_items_tokens is None: + n_items_tokens = valid_mask.sum() + n_items_tokens = max(n_items_tokens, 1) # Avoid division by zero # Handle edge case: no valid tokens if valid_mask.sum() == 0: @@ -862,6 +891,7 @@ def compute_asft_loss( ce_losses = ce_losses.view(batch_size, seq_len) # Initialize token losses + dft_weights = None if asft_mode == "sft": # Standard SFT: just CE token_loss = ce_losses @@ -911,6 +941,7 @@ def compute_asft_loss( logit_softcapping, logit_scaling, streaming_config.force_fp32_kl, + kl_direction, ) elif streaming_enabled and ref_strategy == "seq_kv_cache": seq_chunk_size = streaming_config.seq_chunk_size @@ -935,6 +966,7 @@ def compute_asft_loss( logit_softcapping = logit_softcapping, logit_scaling = logit_scaling, force_fp32 = streaming_config.force_fp32_kl, + kl_direction = kl_direction, ) else: # Full reference forward @@ -947,6 +979,7 @@ def compute_asft_loss( logit_softcapping, logit_scaling, streaming_config.force_fp32_kl, + kl_direction, ) kl = kl.view(batch_size, seq_len) del ref_logits @@ -972,8 +1005,14 @@ def compute_asft_loss( else: raise ValueError(f"Unknown asft_mode: {asft_mode}") - # Final reduction: sum over valid tokens, divide by n_items - loss = token_loss[valid_mask].sum() / n_items + # Final reduction: sum over valid tokens, divide by chosen normalizer. + normalizer = n_items_tokens + if normalize_by == "weights" and dft_weights is not None: + weight_sum = dft_weights[valid_mask].sum() + normalizer = weight_sum.clamp_min(1e-8) + elif normalize_by != "tokens": + raise ValueError(f"Unknown normalize_by: {normalize_by}") + loss = token_loss[valid_mask].sum() / normalizer if return_outputs: return loss, outputs diff --git a/unsloth/trainer.py b/unsloth/trainer.py index 04b670b7d6..128d8f40ed 100644 --- a/unsloth/trainer.py +++ b/unsloth/trainer.py @@ -230,8 +230,10 @@ class ASFTTrainer(UnslothTrainer): asft_enabled: bool = False, asft_mode: Literal["sft", "dft", "sft+kl", "asft"] = "asft", kl_weight: float = 0.0, + kl_direction: Literal["forward", "reverse"] = "forward", reference_policy: Literal["disable_adapter", "frozen_copy"] = "disable_adapter", asft_streaming: Optional[ASFTStreamingConfig] = None, + normalize_by: Literal["tokens", "weights"] = "tokens", **kwargs, ): """Initialize ASFTTrainer. @@ -246,10 +248,12 @@ class ASFTTrainer(UnslothTrainer): - "sft+kl": CE + KL divergence from reference - "asft": Full ASFT (DFT + KL) kl_weight: Weight for KL term (used in sft+kl and asft modes). + kl_direction: "forward" for KL(p_ref || p_cur), "reverse" for KL(p_cur || p_ref). reference_policy: How to compute reference distribution: - "disable_adapter": Use model with LoRA adapters disabled - "frozen_copy": Use a frozen deepcopy of the model asft_streaming: Optional streaming config for VRAM reduction. + normalize_by: "tokens" (default) or "weights" for DFT/ASFT normalization. **kwargs: Keyword arguments for parent trainer. """ super().__init__(*args, **kwargs) @@ -257,8 +261,10 @@ class ASFTTrainer(UnslothTrainer): self.asft_enabled = asft_enabled self.asft_mode = asft_mode self.kl_weight = kl_weight + self.kl_direction = kl_direction self.reference_policy = reference_policy self.asft_streaming = asft_streaming or ASFTStreamingConfig() + self.normalize_by = normalize_by # Will be lazily initialized if needed self._asft_original_model = None @@ -294,6 +300,18 @@ class ASFTTrainer(UnslothTrainer): and not hasattr(model, "disable_adapter") ) if needs_frozen_copy and self._asft_original_model is None: + if self.reference_policy == "frozen_copy": + warnings.warn( + "Unsloth: Creating a frozen copy of the model for ASFT. " + "This doubles VRAM usage. Use 'disable_adapter' if using LoRA.", + stacklevel = 2, + ) + elif self.reference_policy == "disable_adapter": + warnings.warn( + "Unsloth: 'disable_adapter' is unavailable; falling back to a " + "frozen copy for ASFT. This doubles VRAM usage.", + stacklevel = 2, + ) self._asft_original_model = deepcopy(model) self._asft_original_model.eval() self._asft_original_model.requires_grad_(False) @@ -304,9 +322,11 @@ class ASFTTrainer(UnslothTrainer): inputs = inputs, asft_mode = self.asft_mode, kl_weight = self.kl_weight, + kl_direction = self.kl_direction, reference_policy = self.reference_policy, streaming_config = self.asft_streaming, original_model = self._asft_original_model, + normalize_by = self.normalize_by, return_outputs = return_outputs, )