Disallow use_indexer config combinations that zero the entire training objective

[AusarYao28]

Description

With use_indexer=True and indexer_sparse_training=False, train.py runs the indexer dense warm-up:
loss_fn sets xent_sum = 0.0 and the only signal is the indexer KL loss scaled by
indexer_loss_scaling_factor. base.yml defaults those to false / 0.0, so a model config that enabled
use_indexer without overriding either flag (deepseek4-284b, deepseek4-tiny) trained on total_loss == 0.0
without any error.

• validate_train_config rejects use_indexer=True with indexer_sparse_training=False and
  indexer_loss_scaling_factor <= 0.0.
• deepseek4-284b.yml and deepseek4-tiny.yml default indexer_sparse_training: true, so a plain model_name
  run trains on the LM loss. indexer_loss_scaling_factor stays 0.0: with a positive default, the
  types.py check indexer_topk < max_target_length // 4 (config-level, not train-only) rejects
  deepseek4-284b (indexer_topk 512) at base.yml's max_target_length 2048 and at the decode stage's 512.
• Tests added to tests/unit/train_utils_test.py::TestValidateTrainConfig.

FIXES: b/563046810

Tests

• tests/unit/train_utils_test.py::TestValidateTrainConfig: 13 passed.
• train_utils_test + attention_compressed_test + csa_streamindex_test + deepseek_v4_rope_test +
  deepseek_v4_indexer_loss_test: 57 passed, 2 skipped.
• pyconfig + validate_train_config with plain model_name=deepseek4-284b / deepseek4-tiny: pass with this
  change; without the yml default, deepseek4-284b raises "zeroes the entire training objective".
• E2E (v5p-128, pre-review code): https://cloudlogging.app.goo.gl/btd35GVEi93uhmUq5

Checklist

• 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.

[gemini-code-assist]

Code Review

This pull request introduces validation logic to prevent training configurations where 'use_indexer' is enabled without 'indexer_sparse_training' and with a non-positive 'indexer_loss_scaling_factor', which would otherwise result in a zeroed training objective. The validation is implemented in 'validate_train_config' and supported by a new unit test. The reviewer correctly identified that the redundant check added to 'loss_fn' in 'train.py' should be removed to maintain clean code, avoid duplication, and prevent potential issues with raising exceptions inside JIT-compiled functions.

[gemini-code-assist]

This configuration check is redundant because the exact same validation is already performed at startup in validate_train_config (inside src/maxtext/utils/train_utils.py). Centralizing configuration validation in validate_train_config is the standard pattern in MaxText, keeping the loss function clean and avoiding duplication. Additionally, raising Python exceptions inside JIT-compiled functions (like loss_fn during tracing) is best avoided for static configuration parameters. We should remove this check from loss_fn.

    xent_sum = 0.0
    total_z_loss = 0.0

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[Rohan-Bierneni]

Add license header here

[AusarYao28]

File deleted; the tests now live in tests/unit/train_utils_test.py, which already has the header.

[Rohan-Bierneni]

These tests can be moved into tests/unit/train_utils_test.py::TestValidateTrainConfig

[AusarYao28]

The tests moved into train_utils_test.py::TestValidateTrainConfig, and the standalone file was deleted, so no new file needs a header.

[Rohan-Bierneni]

deepseek4-284b.yml sets use_indexer: true but inherits indexer_sparse_training=false / indexer_loss_scaling_factor=0.0 from base.yml, so a plain model_name=deepseek4-284b run now fails validation.

Note that main's 2_test_deepseek.sh pre-train stage uses exactly those defaults, so this needs to land with/after #5283 or add indexer_sparse_training=true to that script here. Optionally, consider defaulting one of these in deepseek4-284b.yml so the model config is self-consistent.

Suggest defaulting indexer_sparse_training: true / indexer_loss_scaling_factor: 1.0 in deepseek4-284b.yml (and -tiny)

[AusarYao28]

Confirmed and fixed in this PR, so it no longer depends on #5283 landing first. Both DS4 ymls now set
indexer_sparse_training: true; main's 2_test_deepseek.sh pretrain stage (no indexer flags,
max_target_length=2048) passes validation with it.

