FSDP optimizations to try

[epwalsh]

We should see what the perf impact of this all of these changes are.

• All-reduce in FP32. Reduce gradients in fp32 #291
• We might not need the autocast() context manager due to FSDP native mixed precision. We should try without it. This is good by default because it forces certain ops to run in FP32, like CE loss.
• Try different wrapping strategies, like wrapping groups of OlmoBlocks together. This would require refactoring the model a a little bit to make it possible to group more than one blocks together. Add option to FSDP wrap by groups of blocks #340
• Store buffers in FP32. This would make the alibi or rope buffers be full precision, which we might want. We're not actually using buffers, so this doesn't affect us.