From eabc6debd21b80023756bafff6b6deb4c5896c43 Mon Sep 17 00:00:00 2001 From: hengtaoguo Date: Fri, 24 Jul 2026 23:21:23 +0000 Subject: [PATCH] Fix roundtrip checkpoint conversion test --- src/maxtext/checkpoint_conversion/to_maxtext.py | 1 + src/maxtext/layers/attention_op.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index ca50d67680..35f1ce38ca 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/src/maxtext/checkpoint_conversion/to_maxtext.py @@ -1155,6 +1155,7 @@ def _eager_getter(key): local_args, remaining_args = parser.parse_known_args() # Reconstruct model_args (script name + the args MaxText needs) model_args = [sys.argv[0]] + remaining_args + sys.argv = model_args # Set jax environment jax.config.update("jax_platforms", "cpu") diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index cf076b4794..dd73c8cc88 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -1124,8 +1124,8 @@ def apply_attention( return out, None, None else: - if model_mode == MODEL_MODE_AUTOREGRESSIVE: - # fallback to dot_product as pallas gpu flash attention doesn't support decode stage + if model_mode == MODEL_MODE_AUTOREGRESSIVE or target_hardware == "cpu": + # fallback to dot_product as pallas gpu flash attention doesn't support decode stage or CPU return self.apply_attention_dot( query, key,