Enabled auxilillary loss free load balancing and sequence wise load b…

[dipakg-lang]

##Description
Enabled auxilillary loss free load balancing and sequence wise load balancing for Deepseek.

We require these changes to be compliant with the specification and also this will help us to have better use to hardware due to the load balancing between the experts.

FIXES: b/509933890
FIXES: b/521990776

##Tests
Tested by running training loop with new tiny Deeepseek V4 model added as part of the commit,
here are the logs for testing and commands used for this :

export JAX_PLATFORM_NAME=cpu & export XLA_FLAGS=--xla_force_host_platform_device_count=8 & python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml override_model_config=True model_name=deepseek4-tiny enable_checkpointing=False base_output_directory=/tmp/maxtext_output/ dataset_type=synthetic hardware=cpu skip_jax_distributed_system=True attention=dot_product per_device_batch_size=1 steps=100 max_target_length=256 async_checkpointing=false dtype=bfloat16 weight_dtype=bfloat16 megablox=False sparse_matmul=False ici_expert_parallelism=-1 sharding_tolerance=1.0 ici_fsdp_parallelism=1 indexer_topk=16 routed_bias=True load_balance_loss_weight=0.0001 routed_bias_update_rate=0.001

Without load balancing active logs : https://paste.googleplex.com/6421399878107136

with load balancing logs : https://paste.googleplex.com/6551357300539392

Here are the results actived for reducing the varience :

1 === DeepSeek V4 Load Balancing Variance Analysis (Step 0 vs Step 20) ===
2
3 | Layer Index | Routing Type | Step 0 Var (Baseline) | Step 20 Var (Run A) | Step 20 Var (Run B) | Improvement (A vs B) |
4 |-------------|--------------|-----------------------|---------------------|---------------------|----------------------|
5 | 0 | Hash Routed | 3932160.00 | 3932160.00 | 3932160.00 | 0.00% |
6 | 1 | Hash Routed | 3932160.00 | 3932160.00 | 3932160.00 | 0.00% |
7 | 2 | Hash Routed | 3932160.00 | 3932160.00 | 3932160.00 | 0.00% |
8 | 3 | Top-K Routed | 7409.38 | 7509.25 | 3672.12 | 51.10% |
9 | 4 | Top-K Routed | 3158.38 | 3230.12 | 1216.00 | 62.35% |
10 | 5 | Top-K Routed | 5713.38 | 5772.75 | 2359.38 | 59.13% |
11 | 6 | Top-K Routed | 8295.25 | 8082.50 | 3674.12 | 54.54% |
12 | 7 | Top-K Routed | 4765.62 | 4614.62 | 1212.75 | 73.72% |
13 | 8 | Top-K Routed | 4960.75 | 4923.12 | 1663.50 | 66.21% |
14 | 9 | Top-K Routed | 3905.50 | 3816.25 | 1316.88 | 65.49% |
15 | 10 | Top-K Routed | 5057.00 | 4981.12 | 2257.75 | 54.67% |
16 | 11 | Top-K Routed | 10446.62 | 10381.62 | 5565.75 | 46.39% |
17 | 12 | Top-K Routed | 9538.50 | 9529.25 | 5319.12 | 44.18% |
18 | 13 | Top-K Routed | 7031.38 | 7131.25 | 3270.25 | 54.14% |
19 | 14 | Top-K Routed | 4852.00 | 4900.12 | 1906.88 | 61.09% |
20 | 15 | Top-K Routed | 9306.12 | 9342.88 | 4733.75 | 49.33% |
21 | 16 | Top-K Routed | 5811.25 | 5749.50 | 2110.88 | 63.29% |
22 | 17 | Top-K Routed | 6715.62 | 6874.25 | 2664.12 | 61.24% |
23 | 18 | Top-K Routed | 8145.50 | 7869.25 | 3383.75 | 57.00% |
24 | 19 | Top-K Routed | 6042.12 | 5908.62 | 2353.00 | 60.18% |
25 | 20 | Top-K Routed | 8559.88 | 8158.25 | 4333.38 | 46.88% |
26 | 21 | Top-K Routed | 11742.25 | 11943.62 | 7563.50 | 36.67% |
27 | 22 | Top-K Routed | 4959.62 | 5014.88 | 1998.62 | 60.15% |
28 | 23 | Top-K Routed | 7717.12 | 7751.88 | 3879.88 | 49.95% |
29 | 24 | Top-K Routed | 9017.75 | 9307.88 | 4702.75 | 49.48% |
30 | 25 | Top-K Routed | 14127.12 | 14111.25 | 8079.25 | 42.75% |
31 | 26 | Top-K Routed | 5074.25 | 5194.12 | 1675.50 | 67.74% |
32 | 27 | Top-K Routed | 11919.50 | 11204.38 | 6470.75 | 42.25% |
33 | 28 | Top-K Routed | 12241.75 | 12998.62 | 7624.12 | 41.35% |
34 | 29 | Top-K Routed | 9384.50 | 9005.00 | 5052.00 | 43.90% |
35 | 30 | Top-K Routed | 9698.62 | 9678.25 | 5231.75 | 45.94% |
36 | 31 | Top-K Routed | 12244.25 | 12392.75 | 7249.25 | 41.50% |
37 | 32 | Top-K Routed | 10030.00 | 9972.62 | 4755.50 | 52.31% |
38 | 33 | Top-K Routed | 7265.00 | 6973.62 | 3271.75 | 53.08% |
39 | 34 | Top-K Routed | 11945.50 | 11940.62 | 6076.88 | 49.11% |
40 | 35 | Top-K Routed | 12917.50 | 13740.00 | 7210.62 | 47.52% |
41 | 36 | Top-K Routed | 15011.62 | 15083.00 | 8870.62 | 41.19% |
42 | 37 | Top-K Routed | 10294.12 | 10176.25 | 5907.50 | 41.95% |
43 | 38 | Top-K Routed | 8928.62 | 9236.00 | 5136.62 | 44.38% |
44 | 39 | Top-K Routed | 15633.62 | 15171.00 | 9684.75 | 36.16% |
45 | 40 | Top-K Routed | 7687.75 | 7658.12 | 4521.25 | 40.96% |
46 | 41 | Top-K Routed | 12485.12 | 12270.38 | 6933.25 | 43.50% |
47 | 42 | Top-K Routed | 17641.25 | 17163.50 | 10974.12 | 36.06% |
48 |-------------|--------------|-----------------------|---------------------|---------------------|----------------------|
49 | TOTAL/AVG | Top-K Only | 357681.12 | 356762.50 | 185883.62 | 47.90% |

Raw data collected for this analysis:

https://paste.googleplex.com/5060754624610304

https://paste.googleplex.com/5473518849490944

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

[x] I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
[x] I have necessary comments in my code, particularly in hard-to-understand areas.
[x] I have run end-to-end tests tests and provided workload links above if applicable.
[x] 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

❌ Patch coverage is 69.33333% with 23 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/trainers/pre_train/train.py	55.31%	19 Missing and 2 partials ⚠️
...axtext/checkpoint_conversion/inspect_checkpoint.py	0.00%	1 Missing ⚠️
src/maxtext/common/metric_logger.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!