sft_lora.py

# Copyright 2020-2025 The HuggingFace 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.

# /// script
# dependencies = [
#     "trl",
#     "peft",
#     "trackio",
#     "kernels",
# ]
# ///

"""
LoRA-optimized SFT script following best practices from "LoRA Without Regret" (Schulman et al. 2025).
Source: https://thinkingmachines.ai/blog/lora/

Key findings implemented:
1. Apply LoRA to ALL weight matrices (not just attention) for optimal performance
2. Use learning rates similar to full fine-tuning - LoRA can match full FT sample efficiency
3. Choose rank based on dataset size and type
4. LoRA uses ~67% of FLOPs compared to full fine-tuning while matching performance

# Small instruction dataset (10K-100K examples, e.g., Capybara)
# Use rank 32-64 with all linear layers
```
python trl/scripts/sft_lora.py \
    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
    --dataset_name trl-lib/Capybara \
    --learning_rate 2.0e-5 \
    --num_train_epochs 1 \
    --packing \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --gradient_checkpointing \
    --eos_token '<|im_end|>' \
    --eval_strategy steps \
    --eval_steps 100 \
    --use_peft \
    --lora_r 64 \
    --lora_alpha 16 \
    --lora_target_modules all-linear \
    --output_dir Qwen2-0.5B-SFT-LoRA \
    --push_to_hub
```

# Medium instruction dataset (100K-1M examples, e.g., Tulu3)
# Use rank 64-128 with all linear layers
```
python trl/scripts/sft_lora.py \
    --model_name_or_path meta-llama/Llama-3.2-1B-Instruct \
    --dataset_name allenai/tulu-3-sft-mixture \
    --learning_rate 2.0e-5 \
    --num_train_epochs 1 \
    --packing \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --gradient_checkpointing \
    --eval_strategy steps \
    --eval_steps 100 \
    --use_peft \
    --lora_r 128 \
    --lora_alpha 16 \
    --lora_target_modules all-linear \
    --output_dir Llama-3.2-1B-Tulu3-LoRA \
    --push_to_hub
```

# Large reasoning dataset (>1M examples, e.g., OpenThoughts)
# Use rank 256+ with all linear layers
```
python trl/scripts/sft_lora.py \
    --model_name_or_path Qwen/Qwen2.5-3B-Instruct \
    --dataset_name open-thoughts/OpenThoughts \
    --learning_rate 2.0e-5 \
    --num_train_epochs 1 \
    --packing \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 16 \
    --gradient_checkpointing \
    --eval_strategy steps \
    --eval_steps 100 \
    --use_peft \
    --lora_r 256 \
    --lora_alpha 16 \
    --lora_target_modules all-linear \
    --output_dir Qwen2.5-3B-OpenThoughts-LoRA \
    --push_to_hub
```

# With quantization (QLoRA) for memory efficiency
# Note: May need slightly higher learning rate with quantization
```
python trl/scripts/sft_lora.py \
    --model_name_or_path meta-llama/Llama-3.1-8B-Instruct \
    --dataset_name trl-lib/Capybara \
    --learning_rate 3.0e-5 \
    --num_train_epochs 1 \
    --packing \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 16 \
    --gradient_checkpointing \
    --eval_strategy steps \
    --eval_steps 100 \
    --use_peft \
    --lora_r 64 \
    --lora_alpha 16 \
    --lora_target_modules all-linear \
    --load_in_4bit \
    --bnb_4bit_quant_type nf4 \
    --use_bnb_nested_quant \
    --output_dir Llama-3.1-8B-QLoRA \
    --push_to_hub
```
"""

import argparse
import os
from typing import Optional

from accelerate import logging
from datasets import load_dataset
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from transformers.models.auto.modeling_auto import (
    MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES,
)

from trl import (
    DatasetMixtureConfig,
    ModelConfig,
    ScriptArguments,
    SFTConfig,
    SFTTrainer,
    TrlParser,
    get_dataset,
    get_kbit_device_map,
    get_peft_config,
    get_quantization_config,
)


logger = logging.get_logger(__name__)

# Enable logging in a Hugging Face Space
os.environ.setdefault("TRACKIO_SPACE_ID", "trl-trackio")


