Reformat ASFT test and loss code to improve readability

Reformats function signatures and multi-line expressions in `test_asft.py` and `unsloth/losses/asft.py` to improve readability. Splits long function signatures across multiple lines, uses parenthesized context managers for multiple `patch()` calls, and breaks complex conditionals/expressions. No functional changes.

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

@@ -40,6 +40,7 @@ def run(args):
     from trl import SFTTrainer, SFTConfig
     from unsloth import is_bfloat16_supported
     from unsloth.models.loader_utils import prepare_device_map
+    from unsloth.trainer import ASFTTrainer, ASFTStreamingConfig
     import logging
     from unsloth import RawTextDataLoader
 
@@ -155,13 +156,33 @@ def run(args):
         packing = args.packing,
     )
 
-    # Initialize trainer
-    trainer = SFTTrainer(
-        model = model,
-        processing_class = tokenizer,
-        train_dataset = dataset,
-        args = training_args,
-    )
+    # Initialize trainer - use ASFTTrainer if ASFT is enabled
+    asft_enabled = getattr(args, "asft", False)
+    if asft_enabled:
+        # Build ASFT streaming config
+        asft_streaming = ASFTStreamingConfig(
+            mode = getattr(args, "asft_streaming", None),
+            ref_microbatch_size = getattr(args, "ref_microbatch_size", None),
+            seq_chunk_size = getattr(args, "seq_chunk_size", None),
+        )
+        trainer = ASFTTrainer(
+            model = model,
+            processing_class = tokenizer,
+            train_dataset = dataset,
+            args = training_args,
+            asft_enabled = True,
+            asft_mode = getattr(args, "asft_mode", "asft"),
+            kl_weight = getattr(args, "kl_weight", 0.0),
+            reference_policy = getattr(args, "reference_policy", "disable_adapter"),
+            asft_streaming = asft_streaming,
+        )
+    else:
+        trainer = SFTTrainer(
+            model = model,
+            processing_class = tokenizer,
+            train_dataset = dataset,
+            args = training_args,
+        )
 
     trainer.train()
 
@@ -206,7 +227,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!"
     )
@@ -469,5 +490,73 @@ if __name__ == "__main__":
         "--stride", type = int, default = 512, help = "Overlap between chunks"
     )
 
