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.

The rest of the description includes relevant details and context, examples:

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 :

=== DeepSeek V4 Load Balancing Variance Analysis (Step 0 vs Step 20) ===

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

• 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

❌ Patch coverage is 72.00000% with 21 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/trainers/pre_train/train.py	59.57%	18 Missing and 1 partial ⚠️
...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!