FP8 Weights

[ClashLuke]

Goal: Low precision weights allow testing a 405B on one H100 node

Currently, the weights are stored in bf16 (see #100 for a fix on that front). However, jax supports other datatypes, and further libraries such as https://github.com/neuralmagic/compressed-tensors support yet more.
Adding support for alternative dtypes (and, ideally, lower-precision datatypes) would significantly simplify testing at scale.