Megatron AMP fix for scheduler step counter
# Bugfix
- Attempt to prevent a step counter update of the scheduler when skipping a step in the optimizer
- Add training steps when nan/inf gradients are detected
- Properly support LR scheduler with "constant lr" at end of training. Refactor CosineAnnealing model for such support.

# Details

Since there are no pre-existing hooks in PTL to modify the scheduler step before the optimizer step, ive had to splice the support into the current GradScaler and then use the on_train_batch_end hook and manually monitor when the optimizer step occurs via the flag  optimizer_update_skipped - None means optimizer didnt step this time (due to accumulated gradient), and True means grads were accumulated but optimizer did not perform step() - skipped step. In such case, perform a no-op for the scheduler too.

Signed-off-by: smajumdar <titu1994@gmail.com>