test script for deepseek

[shuningjin]

Description

Main change
onboard test for deepseek2-16b, plan to run full workflow in XLML on v5p-8

• end_to_end/tpu/deepseek/v2-16b/test_deepseek.sh: ckpt conversion, logit check, train, finetune, decode

onboard test for deepseek3-671b, plan to run step 2 in XLML on v5p-128

• end_to_end/tpu/deepseek/v3-671b/1_test_deepseek.sh
• end_to_end/tpu/deepseek/v3-671b/2_test_deepseek.sh: logit check (scan), decode (unscan). Note we leave out training and finetuning, as is covered by ubench nightly test.

XLML PR: GoogleCloudPlatform/ml-auto-solutions#1112

FIXES: b/423057893

Other change

• Minor format change to other test scripts.
• max_logging.log behavior is changed due to PR#2873, need to set absl.logging.level to make logs visible for multiple scripts

Tests

preparation

• The checkpoints uploaded
  • gs://maxtext-deepseek/deepseek2-16b (hf-bf16)
  • gs://maxtext-deepseek/deepseek3-671b (hf-bf16, scanned, unscanned)
• golden logits uploaded
  • gs://maxtext-test-assets/golden_data_deepseek2-16b.jsonl
  • gs://maxtext-test-assets/golden_data_deepseek3-671b.jsonl

bash end_to_end/tpu/deepseek/v3-671b/2_test_deepseek.sh

• xpk on v5p-128: https://paste.googleplex.com/5746766743994368
• log: https://cloudlogging.app.goo.gl/38jsiXx6y3miv7TL7

bash end_to_end/tpu/deepseek/v2-16b/test_deepseek.sh

• xpk on v5p-8: https://paste.googleplex.com/6732499990478848
• log, 30min: https://cloudlogging.app.goo.gl/J7SA4eLjMvjNyNGm8

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!

[github-actions]

🤖 Hi @shuningjin, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

📋 Review Summary

This Pull Request introduces new test scripts for DeepSeek models (v2-16b and v3-671b) and refactors existing test scripts for GPT-OSS, Llama4, and Mixtral. The changes aim to onboard DeepSeek models, improve consistency in script practices, and adjust logging verbosity.

🔍 General Feedback

• The refactoring of test scripts into modular 1_test_deepseek.sh and 2_test_deepseek.sh for DeepSeek v3 is a positive step for organization and clarity.
• Consistent application of absl.logging.set_verbosity(absl.logging.INFO) across relevant Python files is good for debugging and monitoring.
• The general updates to how BASE_OUTPUT_PATH is handled in the shell scripts improve robustness.
• Consider reviewing the commented-out training and fine-tuning steps in the DeepSeek test scripts to ensure they are intentionally disabled and to add explanatory comments if needed.
• Ensure all golden logit paths are robust and not dependent on transient local paths or specific dates.

[RissyRan]

Great work! Also thanks for keeping them up-to-date! A few minor comments

[RissyRan]

Do we need to set this from XL ML? or we cold direct export MODEL_VARIATION="llama4-17b-16e" without if/else

[shuningjin]

For other models, we have separate test scripts for each variant. For llama4, this serves as the united script. Do we need to account for Maverick as well?

[RissyRan]

I see. I think we should do something similar to GPT-OSS models. By passing the type instead of duplicate the scripts. No need to be in this PR.

[shuningjin]

Different variant may need customization in tests, for instance

• forward logit check criteria
• parallelism as they are typically run on different devices (e.g., v5p-8 vs. v5p-128)

I would prefer keeping disentangled scripts for different variants when possible.

Actually, for the llama4 case, the logit criteria is set for scout. this might not work for maverick. Need to follow this up.

[RissyRan]

Thanks for the updates!

[RissyRan]

I see. I think we should do something similar to GPT-OSS models. By passing the type instead of duplicate the scripts. No need to be in this PR.

[shuningjin]

For deepseek script, updated moe strategy to tokamax_gmm for training and decoding. Tests pass:

• ds3: https://cloudlogging.app.goo.gl/onW4CPgVkZ5SZhBVA
• ds2: https://cloudlogging.app.goo.gl/z8u5cMrpykniX7ji6

[parambole]

LGTM. Thank you for making these changes. I have left a nit.