+    # ASFT Options
+    asft_group = parser.add_argument_group(
+        "🎯 ASFT Options",
+        "Anchored Supervised Fine-Tuning loss configuration (off by default)",
+    )
+    asft_group.add_argument(
+        "--asft",
+        action = "store_true",
+        help = "Enable ASFT (Anchored Supervised Fine-Tuning) loss computation",
+    )
+    asft_group.add_argument(
+        "--asft_mode",
+        type = str,
+        default = "asft",
+        choices = ["sft", "dft", "sft+kl", "asft"],
+        help = (
+            "ASFT loss mode: 'sft' (standard CE), 'dft' (CE weighted by confidence), "
+            "'sft+kl' (CE + KL from reference), 'asft' (DFT + KL). Default: 'asft'"
+        ),
+    )
+    asft_group.add_argument(
+        "--kl_weight",
+        type = float,
+        default = 0.0,
+        help = "Weight for KL divergence term in sft+kl and asft modes. Default: 0.0",
+    )
+    asft_group.add_argument(
+        "--reference_policy",
+        type = str,
+        default = "disable_adapter",
+        choices = ["disable_adapter", "frozen_copy"],
+        help = (
+            "How to compute reference distribution: 'disable_adapter' (use model with LoRA disabled), "
+            "'frozen_copy' (use frozen deepcopy). Default: 'disable_adapter'"
+        ),
+    )
+    asft_group.add_argument(
+        "--asft_streaming",
+        nargs = "?",
+        const = "auto",
+        default = "off",
+        choices = ["off", "auto", "batch", "seq", "hybrid"],
+        help = (
+            "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 or seq_kv_cache strategy",
+    )
+    asft_group.add_argument(
+        "--seq_chunk_size",
+        type = int,
+        default = None,
+        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/__init__.py

@@ -0,0 +1,33 @@
+# 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.
+
+"""Loss functions for Unsloth training."""
+
+from .asft import (
+    ASFTStreamingConfig,
+    compute_asft_loss,
+    effective_logits,
+    fast_cross_entropy_loss_per_token,
+    build_shift_labels,
+    get_reference_forward_callable,
+)
+
+__all__ = [
+    "ASFTStreamingConfig",
+    "compute_asft_loss",
+    "effective_logits",
+    "fast_cross_entropy_loss_per_token",
+    "build_shift_labels",
+    "get_reference_forward_callable",
+]

unsloth/trainer.py

@@ -17,7 +17,8 @@ import os
 import psutil
 import warnings
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Literal, Optional, Union
+from copy import deepcopy
 from functools import wraps
 
 import trl
@@ -40,9 +41,17 @@ from unsloth_zoo.hf_utils import get_transformers_model_type
 from unsloth_zoo.utils import Version
 import dataclasses
 
+# Import ASFT components
+from unsloth.losses.asft import (
+    ASFTStreamingConfig,
+    compute_asft_loss,
+)
+
 __all__ = [
     "UnslothTrainingArguments",
     "UnslothTrainer",
+    "ASFTTrainer",
+    "ASFTStreamingConfig",
     "unsloth_train",
     "_patch_trl_trainer",
     "UnslothVisionDataCollator",
@@ -198,6 +207,130 @@ class UnslothTrainer(SFTTrainer):
         return self.optimizer
 
 
+class ASFTTrainer(UnslothTrainer):
+    """Trainer with ASFT (Anchored Supervised Fine-Tuning) loss support.
+
+    ASFT provides alternative loss functions that weight tokens based on
+    model confidence and/or maintain similarity to a reference model.
+
+    When asft_enabled=False (default), this trainer behaves identically
+    to UnslothTrainer/SFTTrainer with no changes to loss computation.
+
+    Attributes:
+        asft_enabled: Whether to use ASFT loss computation.
+        asft_mode: Loss mode ("sft", "dft", "sft+kl", "asft").
+        kl_weight: Weight for KL divergence term.
+        reference_policy: How to get reference distribution.
+        asft_streaming: Streaming configuration for VRAM reduction.
+    """
+
+    def __init__(
+        self,
+        *args,
+        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.
+
+        Args:
+            *args: Positional arguments for parent trainer.
+            asft_enabled: Whether to enable ASFT loss. Default False preserves
+                          standard SFT behavior completely unchanged.
+            asft_mode: Loss computation mode:
+                - "sft": Standard cross-entropy (for debugging/comparison)
+                - "dft": CE weighted by model's token probability
+                - "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)
+
+        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
+
+    def compute_loss(self, model, inputs, return_outputs = False, **kwargs):
+        """Compute loss with optional ASFT path.
+
+        When asft_enabled=False, delegates entirely to parent compute_loss.
+        When asft_enabled=True, uses ASFT loss computation.
+
+        Args:
+            model: The model to compute loss for.
+            inputs: Input dictionary.
+            return_outputs: Whether to return model outputs.
+            **kwargs: Additional arguments.
+
+        Returns:
+            Loss tensor, or (loss, outputs) tuple if return_outputs=True.
+        """
+        # If ASFT is disabled, use standard path unchanged
+        if not self.asft_enabled:
+            return super().compute_loss(
+                model, inputs, return_outputs = return_outputs, **kwargs
+            )
+
+        num_items_in_batch = kwargs.get("num_items_in_batch")
+        if num_items_in_batch is not None:
+            inputs["num_items_in_batch"] = num_items_in_batch
+
+        if self.asft_mode in ("sft+kl", "asft"):
+            needs_frozen_copy = self.reference_policy == "frozen_copy" or (
+                self.reference_policy == "disable_adapter"
+                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)
+
+        # ASFT-enabled path
+        return compute_asft_loss(
+            model = model,
+            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,
+        )
+
+
 # From `trl>=0.13.0`, they changed how to pass several params to the trainer
 # We need to patch to make the transition smooth
 def _resolve_trainer_params(trainer_class, init_fn):