Most LLM libraries hand you a wrapper around someone else’s model class and call it a day. LitGPT took the opposite bet: reimplement every popular open model from scratch - Llama 3, Gemma 2/3, Phi 4, Qwen2.5 and Qwen3, Mistral, Mixtral, Falcon, OLMo, DeepSeek-R1 distills and more - in plain, readable PyTorch with zero abstraction layers, then build the whole lifecycle on top. litgpt from Lightning AI is that bet, Apache 2.0 licensed, at version 0.5.13, and it powers real projects: TinyLlama was trained with it, Microsoft’s Samba research built on it, and it was the official starter kit of the NeurIPS 2023 LLM Efficiency Challenge.

The result reads like a textbook that happens to run. One GPT class implements every attention variant, a 3194-line config file encodes every supported architecture as data, and the CLI exposes 24 commands covering download, pretrain, five finetuning flavors, six generation modes, evaluation, and serving. In this tour we read the model core, the LoRA machinery, the generation modes including speculative decoding, and the serving layer. As always in this series, what follows is an educational tour of published source code.

Overview of the litgpt architecture: a 24-command CLI and the LLM facade sit over the config registry; the model core holds attention and MoE machinery; training workflows and inference/serving wrap it from both sides.

Reading the overview from left to right:

• litgpt/model.py is the transformer itself at 1325 lines: GPT starts at line 22 and CausalSelfAttention at line 390, with the MoE block at line 842.
• litgpt/config.py is the architecture registry at 3194 lines: the Config dataclass at line 26 encodes every supported model as plain fields.
• litgpt/api.py is the Python facade at 735 lines: LLM at line 38 with load at line 149, distribute at line 257, and generate at line 462.
• litgpt/pretrain.py is the training loop at 532 lines: fit at line 287 with the cosine schedule helper get_lr at line 463.
• litgpt/finetune/lora.py is the LoRA recipe at 578 lines: fit at line 285 and the parallelize_fn hook at line 522.
• litgpt/lora.py is the PEFT layer zoo at 666 lines: LoRALinear at line 89 and the fused LoRAQKVLinear at line 175.
• litgpt/generate/base.py is the token loop at 559 lines: sample at line 54, next_token at line 77, and generate at line 373.
• litgpt/deploy/serve.py is the serving layer at 324 lines: BaseLitAPI at line 21 and the OpenAI-compatible OpenAISpecLitAPI at line 166.

Why You Need This

If you want to understand transformers rather than rent them, this is the codebase to read. Every model is a single readable file, every trick is commented with its origin paper, and there is no inheritance maze between you and the matmul. The attention block alone is a masterclass: one combined QKV projection, the MHA/GQA/MQA tradeoff explained with an ASCII diagram in the source, rotary embeddings applied per layer, sliding-window masks built by hand, and the KV cache sliced by an input_pos tensor you control. When a paper introduces something new - Qwen3’s grouped expert routing, Gemma’s unit-offset norm, DeepSeek’s latent attention - the implementation lands here in reviewable form within weeks.

Second, it is a complete product, not a pile of scripts. The same config object drives a 1-GPU chat session and a 1000-GPU pretrain run through Lightning Fabric, with FSDP sharding, quantized 4-bit and 8-bit modes, TPU/XLA and Thunder extensions, and validated YAML recipes in the config_hub for known-good hyperparameters. The workflows are symmetric: whatever you train with litgpt finetune, you chat with litgpt chat, evaluate with litgpt evaluate, and deploy with litgpt serve - all reading the same checkpoint folder.

Third, it is the pragmatic middle ground between nanoGPT’s 300 lines and a 100k-file training platform. Compared with the full-scale torchtitan trainer, litgpt trades N-D parallelism composability for approachability; compared with C++ runtimes like the CTranslate2 inference engine, it trades raw throughput for hackability. If your job is research experiments, teaching, or product prototypes where you must touch the model internals, litgpt is the fastest path from idea to measured result.

Detailed view: the facade group shows tokenizer and prompt styles; the model core exposes attention variants, latent attention, routers, RoPE and the KV cache; training and inference groups show the recipe files and generation modes.

Start with the config, because everything else is generated from it. config.py declares the Config dataclass at line 26 with fields spanning block sizes, norm variants (norm_qk with an olmo2 type at lines 39 and 40), sliding_window_size at line 76 with a per-layer indicator list, attention_logit_softcapping at line 81 for Gemma 2, rope_adjustments at line 86 that carries either Llama3 low/high frequency factors or YaRN parameters, and the MoE fields n_expert, n_shared_expert, and n_expert_groups at lines 94 through 98. post_init at line 118 pads the vocabulary to a multiple of 512 for hardware efficiency and derives rope_n_elem from rotary_percentage. The mlp_class property at line 223 maps the string mlp_class_name to the actual class in model.py, which is how one GPT class becomes GPT-NeoX-style, LLaMA-style, Gemma-style, or MoE by config alone. from_checkpoint at line 214 prefers a saved model_config.yaml next to the weights, so checkpoints are self-describing.

