RMSNorm

Normalisation is an essential part of LLMs. In the GPT-2 code, we use LayerNorm in a bunch of different places -- in each Transformer block, we do it just before we feed our incoming context vectors into the multi-head attention module, and then again before we feed the result through the FFN: ...and we also do one more round just before we project our context vectors into vocab space via the output head. We do that to make the model more trainable; it helps to prevent vanishing or exploding gradients, by constraining the vectors going into our attention, feed-forward and output head modules so that they have a reasonably stable set of statistical properties. And when (as part of my black-box reimplementation of GPT-2 in JAX) I tried training an almost-complete model without LayerNorm, training completely failed after about 100 steps -- the loss went up to the kind of level that I'd expect from an uninitialised model, and stayed there despite another 700 or so training steps. (Later…