[P] Out-of-the-box FP8 training (nanoGPT demo)

The latest gen of AI chips can do FP8 compute, but making the most of this isn't straightforward - just naïvely inserting FP8 casts causes training to fail (e.g. grads underflow).

To fix this I've been working on a method called *unit scaling*, which I demo in this notebook: [github.com/graphcore-research/out-of-the-box-fp8-training.ipynb](https://github.com/graphcore-research/out-of-the-box-fp8-training/blob/main/out_of_the_box_fp8_training.ipynb)

With a one-line code change (`model = unit_scale(model)`) FP8 training now matches the loss of FP32.

It works by re-scaling operations in the fwd & bwd pass so that training …

ai chips box change chips code compute demo gen line loss machinelearning making nanogpt operations scaling training