AI-Hypercomputer / AI-Hypercomputer/maxtext

Multimodal SFT: Conversion errors to HF Safetensors

未關閉
#4,712 0 則留言 0 個 reaction 已指派 0 人 在 GitHub 檢視
bug
主要語言
Python
星號
2.4k
分支
607
平均合併
2 天 19 小時
30 天內合併 PR
158

描述

### Bug report

python3 -m maxtext.checkpoint_conversion.to_huggingface \
model_name=${MODEL_NAME?} \
load_parameters_path=${POST_TRAIN_PATH?} \
base_output_directory=${HF_EXPORT?} \
scan_layers=True \
use_multimodal=True \
weight_dtype=bfloat16

### Logs/Output

1. I get the following error, while `rngs` and `aqt` should be ignored: [full_error.log](https://github.com/user-attachments/files/30669717/full_error.log).

2. Once this error was fixed by manually stripping those, I got the following error:
```
I0803 16:00:42.236964 126572543616832 utils.py:920] Detected NNX-SFT checkpoint structure
I0803 16:00:42.238894 126572543616832 to_huggingface.py:321]
Proccessing weight...
0%| | 0/571 [00:00", line 198, in _run_module_as_main
File "", line 88, in _run_code
File "/home/dgouju_google_com/maxtext_venv/lib/python3.12/site-packages/maxtext/checkpoint_conversion/to_huggingface.py", line 368, in
app.run(main)
File "/home/dgouju_google_com/maxtext_venv/lib/python3.12/site-packages/absl/app.py", line 367, in run
_run_main(main, args)
File "/home/dgouju_google_com/maxtext_venv/lib/python3.12/site-packages/absl/app.py", line 312, in _run_main
sys.exit(main(argv))
^^^^^^^^^^
File "/home/dgouju_google_com/maxtext_venv/lib/python3.12/site-packages/maxtext/checkpoint_conversion/to_huggingface.py", line 333, in main
processed_params = process_maxtext_param(key, weight, param_map, hook_fn_map, shape_map, config)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/dgouju_google_com/maxtext_venv/lib/python3.12/site-packages/maxtext/checkpoint_conversion/utils/utils.py", line 244, in process_maxtext_param
_process(
File "/home/dgouju_google_com/maxtext_venv/lib/python3.12/site-packages/maxtext/checkpoint_conversion/utils/utils.py", line 183, in _process
numpy_slice = convert_jax_weight_to_numpy(processed_slice, save_dtype).reshape(target_hf_shape)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: 'float' object cannot be interpreted as an integer
```
Sounds like standard/multimodal layer shape translations (the shape_map) where dimensions are derived via standard division `/` rather than floor/integer division `//`, or when floating-point configuration parameters are propagated through the MaxText model configs.

A hotfix was to patch `numpy_slice` creation in `_process()` in `maxtext/checkpoint_conversion/utils/utils.py` with:
```
int_target_hf_shape = tuple(int(dim) for dim in target_hf_shape) if isinstance(target_hf_shape, (tuple, list)) else int(target_hf_shape)
numpy_slice = convert_jax_weight_to_numpy(processed_slice, save_dtype).reshape(int_target_hf_shape)
```

### Environment Information

Installed with `uv pip install maxtext[tpu-post-train]==0.2.2 --resolution=lowest`

### Additional Context

_No response_

貢獻指南

開啟貢獻指南

評估

這個 Issue 還沒有評估資料。

把新 issue 寄到你的電子郵件信箱

精選適合新手參與的 GitHub issue 摘要。