Skip to content

Migrate the test script to use to_maxtext.py for supported models #4541

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
AusarYao28 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Migrate the test script to use to_maxtext.py for supported models #4541

Show file tree
Hide file tree
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
3 changes: 3 additions & 0 deletions tests/assets/golden_logits/golden_data_gpt-oss-20b.jsonl
View file
Open in desktop

@aireenmei aireenmei Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

We don't add golden_logits for new models to tests/assets/golden_logits any more, because it makes the repo size grow too fast. Now we have this instruction (https://github.com/AI-Hypercomputer/maxtext/pull/4541/changes#diff-af585533e794d79e01037bc09d8f4b879e8328b0a577f037579d6f2b6f45c4c4L57) in the test script. You can download the existing golden logits from gs://maxtext-test-assets/golden_data_gpt-oss-20b.jsonl and test against it.

Large diffs are not rendered by default.

22 changes: 20 additions & 2 deletions tests/end_to_end/tpu/deepseek/v2-16b/test_deepseek.sh
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,10 +41,28 @@ if [ -z "${CKPT_DISK_LOCATION}" ]; then
fi

# 1.1 Convert checkpoint to `scanned` format, more suitable for training
JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.standalone_scripts.convert_deepseek_family_ckpt --base_model_path ${CKPT_DISK_LOCATION} --maxtext_model_path ${BASE_OUTPUT_PATH}/scanned --model_size ${MODEL_NAME}
python3 -m maxtext.checkpoint_conversion.to_maxtext \
model_name=${MODEL_NAME} \
base_output_directory=${BASE_OUTPUT_PATH}/scanned \
scan_layers=true \
attention=dot_product \

@aireenmei aireenmei Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

The attention type should not affect checkpoint conversion. I suggest leave it the default value to reduce noise.

@AusarYao28 AusarYao28 Jul 20, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Without dot attention we have error: ```
File "/home/yaoyuchen_google_com/maxtext/maxtext_venv/lib/python3.12/site-packages/jax/experimental/pallas/ops/gpu/attention.py", line 243, in mha
raise ValueError(
ValueError: This kernel expects q, k, and v to have the same head dimension, but found q.shape=(192, 2048, 16, 192), k.shape=(192, 2048, 16, 192), v.shape=(192, 2048, 16, 128).


With dot attention we have a different error: ```
Traceback (most recent call last):                                                                                                                                                                                                                                                                              
  File "<frozen runpy>", line 198, in _run_module_as_main                                                                                                                                                                                                                                                       
  File "<frozen runpy>", line 88, in _run_code                                                                                                                                                                                                                                                                  
  File "/home/yaoyuchen_google_com/maxtext/src/maxtext/checkpoint_conversion/to_maxtext.py", line 1162, in <module>                                                                                                                                                                                             
    main(                                                                                                                                                                                                                                                                                                       
  File "/home/yaoyuchen_google_com/maxtext/src/maxtext/checkpoint_conversion/to_maxtext.py", line 1003, in main                                                                                                                                                                                                 
    filtered_map_keys = validate_and_filter_param_map_keys(param_map_mt_to_hf.keys(), maxtext_abstract_dict.keys())                                                                                                                                                                                             
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                                                             
  File "/home/yaoyuchen_google_com/maxtext/src/maxtext/checkpoint_conversion/utils/utils.py", line 118, in validate_and_filter_param_map_keys                                                                                                                                                                   
    raise ValueError(                                                                                                                                                                                                                                                                                           
ValueError: maxtext_state_dict must be a subset of flattened param_map                                                                                                                                                                                                                                          
param map

enable_nnx=false \

@aireenmei aireenmei Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Have you tested enable_nnx=true (default)? We eventually want the nnx path to work since it's the future and we are migrating to it.

hardware=cpu skip_jax_distributed_system=True \
--hf_model_path=${CKPT_DISK_LOCATION} \
--lazy_load_tensors=False \

@aireenmei aireenmei Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Have you tried the default value lazy_load_tensors=True? It minimizes RAM usage, we can always use it.

--eager_load_method='safetensors'

@aireenmei aireenmei Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

'safetensors' is the default value, no need to specify?

@aireenmei aireenmei Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Add these for Pathways compatibility:

checkpoint_storage_use_ocdbt=False
checkpoint_storage_use_zarr3=False


# 1.2 Convert checkpoint to `unscanned` format, more suitable for decoding
JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.standalone_scripts.convert_deepseek_family_unscanned_ckpt --base_model_path ${CKPT_DISK_LOCATION} --maxtext_model_path ${BASE_OUTPUT_PATH}/unscanned --model_size ${MODEL_NAME}
python3 -m maxtext.checkpoint_conversion.to_maxtext \
model_name=${MODEL_NAME} \
base_output_directory=${BASE_OUTPUT_PATH}/unscanned \
scan_layers=false \
attention=dot_product \
enable_nnx=false \
hardware=cpu skip_jax_distributed_system=True \
--hf_model_path=${CKPT_DISK_LOCATION} \
--lazy_load_tensors=False \
--eager_load_method='safetensors'

# Step 2:
# We define the checkpoint paths. This way it is easier to use these paths in the `train.py` and `decode.py` commands
Expand Down
20 changes: 18 additions & 2 deletions tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,10 +39,26 @@ if [ -z "${CKPT_DISK_LOCATION}" ]; then
fi

# 1.1 Convert checkpoint to `scanned` format, more suitable for training
JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.standalone_scripts.convert_gpt_oss_ckpt --base-model-path ${CKPT_DISK_LOCATION} --maxtext-model-path ${BASE_OUTPUT_PATH}/scanned --model-size ${MODEL_NAME}
python3 -m maxtext.checkpoint_conversion.to_maxtext \
model_name=${MODEL_NAME} \
base_output_directory=${BASE_OUTPUT_PATH}/scanned \
scan_layers=true \
attention=dot_product \
hardware=cpu skip_jax_distributed_system=True \
--hf_model_path=${CKPT_DISK_LOCATION} \
--lazy_load_tensors=False \
--eager_load_method='safetensors'

# 1.2 Convert checkpoint to `unscanned` format, more suitable for decoding
JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.standalone_scripts.convert_gpt_oss_unscanned_ckpt --base-model-path ${CKPT_DISK_LOCATION} --maxtext-model-path ${BASE_OUTPUT_PATH}/unscanned --model-size ${MODEL_NAME}
python3 -m maxtext.checkpoint_conversion.to_maxtext \
model_name=${MODEL_NAME} \
base_output_directory=${BASE_OUTPUT_PATH}/unscanned \
scan_layers=false \
attention=dot_product \
hardware=cpu skip_jax_distributed_system=True \
--hf_model_path=${CKPT_DISK_LOCATION} \
--lazy_load_tensors=False \
--eager_load_method='safetensors'

# Step 2:
# We define the checkpoint paths. This way it is easier to use these paths in the `train.py` and `decode.py` commands
Expand Down
Empty file modified tests/end_to_end/tpu/mixtral/8x7b/1_test_mixtral.sh
100644 → 100755
View file
Open in desktop

@aireenmei aireenmei Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I don't think permission change is needed? We usually run bash <script.sh> instead of directly ./script.sh

Empty file.
Empty file modified tests/end_to_end/tpu/mixtral/8x7b/2_test_mixtral.sh
100644 → 100755
View file
Open in desktop
Empty file.
Loading