GPU Performance Optimizations

[fatsmcgee]

This PR introduces the following optimizations, which in total make small batch training a bit faster (batch=256 on a T4 is about 15-20% faster) and large batch training a lot faster:

• Asynchronously load the next batch to the GPU while forward/backward pass is occurring.
• Set fused=True for AdamW (small but non-trivial speedup)
• In Gated MLP use torch.mean(..., keepdim=True) instead of broadcasting with torch.mean(...)[:,None,:] (small but non-trivial speedup)