xla

5174f9ff - Syncfree SGD optimizer for AMP (#3145)

Syncfree SGD optimizer for AMP (#3145) * Initial commit of sync_free_sgd * Refactor gradscaler and sgd_optimizer_step * Address PR comments * Fix multiply by nan bug * Address CR comments