Fix BFloat16 conversion error in eager checkpoint loading

[phu0ngng]

Description

to_maxtext.py crashes with TypeError: Got unsupported ScalarType BFloat16 when converting bf16 HuggingFace models (e.g. Qwen3, Llama 3) in eager mode. PyTorch's .numpy() doesn't support bfloat16 tensors. This fix casts bf16 tensors to float32 before the numpy conversion.
The lazy loading path is unaffected (safetensors handles this internally).

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[shuningjin]

Thanks for the fix!

I believe this error is dependent on transformer version.

• Using older transformer version, from_pretrained by default loads bf16 ckpt in fp32. so it can pass with v.numpy()
• using new transformer version, it loads bf16 ckpt as is, and hence error with v.numpy(). Yes, your fix v.float().numpy() will help.

Lastly, we are working more on loading & type conversion for to_maxtext in #3184

[shuningjin]

move this import to the top?

[phu0ngng]

We can, but then it will require torch in the lazy loading path as well, even though torch is not needed there. I think we should go with the current implementation to avoid torch requirements for lazy load.

[phu0ngng]

@shuningjin what do you think?

[phu0ngng]

@shuningjin, friendly reminder.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[shuningjin]

Thanks for your patience! We re-vamped the loading/conversion/save pipeline for to_maxtext in PR 3184, with careful consideration of speed and memory.

The current hf -> maxtext behavior would be:

• eager loading (auto dtype as specified in config.json) -> conversion -> save with specified dtype (e.g., bf16 or f32, with bf16 as the recommended type to save memory).
• A typical conversion path: torch.bfloat16 -> torch.float32 -> numpy.float32 -> ml_dtypes.bfloat16.

Could you verify if the latest code solve your original issue? Thanks!