Fix DS v3.2 decoding issue

[RissyRan]

Description

Fix DeepSeek V3.2 model decoding issue, see more details in b/500088962.

• We should enable the Dense warm up only when indexer_loss_scaling_factor > 0 like mentioned here.
• We should use model_mode instead of self.model_mode. Updated other branches as well.

Tests

• Before the change: random outputs
• After the change: reasonable decoding outputs

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!

[shuningjin]

Thanks for the investigation and fix! Could you also add a comment on the model mode for future reference?

maxtext/src/maxtext/layers/decoders.py

Line 1117 in b83ff02

elif (cfg.use_indexer and not cfg.indexer_sparse_training) and self.model_mode == MODEL_MODE_TRAIN:

[RissyRan]

Thanks for the investigation and fix! Could you also add a comment on the model mode for future reference?

maxtext/src/maxtext/layers/decoders.py

Line 1117 in b83ff02

elif (cfg.use_indexer and not cfg.indexer_sparse_training) and self.model_mode == MODEL_MODE_TRAIN:

Added.

[Rohan-Bierneni]

LGTM! Thanks for the fix.