[WIP] BFloat16 weights; 2 sec improvement

[leloykun]

From the pytorch trace below, see that the current record is bottlenecked by the backwards pass. Thus, it makes sense to shift our focus on optimizing this part of the run.

Converting the weights to bfloat16 is the lowest hanging fruit we can do & it seems to reduce wallclock time by 2 secs.

[YouJiacheng]

I think a better choice is to also implement a master weight in optimizers so we can get back to FP32 with minimal performance regression when we need FP32.
Note that embeddings are already full BF16, so we only need to deal with lm_head and weights updated by Muon.
BF16 accumulation is safe now (probably) because we only decay the lr to 0.1× peak after previous changes.
But in longer runs (ofc, irrelevant to 3.28 speedrun) we might decay the lr to smaller values.