Lesson 33 Custom Training
Theme:

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:

  1. 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.
  2. 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.
  3. 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.
  4. 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

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:

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:

  1. Collection: Gather raw text from web crawls (CommonCrawl), domain-specific sources, books, academic papers, and code repositories.
  2. 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.
  3. 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.
  4. Safety filtering: Remove toxic content, personally identifiable information (PII), and content that violates licensing terms.
  5. Tokenisation: For from-scratch training, train a custom BPE tokeniser on your data mix. For continued pre-training, reuse the base model's tokeniser.
  6. 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:

FactorHow to Estimate
GPU hoursFLOPs / (GPU TFLOPS × utilisation). A100 = ~312 TFLOPS BF16; typical utilisation = 40–50%.
Cloud costGPU hours × per-GPU-hour rate. A100 80GB ~$2/hr on-demand; H100 ~$3.50/hr.
StorageTraining data + checkpoints. Budget 2–5 TB for a mid-scale run.
NetworkingInter-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:

Further Reading

This is the final lesson of the course. Congratulations on completing "Claude & LLMs: From Developer to ML Practitioner"!