Enable torch.autocast with ZeRO

[tohtana]

This PR supports torch.autocast with ZeRO. You need to enable torch.autocast in DeepSpeed config.

"torch_autocast": {
  "enabled": true,
  "dtype": "bfloat16"
}

You don't need to explicitly call torch.autocast in your code. The grad scaler is also applied in the DeepSpeed optimizer.

All the parameters are maintained in FP32 but certain operators and layers are computed in the specified dtype. With ZeRO enabled, the communication (allreduce/reduce_scatter) for specified layers are also done in the specified dtype (List of modules).
You cannot enable fp16 or bf16 with torch.autocast.

(Currently working on ZeRO3)