Add AMP for TF Albert
# What does this PR do?

This PR adds the following features to TF Albert:

- AMP compliancy
- Loss computation for TFAlbertForPreTraining
- Cleaning source code
I can split this PR into two different ones, but the one on AMP will be very short (only two single line to update, see the review above). Are you agree with a that tiny PR? If it is still ok, I will split this one^^<|||||>It's ok, thanks for showing me the changes!<|||||>@patrickvonplaten feel free to merge if it looks ok for you!