Add llama3.1-70b E2E test scripts by chiajunglien · Pull Request #4279 · AI-Hypercomputer/maxtext

chiajunglien · 2026-06-26T06:16:34Z

Description

This PR introduces end-to-end (E2E) testing and validation pipelines for the Llama3.1 70B model in MaxText.

Tests (Testing on v5p-128 cluster)

Pre-Training:

export RUN_ID=$(date +%Y-%m-%d-%H-%M)
bash tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_to_mt.sh $RUN_ID
bash tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b.sh $RUN_ID

SFT:

export RUN_ID=$(date +%Y-%m-%d-%H-%M)
bash tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_to_mt.sh $RUN_ID
bash tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_sft.sh $RUN_ID

RL:

export RUN_ID=$(date +%Y-%m-%d-%H-%M)
bash tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_to_mt.sh $RUN_ID
bash tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh $RUN_ID

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 · 2026-07-02T09:45:34Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

….1-70b Introduces E2E test configurations and scripts for llama3.1-70b model, covering both pre-training and post-training validation pipelines.

YixuanWang-99 · 2026-07-13T17:42:05Z

+    load_parameters_path=${CKPT_PATH} \
+    base_output_directory=${BASE_OUTPUT_DIRECTORY}/to_huggingface/${scan_status}/${run_id} \
+    use_multimodal=${USE_MULTIMODAL} \
+    scan_layers=$SCAN_LAYERS


should we do forward_pass_logit_checker as well?

I'm not quite sure, as the existing tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_hf.sh doesn't have this check. I followed its pattern for this PR.

YixuanWang-99 · 2026-07-13T17:43:23Z

+# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S)
+# bash test_llama3.1_70b_to_mt.sh $RUN_ID
+# bash test_llama3.1_70b_rl.sh $RUN_ID
+


$USE_PATHWAY is the second arguments?

Yes, the second argument is for Pathways. But we don't need to manually pass it here since the DAG will automatically handle it and configure for us.

YixuanWang-99 · 2026-07-13T17:44:10Z

+set -ex
+
+run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+use_pathways=${2:-false}


Should use_pathways to be true because we are testing with v5p cluster?

We do need Pathways for multi-host VMs.
The reason we kept the default as false here is just for easier single-host/local testing. When running on multi-host through the DAG, the DAG will automatically pass true as the second argument to handle the Pathways setup for us.

YixuanWang-99 · 2026-07-13T17:44:44Z

+set -ex
+
+run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+use_pathways=${2:-false}


Should use_pathways to be true if tested on multi-host vms?

We do need Pathways for multi-host VMs.
The reason we kept the default as false here is just for easier single-host/local testing. When running on multi-host through the DAG, the DAG will automatically pass true as the second argument to handle the Pathways setup for us.

YixuanWang-99 · 2026-07-13T17:44:55Z

+# export HF_TOKEN=<your Hugging Face access token>
+# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S)
+# bash test_llama3.1_70b_to_mt.sh $RUN_ID
+# bash test_llama3.1_70b_sft.sh $RUN_ID


missing argument

Yes, the second argument is for Pathways. But we don't need to manually pass it here since the DAG will automatically handle it and configure for us.

This was referenced Jun 26, 2026

feat: add Gemma 4 26B E2E pre-training pipeline CIeNET-International/ml-auto-solutions-3#286

Closed

feat: Add Llama3.1-70b and Gemma4-26b E2E training pipeline CIeNET-International/ml-auto-solutions-3#287

Merged

RexBearIU reviewed Jun 26, 2026

View reviewed changes

Comment thread tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_to_mt.sh Outdated

RexBearIU reviewed Jun 26, 2026

View reviewed changes

Comment thread tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_to_mt.sh

chiajunglien changed the title ~~Add llama3.1-70b E2E test scripts for pre-training~~ Add llama3.1-70b E2E test scripts Jun 30, 2026

chiajunglien force-pushed the emma/e2e-training-llama3-70b branch from 98951f4 to 20a8cbb Compare July 1, 2026 08:42

chiajunglien marked this pull request as ready for review July 2, 2026 08:06

chiajunglien force-pushed the emma/e2e-training-llama3-70b branch from e8ba9a7 to 04a50c4 Compare July 2, 2026 09:38

chiajunglien force-pushed the emma/e2e-training-llama3-70b branch 5 times, most recently from c3cb1b2 to e4738b7 Compare July 9, 2026 09:23

feat(testing): add end-to-end training and inference tests for llama3…

6bf8e6f

….1-70b Introduces E2E test configurations and scripts for llama3.1-70b model, covering both pre-training and post-training validation pipelines.

chiajunglien force-pushed the emma/e2e-training-llama3-70b branch from 0f3682c to 6bf8e6f Compare July 10, 2026 07:17

SurbhiJainUSC approved these changes Jul 10, 2026

View reviewed changes

YixuanWang-99 reviewed Jul 13, 2026

View reviewed changes

YixuanWang-99 approved these changes Jul 16, 2026

View reviewed changes

hsuan-lun-chiang added the pull ready label Jul 17, 2026

copybara-service Bot merged commit 4b357c1 into AI-Hypercomputer:main Jul 17, 2026
160 of 162 checks passed

Uh oh!

Conversation

chiajunglien commented Jun 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests (Testing on v5p-128 cluster)

Checklist

Uh oh!

Uh oh!

Uh oh!

codecov Bot commented Jul 2, 2026

Codecov Report

Uh oh!

YixuanWang-99 Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

chiajunglien Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

YixuanWang-99 Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

chiajunglien Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

YixuanWang-99 Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

chiajunglien Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

YixuanWang-99 Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

chiajunglien Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

YixuanWang-99 Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

chiajunglien Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

chiajunglien commented Jun 26, 2026 •

edited

Loading

chiajunglien Jul 14, 2026 •

edited

Loading

chiajunglien Jul 14, 2026 •

edited

Loading

chiajunglien Jul 14, 2026 •

edited

Loading

chiajunglien Jul 14, 2026 •

edited

Loading