Fix roundtrip checkpoint conversion test

[hengtaoguo]

Description

The qwen3-30b-a3b roundtrip checkpoint conversion test was failing in CI (logs). There were two underlying causes:

1. Unknown flag save_dtype: This latent argument-isolation bug was recently exposed by the JAX 0.10.2 dependency upgrade, which changed when the Abseil-backed configuration was initialized. --save_dtype was parsed by the converter's local argparse parser but remained in the global sys.argv. When JAX initialized Abseil, Abseil attempted to parse the same argument and raised UnrecognizedFlagError because save_dtype is not an Abseil flag. This change replaces sys.argv with the remaining MaxText arguments before JAX initialization.

2. NotImplementedError on cpu: For flash or autoselected attention kernels, AttentionOp assumed GPU hardware when the platform was not TPU. It did not check if the active platform was actually CPU, leading to it attempting to call the GPU-only PallasTritonFlashAttention and throwing a NotImplementedError on CPU. We fixed this by adding an explicit check for target_hardware == "cpu" to fallback to standard dot-product attention (apply_attention_dot).

   • Previously, when converting models whose default attention kernel is not supported on CPU, we had to manually specify attention=dot_product to avoid crashes. With this patch, AttentionOp now automatically handles this fallback, allowing the conversion to run smoothly on CPU using default configurations..

Tests

Covered in CI (logs).

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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/checkpoint_conversion/to_maxtext.py	0.00%	1 Missing ⚠️
src/maxtext/layers/attention_op.py	0.00%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.