[DeepSeek v3] Add grad mask and update MLA init

[gagika]

Description

When training DeepSeek-V3 671B with load_balance_loss_weight > 0 deterministically NaNs on step 1 (per-layer MLA backward gradient overflows bf16). Adds two opt-in fixes that default to no-op:

1. Per-token gradient mask (grad_mask_threshold): zeros tokens whose feature-axis backward-RMS exceeds threshold.
2. Constant-std MLA init (mla_init_std): N(0, std) with output proj scaled by 1/sqrt(2*num_decoder_layers). No effect when loading a checkpoint.

DeepSeek-V3 model config defaults these to 0.001 and 100.0.

Tests

• E2E test with random weights. https://cloudlogging.app.goo.gl/FU6FtqvXr6pegYv57
• E2E on mlperf ckpt0424: runs to lm_loss 3.622 over 84 steps.
• Unit Tests intests/unit/grad_mask_utils_test.py.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @gagika, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @gagika, but I was unable to process your request. Please see the logs for more details.

[github-actions]

🤖 Hi @gagika, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This Pull Request introduces two stability-focused improvements for DeepSeek-V3 training: a per-token gradient mask and a constant-standard-deviation initialization for MLA projections. These changes are well-integrated into the existing configuration and layer structure, with comprehensive unit tests provided for the new gradient masking utility.

🔍 General Feedback

• Stability: The opt-in gradient mask is a robust defensive mechanism against gradient overflows in bf16, particularly useful for large-scale training.
• Precision: Using float32 for RMS calculation in the gradient mask is a good practice to maintain precision.
• Initialization: The specialized initialization for MLA projections correctly follows the DeepSeek-V3 architecture's requirements.
• Testing: The new unit tests cover edge cases and dtype preservation effectively.

[github-actions]

🤖 Hi @gagika, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @gagika, but I was unable to process your request. Please see the logs for more details.