Fix zero loss bug in DSV3.2 by adding indexer flags

[Rohan-Bierneni]

Description

Update the deepseek v3.2 script used by XLML to fix training loss set to 0

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/519029446

Tests

Verified on mini-config on synthetic data. Details in https://b.corp.google.com/issues/519029446#comment4.

PR Unit tests should pass.

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.

[gemini-code-assist]

Code Review

This pull request updates the pre-training command in 2_test_deepseek.sh to include indexer-related configurations (use_indexer=True, indexer_loss_scaling_factor=0.01, and indexer_sparse_training=True). Feedback suggests wrapping this extremely long command line to comply with the 80-character line limit of the Google Shell Style Guide, which will improve readability and maintainability.

[gemini-code-assist]

The command line is extremely long (~450 characters), which violates the standard 80-character line limit of the Google Shell Style Guide and makes the script difficult to read and maintain. Wrapping the command using backslashes (\) improves readability and makes it easier to see the added flags.

python3 -m maxtext.trainers.pre_train.train \
  ${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}/base.yml \
  base_output_directory=${BASE_OUTPUT_PATH} \
  run_name=tokamax_gmm_pre_training \
  model_name=${MODEL_NAME} \
  tokenizer_type=huggingface \
  tokenizer_path=${TOKENIZER_PATH} \
  dataset_type=synthetic \
  enable_checkpointing=false \
  attention=flash \
  use_tokamax_splash=True \
  sparse_matmul=True \
  use_tokamax_gmm=True \
  dtype=bfloat16 \
  weight_dtype=bfloat16 \
  per_device_batch_size=1 \
  steps=5 \
  max_target_length=4096 \
  ici_fsdp_parallelism=-1 \
  opt_type=sgd \
  use_indexer=True \
  indexer_loss_scaling_factor=0.01 \
  indexer_sparse_training=True

References

1. Google Shell Style Guide recommends a maximum line length of 80 characters. Long lines should be folded using a backslash where appropriate to improve readability. ^(link)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[bvandermoon]

LGTM, thanks @Rohan-Bierneni. Is this the only test/script impacted by this 0-loss issue?