Conversion scripts fail on standard TPU VMs for large models

[zzxslp]

Bug report

The official instruction 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

[hengtaoguo]

Thanks for reporting! We are still optimizing this conversion script and trying to lower the memory peak. Can you try add a --lazy_load_tensors=True and see if it works (reference)? It loads Hugging Face weights on-demand to minimize RAM usage. With this 235B-sized model, the script was meant to run with an even larger CPU machine.

[zzxslp]

Hi, I've set -lazy_load_tensors=True but it didn't help, didn't dig deeper, maybe it is working?

Temporarily I'm using a CPU pod to bypass this issue, here are the memory breakdown to convert qwen3-235B with simulated_cpu_devices_count=16:

INFO:absl:saved a checkpoint at step 0
INFO:absl:[process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
INFO:absl:Elapse for checkpoint save: 1.24 min
INFO:absl:[Program Ends] RAM Usage: 464.04/1891.84 GB (24.5%)
INFO:absl:Conversion complete. Checkpoint saved to /dev/shm/qwen3-235b-a22b-jax
INFO:absl:Overall Elapse: 59.37 min
INFO:absl:Peak Memory: 585.08 GB

[RissyRan]

Hi @hengtaoguo, I see you have started the conversation, I will assign this to you first. Feel free to triage. Thanks!