Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Rohan-Bierneni wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion tests/end_to_end/tpu/deepseek/v3.2-671b/2_test_deepseek.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ python3 -m tests.utils.forward_pass_logit_checker ${MAXTEXT_CONFIGS_DIR:-${MAXTE

# Run pre-training - tokamax_gmm implementation
# Note: use sgd due to memory constraint
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
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

@gemini-code-assist gemini-code-assist Bot Aug 21, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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)


# Run decoding - megablox implementation
# Note: decode requires the access token for huggingface tokenizer even if the model is not gated
Expand Down
Loading