address comments

Author: hengtaoguo

src/maxtext/checkpoint_conversion/to_maxtext.py

@@ -822,7 +822,7 @@ def _merged_getter(key):
 def main(
     args: Sequence[str],
     lazy_load_tensors: bool = False,
-    eager_load_method: str = "transformers",
+    eager_load_method: str = "safetensors",
     hf_model_path: str | None = None,
     revision: str | None = None,
     save_dtype: str = "bfloat16",
@@ -909,9 +909,9 @@ def main(
       #      (e.g., Multi-Token Prediction weights (`layers.61`) in DeepSeek-V3).
       #
       # Recommendation:
-      # - Use 'transformers' as the default for backward compatibility of mapping.
-      # - 'safetensors' is an interchangeable and valid alternative for most models,
-      #   and is strictly required if the model or specific weights lack Transformers support.
+      # - Use 'safetensors' as the default. Since transformers 5.8.0, model initialization
+      #   changed and the 'transformers' method may produce different key structures.
+      # - Use 'transformers' only if explicitly needed for backward-compatible key mapping.
       if eager_load_method == "transformers":
         max_logging.log("Eager load with Transformers backend, from_pretrained with auto dtype")
         # For auto mode, loaded dtype is the same as `dtype` specified in config.json (or `torch_dtype` for older version)

tests/integration/checkpoint_conversion_test.py

@@ -32,30 +32,29 @@ def test_qwen3_30b_a3b_roundtrip_conversion(self):
     model_name = "qwen3-30b-a3b"
     base_num_decoder_layers = 2
     with tempfile.TemporaryDirectory() as tmpdir:
-      cmd = [
+      to_mt_cmd = [
           sys.executable,
           "-m",
           "maxtext.checkpoint_conversion.to_maxtext",
-          "src/maxtext/configs/base.yml",
           f"model_name={model_name}",
           f"base_output_directory={tmpdir}",
           f"base_num_decoder_layers={base_num_decoder_layers}",
           "override_model_config=True",
           "scan_layers=False",
-          "weight_dtype=bfloat16",
           "hardware=cpu",
           "skip_jax_distributed_system=True",
           "checkpoint_storage_use_ocdbt=False",
           "checkpoint_storage_use_zarr3=False",
+          "--save_dtype=bfloat16",
           "--lazy_load_tensors=True",
       ]
       env = os.environ.copy()
       env["JAX_PLATFORMS"] = "cpu"
       env["HF_HOME"] = tmpdir
 
-      print("Running checkpoint conversion command:", " ".join(cmd))
+      print("Running checkpoint conversion command:", " ".join(to_mt_cmd))
       # Inherit stdout and stderr from the parent process to stream logs in real time
-      subprocess.run(cmd, env=env, check=True)
+      subprocess.run(to_mt_cmd, env=env, check=True)
 
       # Verify output directory exists and contains items
       # Output structure: tmpdir/0/items
@@ -65,7 +64,7 @@ def test_qwen3_30b_a3b_roundtrip_conversion(self):
 
       # Roundtrip conversion back to HuggingFace format
       expected_hf_dir = os.path.join(tmpdir, "hf_safetensor", "qwen3-30b-a3b")
-      hf_cmd = [
+      to_hf_cmd = [
           sys.executable,
           "-m",
           "maxtext.checkpoint_conversion.to_huggingface",
@@ -80,8 +79,8 @@ def test_qwen3_30b_a3b_roundtrip_conversion(self):
           "skip_jax_distributed_system=True",
           "--override_model_architecture=True",
       ]
-      print("Running roundtrip checkpoint conversion command (to HF):", " ".join(hf_cmd))
-      subprocess.run(hf_cmd, env=env, check=True)
+      print("Running roundtrip checkpoint conversion command (to HF):", " ".join(to_hf_cmd))
+      subprocess.run(to_hf_cmd, env=env, check=True)
 
       # Verify HF output directory exists and contains safetensors/config files
       self.assertTrue(