LitGPT

Implements minimal-abstraction decoder-only transformer architectures (GPT, Llama, Mistral, Phi, Gemma, Qwen, etc.) using PyTorch with explicit, modifiable code rather than wrapper abstractions. The Config dataclass in litgpt/config.py defines ~100 parameters per model (layer count, embedding dimensions, attention heads, RoPE scaling, GQA variants) that map directly to model instantiation. Supports model sizes from 0.5B to 405B parameters with native support for architectural variants like grouped query attention, sliding window attention, and mixture-of-experts.

Implements Low-Rank Adaptation (LoRA) and Quantized LoRA (QLoRA) fine-tuning via the litgpt/lora.py module, which injects trainable low-rank decomposition matrices (A, B) into attention and linear layers while freezing base model weights. QLoRA variant uses BitsAndBytes 4-bit quantization to reduce base model memory footprint to ~6GB for 70B models. Supports selective layer targeting (e.g., only attention layers or specific transformer blocks) and integrates with PyTorch Lightning's distributed training for multi-GPU LoRA fine-tuning.

Integrates with LitServe (Lightning AI's inference server) to deploy models as HTTP APIs with OpenAI-compatible endpoints (/v1/chat/completions, /v1/completions). Handles request batching, concurrent inference, and automatic scaling across multiple GPUs. Supports streaming responses (Server-Sent Events), request validation, and error handling. Models can be served with quantization, LoRA adapters, or full precision, with automatic device placement and memory management.

Integrates with EleutherAI's lm-evaluation-harness to run standardized benchmarks (MMLU, HellaSwag, ARC, TruthfulQA, etc.) on trained models. Provides evaluation scripts that load LitGPT checkpoints, apply prompt formatting, and compute benchmark metrics. Supports both zero-shot and few-shot evaluation, with configurable number of shots and prompt templates. Results are comparable across models and frameworks, enabling reproducible evaluation.

Implements a unified Tokenizer class (litgpt/tokenizer.py) that wraps both HuggingFace Tokenizers and SentencePiece backends, providing a consistent encode/decode interface. Handles special tokens, padding, truncation, and batch tokenization. Supports loading tokenizers from HuggingFace hub or local paths, with automatic caching. Integrates with model-specific tokenizer configurations (e.g., Llama's special tokens, Mistral's chat tokens).

Implements a Config dataclass system (litgpt/config.py) that defines model architectures via ~100 parameters (num_layers, hidden_size, num_heads, etc.) and training hyperparameters (learning_rate, batch_size, warmup_steps). Provides named configurations for 20+ model families (Llama, Mistral, Phi, etc.) that can be loaded by name or customized. Configs are Python dataclasses, enabling IDE autocomplete, type checking, and programmatic manipulation. Supports config serialization to YAML for reproducibility.

Implements a Prompt system (litgpt/prompts.py) that applies model-specific instruction templates for chat and instruction-following tasks. Supports templates for Llama Chat, Mistral Instruct, Phi, Gemma, and other models. Handles multi-turn conversations, system prompts, and automatic token counting. Templates are defined as Python classes with format() methods, enabling transparent prompt construction and debugging.

LitGPT provides a configuration hub (litgpt/config.py) with pre-defined Config dataclasses for 20+ model families (Llama, Mistral, Phi, Gemma, Qwen, Falcon, OLMo, etc.), each specifying ~100 architectural parameters (layer count, embedding dimensions, attention heads, RoPE, GQA, etc.). Named configurations enable one-line model instantiation without manual parameter specification. The hub is extensible — new models can be added by defining a Config dataclass and registering it.

Implements Adapter modules (litgpt/adapter.py and litgpt/adapter_v2.py) that inject small bottleneck layers into transformer blocks, reducing trainable parameters to 0.5-2% of base model size. Adapter V1 uses sequential down-projection → activation → up-projection, while V2 adds parallel residual connections and layer normalization for improved gradient flow. Adapters are inserted after attention and feed-forward layers, allowing task-specific specialization while keeping base weights frozen.

Enables end-to-end fine-tuning of all model parameters using PyTorch Lightning's training loop with automatic mixed precision (AMP) in FP16 or BF16, gradient accumulation for effective larger batch sizes, and gradient checkpointing to reduce activation memory. Integrates with FSDP (Fully Sharded Data Parallel) for multi-GPU distributed training, automatically sharding model weights, gradients, and optimizer states across devices. Supports learning rate scheduling, warmup, and weight decay configuration.

Supports training models from random initialization on custom datasets using PyTorch Lightning's distributed training infrastructure. Handles datasets up to 3 trillion tokens via streaming data loading and checkpoint resumption. Includes TinyLlama pretraining example (1.1B model trained on 3T tokens) demonstrating end-to-end pretraining workflow. Integrates with custom DataModules for flexible data loading (raw text, JSON, Parquet, HuggingFace datasets) and supports data shuffling, tokenization, and batching across multiple GPUs.

Implements convert_hf_checkpoint.py and convert_lit_checkpoint.py scripts that enable seamless conversion between LitGPT's native checkpoint format and HuggingFace Transformers format. Handles weight mapping, layer name translation, and config serialization/deserialization. Supports converting HuggingFace checkpoints (Llama, Mistral, Phi, etc.) into LitGPT format for training, and exporting LitGPT checkpoints to HuggingFace format for ecosystem compatibility (inference with vLLM, deployment with HuggingFace Inference API).

Integrates BitsAndBytes quantization library to reduce model memory footprint via 4-bit (NF4) and 8-bit quantization. 4-bit quantization reduces a 70B model to ~6GB VRAM, enabling single-GPU inference and fine-tuning (QLoRA). Supports mixed precision quantization (e.g., quantize attention layers to 4-bit, keep feed-forward in FP16) and automatic dequantization during forward passes. Quantization is applied at model loading time via BitsAndBytes config, preserving model architecture and enabling standard inference APIs.

Leverages PyTorch Lightning's FSDP (Fully Sharded Data Parallel) integration to automatically shard model weights, gradients, and optimizer states across multiple GPUs or TPUs. Supports both data parallelism (each GPU processes different data) and model parallelism (model layers distributed across devices). Handles gradient synchronization, communication optimization (gradient compression), and automatic checkpoint saving across distributed ranks. Enables training of 405B+ models by combining FSDP with pipeline parallelism.

Implements generation strategies in the inference module supporting greedy decoding (argmax), temperature-scaled sampling, top-k/top-p filtering, and beam search. Handles prompt formatting via the Prompt system (litgpt/prompts.py) which applies model-specific instruction templates (e.g., Llama Chat, Mistral Instruct). Supports streaming generation (token-by-token output), batch generation, and generation with constraints (max_length, stop tokens). Integrates with the LLM Python API for programmatic text generation.

Provides a high-level LLM class that wraps model loading, tokenization, and generation into a simple Python API. Supports loading models from checkpoint paths or HuggingFace hub, automatic device placement (CPU/GPU), and generation via a single generate() method. Integrates with quantization (4-bit, 8-bit) and LoRA adapters transparently. Enables programmatic fine-tuning via the Trainer class, which handles distributed training setup, checkpoint management, and metric logging.