Module 11: Training on Shakespeare

Learning Objectives

Train a large model on the complete Shakespeare corpus
Apply gradient accumulation for effective large batches
Implement refinement (fine-tuning) on specific sub-corpora
Evaluate with perplexity and qualitative generation

11.1 Dataset

Property	Value
Total size	~5.4 MB of text
Estimated tokens (BPE, ~4 chars/token)	~1.35M tokens
Unique characters	~80
Lines	~124,000

Haskell — corpus analysis

analyzeDataset :: String -> IO ()
analyzeDataset text = do
  let charCount   = length text
      uniqueChars = length (nub text)
      lineCount   = length (lines text)
  putStrLn $ "Total characters:  " ++ show charCount
  putStrLn $ "Unique characters: " ++ show uniqueChars
  putStrLn $ "Lines:             " ++ show lineCount
  putStrLn $ "Est. BPE tokens:   " ++ show (charCount `div` 4)

11.2 Large Model Architecture

Haskell — large model config

largeConfig :: Config
largeConfig = Config
  { vocabSize     = 50257   -- Standard GPT vocabulary (BPE)
  , contextLength = 1024
  , d_model       = 768
  , numHeads      = 12
  , numLayers     = 24
  , ffnDim        = 3072    -- 4× d_model
  , dropout       = 0.1
  }

-- Rough parameter count
calculateParamCount :: Config -> Int
calculateParamCount cfg =
  let embeddingParams    = vocabSize cfg * d_model cfg
      attnParams         = d_model cfg * (d_model cfg + 1) * 4
      ffnParams          = d_model cfg * ffnDim cfg * 2
                         + ffnDim cfg + d_model cfg
      layerNormParams    = d_model cfg * 2 * 2
      blockParams        = attnParams + ffnParams + layerNormParams
      totalTransformer   = blockParams * numLayers cfg
      outputParams       = d_model cfg * vocabSize cfg
  in embeddingParams + totalTransformer + outputParams
  -- ≈ 500M–700M parameters for this config

11.3 Training Configuration

Haskell — distributed training config

data DistributedTrainingConfig = DistributedTrainingConfig
  { globalBatchSize      :: Int      -- Effective batch (via gradient accumulation)
  , microBatchSize       :: Int      -- Per GPU / per step
  , gradAccumSteps       :: Int      -- globalBatchSize / microBatchSize
  , numEpochs            :: Int
  , learningRate         :: Double
  , warmupSteps          :: Int
  , checkpointInterval   :: Int
  }

shakespeareTrainingConfig :: DistributedTrainingConfig
shakespeareTrainingConfig = DistributedTrainingConfig
  { globalBatchSize    = 1024
  , microBatchSize     = 8         -- Adjust to fit VRAM
  , gradAccumSteps     = 128       -- 1024 / 8
  , numEpochs          = 5
  , learningRate       = 6e-4
  , warmupSteps        = 2000
  , checkpointInterval = 1000
  }

11.4 Full Training Script

Haskell — main training entry point

main :: IO ()
main = do
  putStrLn "Loading Shakespeare corpus..."
  shakesText <- readFile "shakespeare.txt"

  let tokenizer = buildTokenizer shakesText
      tokens    = tokenizeText tokenizer shakesText
  putStrLn $ "Total tokens: " ++ show (length tokens)

  let config         = largeConfig
      trainingConfig = shakespeareTrainingConfig
      examples       = createTrainingExamples (contextLength config) tokens
      batches        = createBatches (microBatchSize trainingConfig)
                                     (contextLength config) examples
  putStrLn $ "Batches: " ++ show (length batches)

  putStrLn "Initialising model..."
  model <- initializeModel config
  putStrLn $ "Parameters: " ++ show (calculateParamCount config)

  let initState = TrainingState
        { model      = model
        , optimizer  = adamOptimizer (learningRate trainingConfig)
        , step       = 0
        , loss       = 0
        , tokenCount = 0
        }

  putStrLn "Starting training..."
  (finalModel, finalState) <- trainWithAccumulation initState batches trainingConfig

  saveCheckpoint "checkpoints" finalModel (step finalState)
  putStrLn $ "Done. Steps: " ++ show (step finalState)
  putStrLn $ "Final loss: " ++ show (loss finalState)

11.5 Refinement Training (Fine-Tuning)

After pretraining on the full corpus, optionally fine-tune on specific sub-corpora (e.g., sonnets only, or a specific play) at a much lower learning rate:

Haskell — refinement training

refinementTrain :: LanguageModel -> String -> IO LanguageModel
refinementTrain baseModel specificText = do
  let tokens  = tokenizeText tokenizer specificText
      examples = createTrainingExamples contextLength tokens
      batches  = createBatches microBatchSize contextLength examples

  let refinedState = TrainingState
        { model     = baseModel
        , optimizer = adamOptimizer 1e-5    -- 0.01×–0.1× base LR
        , step      = 0
        , loss      = 0
        }

  -- Only a few epochs — enough to shift distribution, not enough to forget
  (finalModel, _) <- trainLoop refinedState batches 3
  return finalModel

11.6 Evaluation

Haskell — perplexity

-- Perplexity = exp(average negative log-likelihood)
-- Lower is better; random baseline = vocab_size
perplexity :: LanguageModel -> [Batch] -> Double
perplexity model batches =
  let losses  = map (computeBatchLoss model) batches
      avgLoss = average losses
  in exp avgLoss

-- Throughput
tokensPerSecond :: Int -> Double -> Double
tokensPerSecond tokenCount elapsedSeconds =
  fromIntegral tokenCount / elapsedSeconds

Expected Results

Stage	Loss	Perplexity
Random init (vocab=50257)	~10.8	~50257
After 1 epoch	~2.5–3.0	~12–20
After 5 epochs	~1.5–2.0	~4.5–7.4

Generated Example (after full training)

JULIET:
  O Romeo, Romeo! wherefore art thou Romeo?
  Deny thy father and refuse thy name;
  Or, if thou wilt not, be but sworn my love,
  And I'll no longer be a Capulet.

ROMEO:
  Shall I hear more, or shall I speak at this?

Reading & Viewing

Essential papers

Language Models are Unsupervised Multitask Learners (GPT-2, 2019) Radford, Wu, Child, Luan, Amodei, Sutskever PDF

Language Models are Few-Shot Learners (GPT-3, 2020) Brown et al. arxiv.org/abs/2005.14165

Fine-tuning

Fine-Tuning Pretrained Language Models: Weight Initializations, Data Orders, and Early Stopping arxiv.org/abs/2002.06305

Training the Full Shakespeare Model