Apex.amp support

[BramVanroy]

Describe the solution you'd like
It would be great if Apex.amp is supported.

Additional context
If interested, I can work on this. I don't think it should be too much of a problem. I'd replace fp16 to amp instead of simply halfing everything.

[tholor]

Hey @BramVanroy,

Great idea! Apex.amp seems much more user friendly & powerful. I was always a bit annoyed of all specialties introduced into the code by apex fp16 - and the drop in performance ;)

We would highly appreciate, if you could work on this and raised a PR. It should be mainly changes in the Trainer and in initialize_optimizer(). I am also happy to support in case any questions pop up.

After a quick look, I also realized that the current implementation of fp16 in the master branch is not compatible with the latest apex version anyway.

Instead of fixing this we should directly move to apex.amp!