I kept indexer_loss_scaling_factor at 0.0 rather than 1.0: the compressed-attention check in types.py
(indexer_topk >= max_target_length // 4 with scaling > 0 raises) runs for every mode, so a 1.0 default
rejects plain model_name=deepseek4-284b at base.yml's max_target_length 2048 (512 >= 512) and the decode
stage at 512 (512 >= 128). The e2e pretrain stage passes indexer_loss_scaling_factor=1.0 explicitly at
max_target_length=4096. If you prefer 1.0 as the default, the options are adding
indexer_loss_scaling_factor=0.0 to every inference invocation or scoping that check to training.

[shuningjin]

Thanks for enhancing the robustness! Approve to unblock. Please address the comments from other reviewers.

if keeping the current approach: setting indexer_sparse_training: true in deepseek4-*.yml breaks the documented DSv4 warm-up command. The CLI passes indexer_sparse_training=false without override_model_config=true, so pyconfig rejects it as a model-config/CLI conflict. Please add override_model_config=true to the warm-up command in Run_DeepSeek.md#L310-L328. Also note that deepseek3.2-671b and deepseek-custom set use_indexer: true with the defaults, so plain train.py runs on them will now fail the new validator.

alternative: fix the root cause in train.py instead (see inline comment). That makes the YAML edits, the validator, and the doc change unnecessary.

[shuningjin]

Fix the root cause instead of adding a validator.

The dense warm-up gate in train.py ignores indexer_loss_scaling_factor, so it disagrees with the logits-skip gate and with the config docs:

Location	Condition
train.py#L222-L226 (zeroes LM loss)	use_indexer and not sparse
nnx_decoders.py#L2308-L2315 (skips logits)	use_indexer and scale > 0 and not sparse
base.yml#L556-L562 (docs)	indexer_sparse_training "only active when indexer_loss_scaling_factor > 0"

With scale=0 and sparse=false (the defaults), the model computes full logits and then train.py discards them: wasted compute and zero loss.

Suggested change at train.py#L222:

if (config.use_indexer and config.indexer_loss_scaling_factor > 0.0
    and not config.indexer_sparse_training) and is_train:

This aligns all three conditions. scale=0 then means normal LM training with the indexer used for selection only, which covers continued training of converted V3.2/V4 checkpoints. With this change:

• the YAML edits and the new validator in train_utils.py should be reverted;
• deepseek3.2-671b / deepseek-custom need no YAML change;
• the DSv4 warm-up command in Run_DeepSeek.md works unchanged.

[AusarYao28]

Done, I went with the root-cause fix. train.py now gates the warm-up branch on indexer_loss_scaling_factor > 0.0, the same way nnx_decoders.py and attention_compressed.py already do. I reverted the validator and both deepseek4 yml edits. deepseek3.2-671b and deepseek-custom work unchanged, and so does the documented DSv4 warm-up command. I updated the two warm-up tests in train_nnx_test.py to set a positive scale, and added a test that scale 0 keeps the LM loss.

[AusarYao28]

Done, I went with the root-cause fix. train.py now gates the warm-up branch on indexer_loss_scaling_factor > 0.0, the same way nnx_decoders.py and attention_compressed.py already do. I reverted the validator and both deepseek4 yml edits. deepseek3.2-671b and deepseek-custom work unchanged, and so does the documented DSv4 warm-up command. I updated the two warm-up tests in train_nnx_test.py to set a positive scale, and added a test that scale 0 keeps the LM loss.

[AusarYao28]

Thanks for enhancing the robustness! Approve to unblock. Please address the comments from other reviewers.

if keeping the current approach: setting indexer_sparse_training: true in deepseek4-*.yml breaks the documented DSv4 warm-up command. The CLI passes indexer_sparse_training=false without override_model_config=true, so pyconfig rejects it as a model-config/CLI conflict. Please add override_model_config=true to the warm-up command in Run_DeepSeek.md#L310-L328. Also note that deepseek3.2-671b and deepseek-custom set use_indexer: true with the defaults, so plain train.py runs on them will now fail the new validator.

alternative: fix the root cause in train.py instead (see inline comment). That makes the YAML edits, the validator, and the doc change unnecessary.

Both points are moot now: the yml defaults are reverted, so the warm-up command in Run_DeepSeek.md no longer conflicts, and deepseek3.2-671b / deepseek-custom need no change.