Theory
Reading time: ~20 minutes (~1500 words) · Prerequisites: Modules 4–5 (transformers, fine-tuning), familiarity with PyTorch
1. The Training Spectrum
Throughout this course, we have progressed from prompting to RAG to fine-tuning. Custom training sits at the far end of this spectrum: maximum control, maximum cost, and maximum complexity. Before committing to a training run, you should be confident that simpler approaches cannot achieve your goal.
The decision tree is roughly:
- Prompting + RAG — Can you solve the problem by providing the right context at inference time? If yes, stop here. This is the cheapest and most flexible option.
- Fine-tuning (LoRA or full) — Does the base model understand your domain but need to change its behaviour, style, or output format? Fine-tune on hundreds to thousands of examples.
- Continued pre-training — Does the base model lack fundamental knowledge about your domain? Feed it millions of tokens of domain text to build that knowledge into the weights.
- Training from scratch — Is there no suitable base model for your language, modality, or architectural requirements? Build from the ground up.
Most organisations will never go past step 2. Continued pre-training is relevant for specialised industries (biomedical, legal, finance) or underserved languages. Training from scratch is the province of frontier labs and well-funded research teams.
2. Continued Pre-Training
Continued pre-training (also called domain-adaptive pre-training, or DAPT) takes an existing pre-trained model and trains it further on a domain-specific corpus using the same causal language modelling objective: predict the next token.
The critical insight is that pre-trained models have already learned syntax, reasoning patterns, and general world knowledge. Continued pre-training adds domain-specific knowledge on top of this foundation, much like a doctor who first completes a general medical degree before specialising.
Key Parameters
- Learning rate: Use a significantly lower learning rate than from-scratch training. Typical values are 1e-5 to 5e-5, compared to 1e-4 to 3e-4 for pre-training. Too high a learning rate causes catastrophic forgetting.
- Data mixing: Include 5–20% general-domain data (e.g., a subsample of the original pre-training distribution) alongside your domain data. This acts as a regulariser and prevents the model from "forgetting" how to do basic tasks like follow instructions.
- Duration: One to two epochs over your domain corpus is typical. More epochs risk overfitting, especially if your corpus is small relative to the model size.
- Evaluation: Track both domain-specific perplexity (should decrease) and general benchmarks like MMLU or HellaSwag (should remain stable). If general performance degrades sharply, reduce the learning rate or increase the general-data mixing ratio.
Case Study: CodeLlama
Meta's CodeLlama took the Llama 2 base model and continued pre-training on 500 billion tokens of code. The result was a model that retained Llama 2's natural language abilities while dramatically improving on code generation and understanding. This is the gold standard example of continued pre-training done well.
3. Training Infrastructure
Training a model with billions of parameters requires distributing computation across multiple GPUs. Two dominant frameworks have emerged:
DeepSpeed ZeRO
Microsoft's DeepSpeed library implements the ZeRO (Zero Redundancy Optimiser) family of techniques that partition training state across GPUs:
- Stage 1: Partitions optimiser states (e.g., Adam momentum and variance). Reduces memory by ~4x per GPU.
- Stage 2: Additionally partitions gradients. Further reduces memory footprint.
- Stage 3: Additionally partitions model parameters themselves. Each GPU only stores a shard of the model and gathers parameters on-demand during forward and backward passes.
DeepSpeed also supports CPU offloading (move optimiser states to CPU RAM) and NVMe offloading (use SSD as overflow memory), enabling training of models that would not fit even with ZeRO Stage 3 alone.
FSDP (Fully Sharded Data Parallel)
PyTorch's native answer to ZeRO, FSDP shards model parameters, gradients, and optimiser states across GPUs in a way conceptually similar to ZeRO Stage 3. The advantage is that it requires no external library — it ships with PyTorch 2.x and is tightly integrated with the PyTorch ecosystem. FSDP has been gaining adoption rapidly, particularly in research settings.
Choosing Between Them
Both achieve similar memory savings. DeepSpeed has more configuration options and a longer track record at extreme scale. FSDP is simpler to set up if you are already in the PyTorch ecosystem. HuggingFace Trainer supports both as backends, so switching between them requires only a configuration change.
4. Scaling Laws
One of the most important results in modern ML is that language model performance is predictable. The Chinchilla scaling law (Hoffmann et al., 2022) established that for a given compute budget, there is an optimal balance between model size and training data:
For compute-optimal training, the number of training tokens should be approximately 20 times the number of model parameters.
This means a 7 B parameter model should ideally be trained on ~140 B tokens. The original GPT-3 (175 B parameters) was trained on only 300 B tokens — heavily undertrained by Chinchilla standards, which would recommend ~3.5 T tokens.
A useful approximation for total training compute is:
FLOPs ≈ 6 × N × D
Where N is the number of parameters and D is the number of training tokens. The factor of 6 accounts for the multiply-accumulate operations in both the forward pass (2x) and the backward pass (4x, because gradients must be computed for both weights and activations).
Over-Training on Purpose
In practice, many teams intentionally over-train models relative to Chinchilla optimality. Why? Because inference cost is proportional to model size, not training data. A smaller model trained on more data (like Llama 3 8B trained on 15T tokens) may give better cost-performance at inference time than a larger, compute-optimal model. Scaling laws tell you the training-efficient frontier; your deployment economics determine where you actually want to operate.
5. Training Data Pipelines
Data quality is the single largest determinant of model quality. A well-curated 100 B token dataset will produce a better model than a poorly filtered 1 T token dataset. The standard pipeline has six stages:
- Collection: Gather raw text from web crawls (CommonCrawl), domain-specific sources, books, academic papers, and code repositories.
- Deduplication: Remove near-duplicate documents using MinHash or SimHash. Web data typically contains 30–50% near-duplicates. Training on duplicates wastes compute and can cause memorisation of specific passages.
- Quality filtering: Score documents using a quality classifier (often a small model trained on human ratings) or use perplexity filtering with a reference model. Remove low-quality content.
- Safety filtering: Remove toxic content, personally identifiable information (PII), and content that violates licensing terms.
- Tokenisation: For from-scratch training, train a custom BPE tokeniser on your data mix. For continued pre-training, reuse the base model's tokeniser.
- Packing and shuffling: Concatenate documents (with separator tokens) to create sequences that fill the full context length, minimising wasted padding. Shuffle at the document level before packing.
6. Mixed-Precision Training and Checkpointing
Mixed-precision training uses lower-precision floating-point formats for most operations while keeping critical accumulations in FP32. The current standard is BF16 (Brain Float 16), which has the same exponent range as FP32 (avoiding overflow issues that plague FP16) but half the memory footprint. Master weights and optimiser states remain in FP32; forward and backward passes use BF16.
Checkpointing is essential for long training runs. A complete checkpoint includes model weights, optimiser states, learning rate scheduler state, and RNG seeds. For a 7 B model with Adam optimiser, a full checkpoint is approximately 56 GB (model in FP32 = 28 GB, plus two optimiser buffers of 28 GB each). Save checkpoints every 500–1000 steps and retain the most recent 3–5. Use asynchronous writes so the training loop does not stall while checkpoints are flushed to storage.
At cluster scale (hundreds of GPUs), hardware failures happen regularly. Design your training pipeline to detect failures via heartbeat monitoring, restart from the most recent checkpoint, and optionally continue with fewer nodes using elastic training (TorchElastic or similar).
7. Cost Estimation
Before launching a training run, estimate costs to avoid unpleasant surprises:
| Factor | How to Estimate |
|---|---|
| GPU hours | FLOPs / (GPU TFLOPS × utilisation). A100 = ~312 TFLOPS BF16; typical utilisation = 40–50%. |
| Cloud cost | GPU hours × per-GPU-hour rate. A100 80GB ~$2/hr on-demand; H100 ~$3.50/hr. |
| Storage | Training data + checkpoints. Budget 2–5 TB for a mid-scale run. |
| Networking | Inter-node bandwidth matters at scale. Use InfiniBand or equivalent for multi-node runs. |
8. Future Directions
The field of model training is evolving rapidly. Key trends to watch:
- Mixture of Experts (MoE): Models like Mixtral activate only a subset of parameters per token, enabling much larger total parameter counts at the same inference cost. Training MoE models introduces load-balancing challenges but is becoming increasingly well-understood.
- Multi-modal pre-training: Rather than bolting vision or audio onto a text model after the fact, next-generation models are trained jointly on text, images, audio, and video from the start.
- Synthetic data and distillation: Using strong models to generate training data for smaller, specialised models is becoming a standard part of the training pipeline.
- Efficient architectures: Mamba (state-space models), RWKV, and other non-transformer architectures offer linear-time inference, though transformers remain dominant for now.
Further Reading
- Training Compute-Optimal Large Language Models — Hoffmann et al. (2022). The Chinchilla paper.
- Llama 2: Open Foundation and Fine-Tuned Chat Models — Touvron et al. (2023). Detailed training methodology for Llama 2.
- Code Llama: Open Foundation Models for Code — Rozière et al. (2023). Continued pre-training case study.
- DeepSpeed ZeRO Tutorial — Official documentation for ZeRO stages.
- PyTorch FSDP Tutorial — Getting started with Fully Sharded Data Parallel.
- Mixtral of Experts — Jiang et al. (2024). MoE architecture and training details.
This is the final lesson of the course. Congratulations on completing "Claude & LLMs: From Developer to ML Practitioner"!