Training Optimization and Stability Deep Dive
Training frontier-scale models is a controlled fight against instability, inefficiency, and bad extrapolation from small experiments.
Primary sources:
The Training Objective
For pretraining, the model minimizes average next-token cross-entropy:
L = - mean_t log P(x_t | x_<t)
At scale, the goal is not just low training loss. It is stable validation loss improvement under a fixed compute/data/hardware budget.
Optimizer
Most transformer LLMs use Adam or AdamW variants.
Adam tracks:
- First moment of gradients.
- Second moment of gradients.
- Adaptive per-parameter updates.
AdamW decouples weight decay from the adaptive update, making regularization behavior cleaner.
Learning Rate Schedule
Common pattern:
- Warmup.
- Peak learning rate.
- Decay, often cosine or linear.
Why warmup matters:
- Early gradients can be unstable.
- Optimizer moments are not calibrated yet.
- Large models can diverge if learning rate starts too high.
Batch Size And Gradient Accumulation
Effective batch size:
microbatch_size * gradient_accumulation_steps * data_parallel_workers
Tradeoffs:
- Larger batches improve hardware utilization.
- Too-large batches can hurt optimization or require LR retuning.
- Gradient accumulation saves memory but changes throughput.
Mixed Precision
Training often uses BF16/FP16 with FP32 master states or carefully managed accumulations.
Why:
- Lower memory.
- Faster tensor cores.
Risks:
- Overflow/underflow.
- Loss scaling problems in FP16.
- Numerically unstable kernels.
BF16 is often preferred when hardware supports it because it has wider exponent range.
Gradient Clipping
Gradient clipping caps update magnitude:
if norm(g) > threshold:
g = g * threshold / norm(g)
Why:
- Prevents rare spikes from destroying training.
- Especially useful during unstable phases.
Initialization And Parameterization
At scale, hyperparameters that work for one model size may fail for another.
muP asks for parameterization and hyperparameter transfer rules that make small-model experiments more predictive of large-model behavior.
Research importance:
- If small experiments do not transfer, architecture research becomes expensive guessing.
Normalization And Residual Stability
Pre-norm, RMSNorm, residual scaling, and initialization all affect deep training stability.
Failure signs:
- Loss spikes.
- NaNs.
- Activation explosion.
- Dead experts in MoE.
- Attention entropy collapse.
- Sudden validation regression.
Data Pipeline Stability
Training instability is not always math.
Watch for:
- Bad documents.
- Tokenization bugs.
- Duplicated shards.
- Corrupt samples.
- Wrong loss masks.
- Data mixture drift.
- Eval contamination.
Distributed Training
Frontier models do not fit on one GPU.
Parallelism types:
- Data parallelism: replicate model, split batch.
- Tensor parallelism: split matrix operations.
- Pipeline parallelism: split layers.
- Sequence/context parallelism: split sequence dimension.
- Expert parallelism: split MoE experts.
Systems notes:
- Communication can dominate.
- Checkpointing must be fault-tolerant.
- Optimizer state is huge.
- Activation checkpointing trades compute for memory.
ZeRO-Style Memory Partitioning
ZeRO reduces memory by partitioning:
- Optimizer states.
- Gradients.
- Parameters.
This makes larger models trainable across devices.
Training Run Monitoring
Track:
- Training loss.
- Validation loss by domain.
- Gradient norm.
- Learning rate.
- Token throughput.
- MFU/hardware utilization.
- Attention/activation stats.
- Data mixture proportions.
- Loss spikes and recovery.
- Checkpoint evals.
Research Questions
- Which small-scale training signals predict large-scale success?
- How much instability is architecture versus data?
- Can optimizer choices improve reasoning, not just loss?
- How should hyperparameters transfer across dense and MoE models?
- Which training failures are detectable early?