torchtune

torchtune is a PyTorch library designed for authoring, fine-tuning, and experimenting with Large Language Models (LLMs). It is built with a focus on simplicity, correctness, stability, and hardware accessibility.

Key capabilities include:

• Modular LLM Implementations: Native-PyTorch implementations of popular models.
• Model Interoperability: Checkpoint-conversion utilities to work with popular model zoos.
• Training Recipes: Pre-defined pipelines for various fine-tuning techniques.
• Ecosystem Integration: Supports Hugging Face Datasets for training data and EleutherAI's Eval Harness for evaluation.
• Distributed Training: Support for FSDP2.
• Configuration: Uses YAML files to manage training runs without code changes.

torchtune is a PyTorch library designed for authoring, post-training, and experimenting with Large Language Models (LLMs). Key features include:

• Hackable training recipes: Supports SFT, knowledge distillation, DPO, PPO, GRPO, and quantization-aware training (QAT).
• Model Implementations: PyTorch implementations of popular LLMs such as Llama (3.1, 3.2, 3.3, 4), Gemma, Mistral, Phi, and Qwen.
• Performance & Efficiency: Utilizes the latest PyTorch APIs for memory efficiency and scaling.
• YAML Configuration: Uses YAML files to easily configure training, evaluation, quantization, or inference recipes.

The Async GRPO recipe uses Ray to coordinate multiple remote workers to overlap generation and training, maximizing GPU utilization.

Key components include:

• Trainer: Iterates over a torchrl.data.ReplayBuffer (populated by the generator) instead of a standard dataloader.
• Generator: A torchrl.collectors.SyncDataCollector instance that manages a vLLM instance and environments (interactive or dataset-based). It pushes GRPOTrajectory objects to the ReplayBuffer.
• VLLMParameterServer: Manages model parameter syncing between the trainer and generator. The trainer pushes updates every $n$ steps, and the generator pulls updates every $m$ steps.

Note: The VLLMParameterServer currently only supports single-GPU, limiting the recipe to models that fit on a single GPU.

torchtune supports fine-tuning LLMs and VLMs using datasets from the Hugging Face Hub, local files, or remote URLs. It provides built-in dataset builders for common workflows such as:

• Text supervised fine-tuning: Instruct tuning and chat-based tuning.
• Multimodal supervised fine-tuning: Training with images and text.
• RLHF: Preference alignment workflows.
• Continued pre-training: Text completion tasks.
• Custom datasets: Full customizability to train on any data format or schema.

torchtune relies on two primary abstractions to manage training workflows:

Configs

Configs are YAML files used to configure training settings and hyperparameters. They allow you to specify settings for the dataset, model, and checkpoint, as well as hyperparameters like batch size and learning rate, without modifying the underlying Python code.

Recipes

Recipes are targeted end-to-end pipelines for training and optionally evaluating LLMs. A recipe implements a specific training method (e.g., full fine-tuning) and applies a set of features (e.g., FSDP2, Activation Checkpointing, Gradient Accumulation, and Reduced Precision training) to a specific model family (e.g., Llama3.1).

Torchtune recipes are designed as hackable, singularly-focused scripts for interacting with LLMs. Every recipe is composed of three distinct parts:

1. Configurable parameters: Defined via YAML configuration files or provided through command-line overrides.
2. Recipe script: The entry-point script that handles configuration parsing, validation, environment setup, and execution logic.
3. Recipe class: The core logic implementation (e.g., for fine-tuning) exposed through a set of APIs.

LoRA is a parameter-efficient finetuning (PEFT) technique that adds trainable low-rank decomposition matrices to specific layers of a neural network while freezing the original pretrained parameters.

Key Benefits:

• Memory Savings: Substantial reduction in the number of parameters requiring gradients and optimizer states (e.g., using AdamW).
• Efficiency: For a 7B Llama2 model, a rank r=8 can reduce trainable parameters for a projection from ~15M to ~65K (over 99% reduction).

How it works: LoRA replaces weight updates with two matrices, A and B. Matrix A projects inputs down to a small rank r, and matrix B projects them back to the original output dimension. The final output is the sum of the frozen original output and the scaled LoRA output: frozen_out + (alpha / rank) * lora_out.

torchtune checkpointers are composable components designed to be plugged into any recipe, including training, evaluation, or generation. They are built to be "state-dict invariant," meaning they manage the complexities of different weight formats automatically.

Key benefits of the torchtune checkpointer include:

• Format Agnostic Loading: You can load checkpoints from multiple sources and formats without manual conversion.
• Format Preserving Saving: When saving, torchtune converts the state_dict back into the original source format, including splitting keys and weights across the same number of files as the original.
• Ecosystem Interoperability: Because torchtune preserves the source format, fine-tuned checkpoints can be used with other post-training tools (quantization, eval, inference) that support the original format without additional conversion scripts.

The tune CLI is used to download models, list recipes and configs, copy built-in files for customization, run training recipes, validate configurations, and inspect config files. Use --help to see all available subcommands or tune <subcommand> --help for specific command details.

$ tune --help
usage: tune [-h] {download,ls,cp,run,validate,cat} ...

Welcome to the torchtune CLI!

options:
-h, --help            show this help message and exit

subcommands:
  {download,ls,cp,run,validate,cat}
    download            Download a model from the Hugging Face Hub.
    ls                  List all built-in recipes and configs
    ...