AI-Hypercomputer / AI-Hypercomputer/maxtext
Multimodal SFT: Conversion errors to HF Safetensors
- Ngôn ngữ chính
- Python
- Star
- 2.4k
- Fork
- 607
- Merge trung bình
- 2 ngày 19 giờ
- Pull request đã merge (30 ngày)
- 158
Mô tả
### 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_
Hướng dẫn đóng góp
Đánh giá
Issue này chưa được đánh giá.