Training speed ups: +10-15% tokens/second/gpu

[achalddave]

This PR provides 3 changes to improve the tokens/second/gpu for open_lm. Each change is guarded by a flag that is off by default, since these changes likely affect model convergence and downstream accuracy. We can enable them once we train large enough models with these changes.

Below, I'm quoting rough improvements in tokens/second/gpu for a 7b model, batch size 16, 2048 token length on a machine with 8 80gb A100s.

1. Low precision layer norm: 3340 -> 3470
2. Fused cross entropy: 3470 -> 3490
3. SwiGLU -> GELU: 3490 -> 3690
4. Use bf16 for gradient reduction precision: no improvement on 1 node, but likely important for multi-node

These tests were done with xformers==0.0.19 and torch==2.0.0+cu117. Further upgrading to xformers==0.0.22 (+200t/s/g) and torch==2.0.1+cu118 (+100t/s/g) gets us to about 4000 tokens/second/gpu on one node.

Next steps: Multi-node. With the above changes, our 2 node performance is at ~3700 t/s/g, but this goes up to ~3840 t/s/g when we double the batch size to 32. Can we close the gap to 4k? Might be hard since we don't have the same EFA as mosaic reports here: https://github.com/mosaicml/llm-foundry/tree/main/scripts/train/benchmarking#a100-80gb-with-1600-gbps-node-node-interconnect-roce

Closes #22 #27

[sagadre]

maybe attach a training curve / eval perplexity for some small scale experiments (e.g., 87m params, 1.7B tokens), just to make sure there is no degradation in performance?

[achalddave]

Some results on the 160m param scale (this is with the earlier initialization from mosaic and bias in the mlps). At least low precision layer norm and fused cross-entropy have basically no effect on train loss / val ppl, while giving us a good speedup. Unclear if gelu and float16 reduce are affecting convergence, we will test at 1b next.