AI-Hypercomputer / AI-Hypercomputer/maxtext

Conversion scripts fail on standard TPU VMs for large models

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

描述

### Bug report

The [official instruction]([https://maxtext.readthedocs.io/en/latest/guides/checkpointing_solutions/convert_checkpoint.html#hugging-face-to-maxtext) for model conversion failed on a standard TPU-v5p VM for large models such as QWen3-235B, with CPU OOM errors when sharding MoE layers. Since on GCP the CPU memory is fixed (400GB), I wonder if we can improve the script to bypass this issue, or is the doc outdated?

Also for the model conversion the sharding process with default simulated_cpu_devices_count=16 is very very slow (even for 30B model).

### Logs/Output

_No response_

### Environment Information

_No response_

### Additional Context

_No response_

貢獻指南

開啟貢獻指南

評估

這個 Issue 還沒有評估資料。

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

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