Scaling to Larger Models (500M–1B parameters)

9.1 Chinchilla Scaling Laws

How should you allocate a compute budget between model size and training data? The Chinchilla paper (2022) answers this empirically:

9.2 Mixed Precision Training

• Use FP16 (16-bit floats) for forward/backward passes
• Keep FP32 master weights for stability
• Result: ~2× faster, half memory, minimal accuracy loss

data MixedPrecisionTrainer = MixedPrecisionTrainer
  { fp32Weights :: [Double]   -- Master weights
  , fp16Weights :: [Float]    -- Working copy
  , lossScale   :: Double     -- Prevent FP16 underflow
  }

-- Scale loss before backprop to prevent FP16 underflow gradients
-- Unscale gradients before updating FP32 master weights
backpropMixedPrecision :: MixedPrecisionTrainer -> [Float] -> IO MixedPrecisionTrainer
backpropMixedPrecision trainer gradients =
  let scaledGrads  = map (* realToFrac (lossScale trainer)) gradients
      fp32Grads    = map realToFrac scaledGrads :: [Double]
      unscaledGrad = map (/ lossScale trainer) fp32Grads
      newFp32      = zipWith (-) (fp32Weights trainer)
                       (map (* learningRate) unscaledGrad)
      newFp16      = map realToFrac newFp32
  in return trainer { fp32Weights = newFp32, fp16Weights = newFp16 }

9.3 Gradient Accumulation

When GPU memory limits your batch size, accumulate gradients over N micro-batches before updating — effectively simulating a batch N× larger.

data AccumulatedGradients = AccumulatedGradients
  { gradAccum :: [Double]
  , steps     :: Int
  }

accumulateGradient :: AccumulatedGradients -> [Double] -> AccumulatedGradients
accumulateGradient acc grads =
  acc { gradAccum = zipWith (+) (gradAccum acc) grads
      , steps     = steps acc + 1 }

applyAccumulatedGradients
  :: TrainingState -> AccumulatedGradients -> Int -> TrainingState
applyAccumulatedGradients state accum updateFreq =
  if steps accum >= updateFreq
  then let avgGrads = map (/ fromIntegral updateFreq) (gradAccum accum)
       in updateParameters state avgGrads
  else state

9.4 Gradient Checkpointing

Trade compute for memory: instead of storing all intermediate activations during the forward pass, recompute them during backpropagation. This typically reduces memory by 30–60% at the cost of ~33% more compute.

Applies the principle: "it's cheaper to recompute than to store" — useful when VRAM is the bottleneck.

9.5 Multi-GPU: Data Parallelism

-- Each GPU processes a different mini-batch
-- Parameters are identical across GPUs
-- After each step: average gradients across GPUs (all-reduce)

data DataParallelTraining = DataParallelTraining
  { numGPUs        :: Int
  , modelPerGPU    :: [LanguageModel]
  , gradientBuffer :: [Double]
  }

allReduceGradients :: DataParallelTraining -> IO DataParallelTraining
allReduceGradients training =
  -- 1. Collect gradients from all GPUs
  -- 2. Average
  -- 3. Distribute back
  -- (Implementation depends on communication backend)
  undefined

Model parallelism (for models too large for one GPU) splits layers across GPUs. More complex — requires pipeline scheduling to avoid idle GPUs.