Add Z-Loss Support for LM Pretraining

Currently Z-Loss is only supported for RM Training. It shouldn't be too hard to also add in support for LM-Pretraining. Creating this issue to track progress on this.