AI-Hypercomputer / AI-Hypercomputer/maxtext

Multimodal SFT: Conversion errors to HF Safetensors

Đang mở
#4,712 0 bình luận 0 reaction 0 người được giao Xem trên GitHub
bug
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

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.