Fix a few issues found during fine-tuning (#2582)

* Fix wandb log

* fix validation loss is not logged

batch_idx got reset for each epoch.
use the global step counter instead

* LR should only be updated per step, not per step+ per epoch

* add early stopping

* Fix bf16 handling

scaler is only needed for fp16

* more logs

---------

Co-authored-by: Tony Mak <tony@Tonys-MacBook-Air-1800.local>