def validate_lora_config(
    model_args, training_args, dataset_size_estimate: Optional[int] = None
):
    """
    Validate and provide guidance on LoRA configuration based on best practices.

    Based on "LoRA Without Regret" (Schulman et al. 2025):
    https://thinkingmachines.ai/blog/lora/
    """
    if not model_args.use_peft:
        return

    # Check if target_modules is set appropriately
    if model_args.lora_target_modules is None:
        logger.warning(
            "⚠️  No lora_target_modules specified. For best performance, set --lora_target_modules all-linear "
            "to apply LoRA to ALL weight matrices (not just attention). Research shows that attention-only "
            "LoRA underperforms even when using higher rank to match parameter count."
        )
    elif isinstance(model_args.lora_target_modules, (list, str)):
        target_str = str(model_args.lora_target_modules)
        if "all-linear" not in target_str and "q_proj" in target_str:
            logger.warning(
                "⚠️  Detected attention-only LoRA configuration. For best performance, use --lora_target_modules "
                "all-linear to apply LoRA to ALL weight matrices including MLP layers. This significantly improves "
                "performance compared to attention-only LoRA."
            )

    # Check learning rate
    if training_args.learning_rate > 5e-4:
        logger.warning(
            f"⚠️  Learning rate {training_args.learning_rate} seems high. Research shows LoRA can match full "
            f"fine-tuning sample efficiency with similar learning rates (typically 1e-5 to 5e-5). Consider "
            f"reducing the learning rate unless you have a specific reason for this choice."
        )

    # Provide rank guidance based on dataset size if available
    if dataset_size_estimate:
        if dataset_size_estimate < 10_000 and model_args.lora_r > 64:
            logger.info(
                f"💡 Dataset size estimate: {dataset_size_estimate} examples. For small datasets (<10K), "
                f"rank 32-64 is typically sufficient. Current rank: {model_args.lora_r}"
            )
        elif 10_000 <= dataset_size_estimate < 1_000_000 and model_args.lora_r < 32:
            logger.warning(
                f"⚠️  Dataset size estimate: {dataset_size_estimate} examples. For medium datasets (10K-1M), "
                f"rank 64-128 is recommended for best performance. Current rank: {model_args.lora_r} may be too low."
            )
        elif dataset_size_estimate >= 1_000_000 and model_args.lora_r < 128:
            logger.warning(
                f"⚠️  Dataset size estimate: {dataset_size_estimate} examples. For large datasets (>1M), "
                f"rank 256+ is recommended for best performance. Current rank: {model_args.lora_r} may be too low."
            )

    # Check batch size - LoRA may be less tolerant of very large batch sizes
    total_batch_size = (
        training_args.per_device_train_batch_size
        * training_args.gradient_accumulation_steps
        * training_args.world_size
    )
    if total_batch_size > 256:
        logger.warning(
            f"⚠️  Large effective batch size detected: {total_batch_size}. Research shows LoRA may be less "
            f"tolerant of very large batch sizes compared to full fine-tuning. Consider reducing batch size "
            f"if you observe suboptimal performance."
        )

    # Log configuration summary
    logger.info("=" * 80)
    logger.info("LoRA Configuration Summary (based on 'LoRA Without Regret'):")
    logger.info(f"  Rank (r): {model_args.lora_r}")
    logger.info(f"  Alpha: {model_args.lora_alpha}")
    logger.info(f"  Alpha/r ratio: {model_args.lora_alpha / model_args.lora_r:.2f}")
    logger.info(f"  Target modules: {model_args.lora_target_modules}")
    logger.info(f"  Dropout: {model_args.lora_dropout}")
    logger.info(f"  Learning rate: {training_args.learning_rate}")
    logger.info(f"  Effective batch size: {total_batch_size}")
    logger.info(
        f"  Quantization: {'4-bit' if model_args.load_in_4bit else '8-bit' if model_args.load_in_8bit else 'None'}"
    )
    logger.info("=" * 80)


