OpenScholar

torchtune is a PyTorch library designed for authoring, fine-tuning, and experimenting with Large Language Models (LLMs). It is built as a native-PyTorch library emphasizing simplicity, extensibility, and correctness.

Key features include:

• Modular LLM Implementations: Native-PyTorch implementations of popular models.
• Model Interoperability: Utilities for converting checkpoints from popular model zoos.
• Training Recipes: Pre-defined pipelines for various fine-tuning techniques.
• Ecosystem Integration: Works with Hugging Face Datasets for training data and EleutherAI's Eval Harness for evaluation.
• Distributed Training: Support for FSDP (Fully Sharded Data Parallel).
• Configuration: Uses YAML files to manage training runs without code changes.

QLoRA (Quantized LoRA) is an enhancement of LoRA that reduces memory usage by quantizing the frozen base model parameters into a 4-bit NormalFloat (NF4) data type.

Key concepts:

• Storage vs. Compute Dtypes: QLoRA uses a low-precision 'storage dtype' (4-bit NF4) for base model parameters to save memory, while using a higher-precision 'compute dtype' (typically fp32 or bf16) for activations, gradients, and optimizer states to maintain accuracy.
• NF4 Abstraction: torchtune utilizes the NF4Tensor abstraction from the torchao library to implement these quantized components.
• Memory Savings: By quantizing base parameters, you can achieve 4-8x less parameter memory usage compared to standard LoRA, where both base parameters and adapters are held in higher precision.

The torchtune.data module provides various templates to format data for instruct and chat prompts. These templates ensure that datasets are correctly formatted for specific models and instruction styles.

Available templates include:

• InstructTemplate
• AlpacaInstructTemplate
• GrammarErrorCorrectionTemplate
• SummarizeTemplate
• StackExchangedPairedTemplate
• PromptTemplate
• PromptTemplateInterface
• ChatMLTemplate

To mix different types of data (e.g., combining instruction datasets with chat datasets), use the torchtune.datasets.ConcatDataset interface. In a YAML configuration, you can provide a list of dataset components under the dataset key.

dataset:
  - _component_: torchtune.datasets.instruct_dataset
    source: vicgalle/alpaca-gpt4
    template: torchtune.data.AlpacaInstructTemplate
    split: train
    train_on_input: True
  - _component_: torchtune.datasets.instruct_dataset
    source: samsum
    template: torchtune.data.SummarizeTemplate
    column_map:
      output: summary
    split: train
    train_on_input: False
  - _component_: torchtune.datasets.chat_dataset
    ...

torchtune relies on two primary abstractions to manage training workflows:

Configs

Configs are YAML files used to configure training settings and hyperparameters without modifying the underlying Python code. They allow you to specify:

• Dataset settings
• Model settings
• Checkpoint settings
• Hyperparameters (e.g., batch_size, learning_rate)

Recipes

Recipes are targeted, end-to-end pipelines for training and optionally evaluating LLMs. A recipe implements a specific training method (such as full fine-tuning) and applies a set of optimized features to a specific model family (such as Llama2). Common features included in recipes are:

• FSDP (Fully Sharded Data Parallel)
• Activation Checkpointing
• Gradient Accumulation
• Reduced Precision training

QLoRA significantly reduces peak memory usage compared to standard LoRA.

• Model Initialization: QLoRA reduces peak memory by approximately 35%.
• Training: QLoRA reduces peak memory by approximately 40%.

To compare, you can run standard LoRA using:

tune run lora_finetune_single_device --config llama2/7B_lora_single_device

Monitor the logs for GPU peak memory reserved during initialization and every 100 iterations to observe the savings.

To avoid duplication and ensure consistency, you can use interpolations to reference the value of another field within the same YAML config. The instantiate API will automatically resolve these references.

Example: referencing an output_dir in a logger configuration:

output_dir: /tmp/alpaca-llama2-finetune
metric_logger:
  _component_: torchtune.utils.metric_logging.DiskLogger
  log_dir: ${output_dir}

torchtune checkpointers are designed to be "state-dict invariant." This means they manage the complexities of different model weight formats automatically:

• Loading: torchtune accepts checkpoints from multiple sources and formats (e.g., Meta or Hugging Face) without requiring manual conversion.
• Saving: torchtune produces checkpoints in the same format as the source. It converts the internal state dict back to the original form, including splitting keys and weights across the correct number of files.

This invariance allows you to use fine-tuned checkpoints from torchtune with other post-training tools (quantization, evaluation, inference) that support the original format without additional conversion scripts.

torchtune Checkpointers handle two distinct checkpointing scenarios:

1. End-of-training Checkpointing: Writes model weights to file at the end of a completed run. The output files maintain the same keys and partitioning (number of files) as the original input checkpoint.
2. Mid-training Checkpointing: Used for resuming training. In addition to model weights, it outputs a recipe_state.pt file (typically at the end of each epoch) containing optimizer state, epoch count, and other metadata. To prevent directory flooding, recipe_state.pt is overwritten at the end of each epoch.

Data Formats:

• Model weights: A dictionary of {"key": weight}.
• Recipe State: A dictionary containing {"optimizer": ..., "epoch": ..., ...}.

To handle conversational data for specific model architectures, use the provided ChatFormat implementations. These define how messages are structured for models like ChatML, Llama2, or Mistral.

Available formats include:

• ChatFormat
• ChatMLFormat
• Llama2ChatFormat
• MistralChatFormat