The model core is where the reading pays off. model.py implements GPT at line 22 whose forward at line 85 accepts input_pos for cached decoding and input_pos_maxp1 so the mask can be sliced to just the live window - the docstring at lines 92 through 102 explains this avoids graph breaks under torch.compile - and lm_head_chunk_size at line 177 splits the output projection into chunks to cut autograd peak memory. CausalSelfAttention starts at line 390: its forward at line 430 does one fused qkv matmul at line 472, splits heads at lines 490 through 499, applies RoPE at lines 506 through 513 (choosing the interleave variant for models that need it), consults the KVCache at line 519, repeats keys and values for grouped-query balance at lines 536 through 539, builds the sliding-window mask at lines 541 through 563 with a hand-drawn diagram in the comment, and lands in scaled_dot_product_attention at line 576 - the flash-attention path, connecting to the FlashAttention kernel lineage we toured earlier. Exotics live beside it: MultiheadLatentAttention at line 649 implements DeepSeek’s compressed KV attention, GroupedTopkRouter at line 888 serves Qwen3-style expert routing, build_rope_cache at line 939 handles YaRN and Llama3 scaling, apply_rope at line 1144 and apply_rope_interleave at line 1176 are the two rotation conventions, KVCache at line 1221 is a plain nn.Module holding one preallocated buffer, and RMSNorm at line 1301 carries the Gemma add_unit_offset flag.

Training comes in two flavors that share the same spine. pretrain.py runs main at line 176 into fit at line 287: gradient accumulation, validation between steps, cosine decay computed by get_lr at line 463, and weight init at line 477. The finetuning recipes mirror it - finetune/lora.py defines setup at line 49, main at line 183, fit at line 285, and a parallelize_fn at line 522 that applies device-mesh sharding plus activation checkpointing before training - while finetune/full.py runs the same loop without wrappers. The PEFT layers themselves are lora.py: LoRALinear at line 89 freezes the base weight and learns a low-rank delta, LoRAQKVLinear at line 175 does the same for the fused attention projection, mark_only_lora_as_trainable at line 417 flips requires_grad flags by policy, and merge_lora_weights at line 662 folds the deltas back for export. The adapter variants in adapter_v2.py take the parallel-residual route instead. Data is pluggable behind data/base.py: Alpaca, LIMA, TinyStories, raw JSON, and lit-data streaming all satisfy one interface, and prompt formatting is delegated to prompts.py, where PromptStyle at line 17 has 38 subclasses - Llama3 at line 170, the ChatML family from line 371 covering Qwen - resolved by model_name_to_prompt_style at line 452.

Inference is where litgpt quietly ships research-grade features. generate/base.py defines multinomial_num_samples_1 at line 31, top-p sampling at line 39, the sample routine at line 54 combining top-k and top-p, and generate at line 373 that pre-fills the prompt then loops next_token at line 77 with the cached positions. On top of it sit four exotic modes: speculative_decoding.py runs a small draft model ahead of the target and verifies tokens in one forward - the core loop is speculative_decoding at line 61; sequentially.py executes transformer blocks one per GPU for models too big to replicate; tp.py rewrites linears into column/row-parallel shards with tensor_parallel at line 84 and an all-reduce output hook at line 78; and adapter-aware generators serve PEFT checkpoints. The interactive chat/base.py wraps this in a REPL with multiline input via interact at line 124, and deploy/serve.py turns it into an HTTP product with SimpleLitAPI at line 76, StreamLitAPI at line 121, and OpenAISpecLitAPI at line 166 so existing OpenAI clients work unmodified. The shared plumbing is utils.py at 1118 lines: chunked_cross_entropy at line 302 forces the loss to deallocate as it goes, incremental_save at line 249 streams checkpoint pickles without double memory, auto_download_checkpoint at line 698 fetches weights on demand, and validate_checkpoint at line 959 checks integrity before you burn GPU hours. Conversion is handled per family by scripts/convert_hf_checkpoint.py - copy_weights_hf_llama at line 138, copy_weights_qwen_3 at line 641, and the qkv_reassemble helper at line 726 - so Hugging Face weights load as litgpt checkpoints in one command.

From Install to Chat

The happy path is four commands, with weights downloaded automatically:

pip install 'litgpt[extra]'
litgpt chat microsoft/phi-2
litgpt finetune microsoft/phi-2 \
  --data JSON \
  --data.json_path my_custom_dataset.json \
  --out_dir out/custom-model
litgpt serve out/custom-model/final

Or drive it from Python, where the facade hides the whole loop:

from litgpt import LLM

llm = LLM.load("microsoft/phi-2")
text = llm.generate("Fix the spelling: Every fall, the family goes to the mountains.")
print(text)
# Corrected Sentence: Every fall, the family goes to the mountains.

Honest limits: litgpt is single-model focused - it trains and serves one model per process, so there is no continuous batching scheduler, no N-D parallelism mesh, and no distributed stage schedules; very large pretrains will want torchtitan or Megatron-style stacks instead. Attention goes through PyTorch’s SDPA rather than hand-tuned kernels for every architecture, softcapped models fall back to slower masked attention, and the quantization options lean on bitsandbytes rather than shipping their own kernels. The project’s pace of model additions means the config registry grows monthly, so pin versions in production. As a readable reference implementation and a serious finetuning workhorse, though, it has no real competition in its weight class.

Enjoyed this post? Never miss out on future posts by following us