Add methods to improve model stability by controlling the magnitude of linear layers outputs (including logits)

Methods drawn from Section 4 of Methods of improving LLM training stability by NVIDIA folks

They experimented with the following methods (derived from different prior works) -

Method 1 - σReparam Method 2 - Softmax temperature (soft_temp) Method 3 - Softmax capping (soft_cap) Method 4 - Softmax clipping (soft_clip) Method 5 - LayerScale Method 6 - QK layer normalization (QK_norm) Method 7 - Combination of QK layer normalization with softmax capping (QK_norm_cap) Method 8 - Layer normalization after QKV layers (QKV_norm) Method 9 - Layer normalization after QK, Proj and FC2 layers (QK_FC_norm)

Their findings -

Based on this and as models like Gemma 4 use Logit Soft-Capping, we should maybe prioritize some of these? I feel like adding soft_cap & soft_clip might be a nice starting point as they can also be composed with existing QK Norm to get QK_norm_cap