torchtune

Provides pre-built, composable training recipes (full fine-tuning, LoRA, QLoRA, DPO, PPO, knowledge distillation) that encapsulate complete training workflows with built-in support for distributed training, checkpointing, and metric logging. Each recipe is a targeted end-to-end pipeline that combines model loading, data processing, training loop, and evaluation into a single executable unit registered in a recipe registry system.

Implements a configuration layer that uses YAML files to specify all training parameters (model, optimizer, data, scheduler, etc.) with support for CLI overrides and dynamic component instantiation. The system resolves component dependencies, instantiates objects from configuration specs, and enables parameter sweeps without code modification. Configuration files support inheritance and composition patterns for reusability.

Provides a metric logging abstraction that integrates with popular experiment tracking platforms (Weights & Biases, TensorBoard, MLflow) to log training metrics (loss, accuracy, learning rate, gradient norms) at configurable intervals. Metrics are logged from all distributed ranks and aggregated, with support for custom metrics via callback hooks. Logging is decoupled from training logic via a logger interface.

Provides utilities to convert model weights between different formats (HuggingFace safetensors, PyTorch .pt, GGUF) and handle weight name mapping between different implementations. Conversion handles layer name mismatches, missing keys, and shape incompatibilities. Supports downloading models from HuggingFace Hub and converting them to torchtune format.

Provides inference utilities for generating text from fine-tuned models with support for KV-cache (key-value cache) optimization to reduce memory and compute during autoregressive generation. Supports sampling strategies (greedy, top-k, top-p, temperature), beam search, and batch generation. KV-cache is automatically managed and reused across generation steps to avoid redundant computation.

Provides a command-line interface (`tune run`) that executes recipes with YAML configuration files and supports parameter overrides via CLI arguments. The CLI handles argument parsing, configuration merging, and recipe instantiation without requiring Python code. Supports downloading models and datasets via `tune download` command with progress tracking.

Implements multiple attention mechanisms including standard multi-head attention, grouped query attention (GQA) for reduced KV-cache memory, and integration with flash attention kernels for faster computation. Attention implementations are configurable per model and support both training and inference modes with proper gradient computation. Flash attention is automatically used when available, falling back to standard attention otherwise.

Integrates PyTorch's FSDP for distributed training across multiple GPUs/nodes with automatic model sharding, gradient accumulation for larger effective batch sizes, and activation checkpointing to reduce memory footprint. The training infrastructure handles device placement, synchronization, and checkpoint saving across distributed processes transparently through the recipe system.

Provides native PyTorch implementations of LoRA (Low-Rank Adaptation), QLoRA (quantized LoRA), and DoRA (Dominant-Rank Adaptation) that reduce trainable parameters from millions to thousands while maintaining model quality. These implementations integrate with the model builders to selectively apply adapters to attention and linear layers, with support for rank, alpha, and dropout configuration per layer.

Provides clean, hackable PyTorch implementations of popular LLM architectures (Llama 2/3, Gemma, Mistral, Phi, Qwen) with support for modern features like grouped query attention (GQA), flash attention, and KV-cache optimization. Models are built via builder functions that instantiate architectures from configuration specs, enabling easy modification of layer counts, hidden dimensions, and attention mechanisms.

Implements a data pipeline that converts raw datasets (JSON, CSV) into training-ready token sequences using a message-based templating system. The pipeline supports custom prompt templates, multi-turn conversation formatting, and dataset builders that handle tokenization, padding, and batching. Messages are structured as role-content pairs (user, assistant, system) and rendered into prompt templates with configurable formatting.

Provides a checkpointing system that saves and restores model weights, optimizer state, training step counters, and random number generator states across distributed training processes. Checkpoints are saved asynchronously to avoid blocking training, with support for multiple checkpoint formats and automatic cleanup of old checkpoints. The system handles rank synchronization to ensure all processes save consistently.

Implements DPO as a complete training recipe that optimizes models to prefer chosen responses over rejected responses without explicit reward modeling. The recipe handles preference pair formatting, loss computation (Bradley-Terry model), and integration with the standard training infrastructure (distributed training, checkpointing, logging). DPO is configured via YAML like other recipes.

Provides quantization-aware training recipes that simulate quantization during training, allowing models to adapt to lower precision (INT8, INT4) before deployment. QAT uses fake quantization (simulating quantization without actually quantizing) during forward passes, enabling gradient flow through quantization operations. Supports both weight-only and activation quantization configurations.

Implements knowledge distillation as a training recipe where a smaller student model learns from a larger teacher model's outputs. The recipe handles teacher-student forward passes, KL divergence loss computation between teacher and student logits, and temperature-scaled softmax for soft targets. Supports both online distillation (teacher and student trained together) and offline distillation (fixed teacher).