[ENHANCEMENT] Deepnorm supporting

We are encountering a significant numerical stability issue when training large-scale MoE (Mixture of Experts) models. The model has a large number of layers (e.g., more than 20 layers). Even with standard stabilization techniques employed, including:

Using Input LayerNorm before each sub-layer (e.g., Attention and Feed-Forward MLP layers), adhering to the Pre-LayerNorm architecture.

In some designs, also experimenting with Pre-MLP LayerNorm after the MLP layer.

And the expert load balance loss specific to MoE models.

However, during training, the hidden state activations (embeddings/representations) in the latter several layers exhibit a sharp increase in magnitude ("explosion"), characterized by cumulative growth in their mean and variance layer by layer. This can ultimately lead to exploding gradients, NaN losses, and severely impact model convergence and performance.

To our knowledge, DeepNorm is a technique proven to effectively address this kind of error accumulation and numerical instability caused by residual connections in Transformer architectures. It systematically scales the residual path and the output of transformer sub-layers, fundamentally stabilizing the forward and backward passes of deep models.

Therefore, we highly recommend considering the integration of or support for the DeepNorm functionality within the framework to tackle the stability challenges in training deep MoE models.