def main(script_args, training_args, model_args, dataset_args):
    ################
    # LoRA Best Practices Information
    ################
    # This script implements best practices from "LoRA Without Regret" (Schulman et al. 2025)
    # Key insights:
    #
    # 1. Apply LoRA to ALL weight matrices for best performance
    #    - Use --lora_target_modules all-linear (not just attention layers)
    #    - Attention-only LoRA underperforms even with matched parameter count
    #
    # 2. Learning rate should be similar to full fine-tuning
    #    - LoRA can achieve same sample efficiency as full FT
    #    - Typical range: 1e-5 to 5e-5 (same as full FT)
    #    - The 1/r scaling in LoRA makes optimal LR approximately rank-independent
    #
    # 3. Choose rank based on dataset size and type:
    #    - Small instruction datasets (<10K examples): rank 32-64
    #    - Medium instruction datasets (10K-1M examples): rank 64-128
    #    - Large reasoning datasets (>1M examples): rank 256+
    #    - RL tasks: rank 8-32 (RL requires very low capacity ~1 bit per episode)
    #
    # 4. Batch size considerations:
    #    - LoRA may be less tolerant of very large batch sizes than full FT
    #    - Keep effective batch size reasonable (<256 for most cases)
    #
    # 5. Computational efficiency:
    #    - LoRA uses ~67% of FLOPs compared to full fine-tuning
    #    - Memory savings enable training on smaller GPU clusters
    #    - Multi-tenant serving: single base model + multiple adapters
    ################

    ################
    # Model init kwargs & Tokenizer
    ################
    model_kwargs = dict(
        revision=model_args.model_revision,
        trust_remote_code=model_args.trust_remote_code,
        attn_implementation=model_args.attn_implementation,
        dtype=model_args.dtype,
    )
    quantization_config = get_quantization_config(model_args)
    if quantization_config is not None:
        # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
        model_kwargs["device_map"] = get_kbit_device_map()
        model_kwargs["quantization_config"] = quantization_config

    # Create model
    config = AutoConfig.from_pretrained(model_args.model_name_or_path)
    valid_image_text_architectures = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values()

    if config.architectures and any(
        arch in valid_image_text_architectures for arch in config.architectures
    ):
        from transformers import AutoModelForImageTextToText

        model = AutoModelForImageTextToText.from_pretrained(
            model_args.model_name_or_path, **model_kwargs
        )
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_args.model_name_or_path, **model_kwargs
        )

    # Create tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        trust_remote_code=model_args.trust_remote_code,
        use_fast=True,
    )

    # Load the dataset
    if dataset_args.datasets and script_args.dataset_name:
        logger.warning(
            "Both `datasets` and `dataset_name` are provided. The `datasets` argument will be used to load the "
            "dataset and `dataset_name` will be ignored."
        )
        dataset = get_dataset(dataset_args)
    elif dataset_args.datasets and not script_args.dataset_name:
        dataset = get_dataset(dataset_args)
    elif not dataset_args.datasets and script_args.dataset_name:
        dataset = load_dataset(
            script_args.dataset_name,
            name=script_args.dataset_config,
            streaming=script_args.dataset_streaming,
        )
    else:
        raise ValueError("Either `datasets` or `dataset_name` must be provided.")

    # Estimate dataset size for validation
    dataset_size_estimate = None
    try:
        if script_args.dataset_train_split in dataset:
            train_dataset = dataset[script_args.dataset_train_split]
            if hasattr(train_dataset, "__len__"):
                dataset_size_estimate = len(train_dataset)
    except Exception:
        pass  # If we can't estimate, that's okay

    # Validate LoRA configuration and provide guidance
    validate_lora_config(model_args, training_args, dataset_size_estimate)

    # Handle eval dataset - check if split exists
    eval_dataset = None
    if training_args.eval_strategy != "no":
        if script_args.dataset_test_split in dataset:
            eval_dataset = dataset[script_args.dataset_test_split]
        else:
            logger.warning(
                f"Evaluation split '{script_args.dataset_test_split}' not found in dataset. "
                f"Available splits: {list(dataset.keys())}. Setting eval_dataset to None."
            )
            # Set eval_strategy to "no" since we don't have an eval dataset
            training_args.eval_strategy = "no"

    # Initialize the SFT trainer
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset[script_args.dataset_train_split],
        eval_dataset=eval_dataset,
        processing_class=tokenizer,
        peft_config=get_peft_config(model_args),
    )

    # Train the model
    logger.info("Starting training...")
    logger.info(
        "💡 Tip: With proper LoRA configuration, you should see similar loss curves to full fine-tuning "
        "while using only ~67% of the compute!"
    )
    trainer.train()

    # Save and push to Hub
    trainer.save_model(training_args.output_dir)
    if training_args.push_to_hub:
        trainer.push_to_hub(dataset_name=script_args.dataset_name)

    logger.info("✅ Training complete!")
    logger.info(
        f"💡 Your LoRA adapter is saved to {training_args.output_dir}. "
        f"It's much smaller than a full model and can be easily shared or loaded for inference."
    )


def make_parser(subparsers: Optional[argparse._SubParsersAction] = None):
    dataclass_types = (ScriptArguments, SFTConfig, ModelConfig, DatasetMixtureConfig)
    if subparsers is not None:
        parser = subparsers.add_parser(
            "sft_lora",
            help="Run the LoRA-optimized SFT training script (following 'LoRA Without Regret' best practices)",
            dataclass_types=dataclass_types,
        )
    else:
        parser = TrlParser(dataclass_types)
    return parser


if __name__ == "__main__":
    parser = make_parser()
    # When using the trl cli, this script may be run with additional arguments, corresponding accelerate arguments.
    # To ensure that their parsing does not interfere with the script arguments, parse the arguments with
    # `return_remaining_strings=True`, then ignore the remaining strings.
    script_args, training_args, model_args, dataset_args, _ = (
        parser.parse_args_and_config(return_remaining_strings=True)
    )
    main(script_args, training_args, model_args, dataset_args)