Migrate the test script to use to_maxtext.py for supported models#4541
Migrate the test script to use to_maxtext.py for supported models#4541AusarYao28 wants to merge 1 commit into
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Migrated test script for Deepseek2-16b, gpt-oss-20b, and Mixtral8x7b
080dab1 to
ba340e3
Compare
aireenmei
left a comment
aireenmei
left a comment
There was a problem hiding this comment.
Could you share the test results, even it's the same error unrelated to this migration? We may want to create a bug further investigate if not already.
| base_output_directory=${BASE_OUTPUT_PATH}/scanned \ | ||
| scan_layers=true \ | ||
| attention=dot_product \ | ||
| enable_nnx=false \ |
There was a problem hiding this comment.
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.
| model_name=${MODEL_NAME} \ | ||
| base_output_directory=${BASE_OUTPUT_PATH}/scanned \ | ||
| scan_layers=true \ | ||
| attention=dot_product \ |
There was a problem hiding this comment.
The attention type should not affect checkpoint conversion. I suggest leave it the default value to reduce noise.
There was a problem hiding this comment.
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
| hardware=cpu skip_jax_distributed_system=True \ | ||
| --hf_model_path=${CKPT_DISK_LOCATION} \ | ||
| --lazy_load_tensors=False \ | ||
| --eager_load_method='safetensors' |
There was a problem hiding this comment.
'safetensors' is the default value, no need to specify?
| enable_nnx=false \ | ||
| hardware=cpu skip_jax_distributed_system=True \ | ||
| --hf_model_path=${CKPT_DISK_LOCATION} \ | ||
| --lazy_load_tensors=False \ |
There was a problem hiding this comment.
Have you tried the default value lazy_load_tensors=True? It minimizes RAM usage, we can always use it.
| hardware=cpu skip_jax_distributed_system=True \ | ||
| --hf_model_path=${CKPT_DISK_LOCATION} \ | ||
| --lazy_load_tensors=False \ | ||
| --eager_load_method='safetensors' |
There was a problem hiding this comment.
Add these for Pathways compatibility:
checkpoint_storage_use_ocdbt=False
checkpoint_storage_use_zarr3=False
There was a problem hiding this comment.
I don't think permission change is needed? We usually run bash <script.sh> instead of directly ./script.sh
There was a problem hiding this comment.
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.
2 participants
Description
Migrated test script for Deepseek2-16b, gpt-oss-20b, and Mixtral8x7b
FIXES: b/534808841
Tests
I've run each test script and compared its output before and after change, deepseekv2 test script fails with the same error message before and after change (checkpoint conversion script passed but forward pass failed)
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.