[perf] PyTorch AMP reverts the grads to fp32, meaning that reduce calls are overweight

[blefaudeux]

🚀 Feature

Either through torch amp or via a postfix .half() cast (not too nice), make sure that we don't reduce fp32 grads when they were actually computed on fp16..

Motivation

The gradient size are double what they should be if torch AMP is being used, which can be a bottleneck depending on internode coms

Pitch

free speed

Alternatives

Not doing anything, this does not affect correctness

Additional context