Fix OOM in to_huggingface conversion and stabilize post-training E2E TPU pipelines

.github/scripts/analyze_code_changes.sh

@@ -112,6 +112,7 @@ EXCLUDED_FILES=(
   '^\.github/scripts/'
   '^tools/'
   '\.md$'
+  '^tests/end_to_end/'
 )
 
 # Loop through every changed file
@@ -156,7 +157,7 @@ while IFS= read -r file; do
   fi
 
   # GPU configs/source/test changes
-  if matches_pattern "$file" "src/maxtext/configs/gpu/|src/maxtext/inference/gpu/|tests/end_to_end/gpu/"; then
+  if matches_pattern "$file" "src/maxtext/configs/gpu/|src/maxtext/inference/gpu/"; then
     echo "GPU files changed, enabling GPU tests."
     enable_flags run_tests run_gpu_tests
     matched=true

src/maxtext/checkpoint_conversion/to_huggingface.py

@@ -139,8 +139,8 @@ def _get_lora_delta(key, lora_state_dict, lora_scaling):
     a_key, b_key = key[7:] + "_lora_a", key[7:] + "_lora_b"
 
   if a_key in lora_state_dict and b_key in lora_state_dict:
-    data_a = jnp.asarray(lora_state_dict[a_key], dtype=jnp.float32)
-    data_b = jnp.asarray(lora_state_dict[b_key], dtype=jnp.float32)
+    data_a = jnp.asarray(lora_state_dict.pop(a_key), dtype=jnp.float32)
+    data_b = jnp.asarray(lora_state_dict.pop(b_key), dtype=jnp.float32)
 
     is_attention = "attention" in key.lower() or "attn" in key.lower()
 
@@ -545,6 +545,8 @@ def main(argv: Sequence[str]) -> None:
 
   # 4. Extract and transform weights for Linen/NNX-SFT/NNX-RL checkpoints
   maxtext_state_dict = detect_and_extract_checkpoint(checkpoint_dict)
+  del checkpoint_dict
+  gc.collect()
 
   # Validate that checkpoint keys match the parameter mapping
   state_keys = {k.replace("_lora_a", "").replace("_lora_b", "") for k in maxtext_state_dict}

src/maxtext/checkpoint_conversion/utils/utils.py

@@ -33,6 +33,7 @@
 from etils import epath
 
 import jax
+import jax.numpy as jnp
 from jax import tree
 from jax.experimental import multihost_utils
 from jaxtyping import Array
@@ -879,10 +880,28 @@ def load_orbax_checkpoint(config) -> dict:
   devices = np.array(jax.devices()).reshape((-1,))
   single_device_mesh = jax.sharding.Mesh(devices, ("x",))
 
+  target_dtype = None
+  if getattr(config, "weight_dtype", None):
+    try:
+      target_dtype = jnp.dtype(config.weight_dtype)
+    except Exception:  # pylint: disable=broad-exception-caught
+      target_dtype = None
+
   def create_restore_args(tree_metadata):
     """Create restore args for unsharded restoration."""
     if hasattr(tree_metadata, "shape"):
-      return ocp.ArrayRestoreArgs(sharding=jax.sharding.NamedSharding(single_device_mesh, jax.sharding.PartitionSpec()))
+      leaf_dtype = getattr(tree_metadata, "dtype", None)
+      restore_dtype = None
+      if target_dtype is not None and leaf_dtype is not None:
+        try:
+          if jnp.issubdtype(leaf_dtype, jnp.floating):
+            restore_dtype = target_dtype
+        except Exception:  # pylint: disable=broad-exception-caught
+          restore_dtype = None
+      return ocp.ArrayRestoreArgs(
+          dtype=restore_dtype,
+          sharding=jax.sharding.NamedSharding(single_device_mesh, jax.sharding.PartitionSpec()),
+      )
     elif isinstance(tree_metadata, dict):
       return {k: create_restore_args(v) for k, v in tree_metadata.items()}
     else:
@@ -901,10 +920,10 @@ def create_restore_args(tree_metadata):
         checkpoint_tree = {"params": checkpoint_tree["params"]}
         max_logging.log(f"Filtering checkpoint to only load 'params' from {path}")
       else:
-        filtered_tree = {k: v for k, v in checkpoint_tree.items() if k not in ("opt_state", "optimizer")}
+        filtered_tree = {k: v for k, v in checkpoint_tree.items() if k not in ("opt_state", "optimizer", "rngs", "step")}
         if len(filtered_tree) < len(checkpoint_tree):
           checkpoint_tree = filtered_tree
-          max_logging.log(f"Filtering checkpoint to exclude optimizer keys from {path}")
+          max_logging.log(f"Filtering checkpoint to exclude non-param keys from {path}")
 
     restore_args = jax.tree_util.tree_map(
         lambda x: create_restore_args(x) if hasattr(x, "shape") else None,

src/maxtext/integration/vllm/maxtext_vllm_rollout.py

@@ -538,6 +538,7 @@ def __init__(
         direct_maxtext_sync=direct_maxtext_sync,
         num_experts=getattr(maxtext_config, "num_experts", 1),
         tensor_parallel_size=rollout_config.tensor_parallel_size,
+        weight_dtype=getattr(maxtext_config, "weight_dtype", None),
     )
 
     self._sampler = MaxTextVllmSampler(

tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh

@@ -23,6 +23,8 @@ export JAX_RANDOM_WEIGHTS='1'
 export VLLM_ENABLE_V1_MULTIPROCESSING='0'
 export SKIP_JAX_PRECOMPILE='1'
 export NEW_MODEL_DESIGN='0'
+export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION='upb'
+export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY='PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'
 
 run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
 use_pathways=${2:-false}
@@ -39,22 +41,28 @@ python3 -m maxtext.inference.vllm_decode \
     load_parameters_path=${UNSCANNED_CKPT_PATH} \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
     hbm_utilization_vllm=0.5 \
+    weight_dtype=bfloat16 dtype=bfloat16 \
     prompt='Suggest some famous landmarks in London.' \
-    use_chat_template=True scan_layers=false enable_single_controller=${use_pathways}
+    use_chat_template=True scan_layers=false enable_single_controller=${use_pathways} \
+    ici_tensor_parallelism=4
 
 # Step 2: Run RL on the converted checkpoint
 python3 -m maxtext.trainers.post_train.rl.train_rl \
     base_output_directory=${BASE_OUTPUT_DIRECTORY}/rl \
-    load_parameters_path=${UNSCANNED_CKPT_PATH} \
-    run_name=${run_id} rl.loss_algo='grpo' scan_layers=false \
+    load_parameters_path=${SCANNED_CKPT_PATH} \
+    run_name=${run_id} rl.loss_algo='grpo' scan_layers=true \
     num_batches=2 batch_size=1 num_test_batches=2 \
     model_name=${MODEL_NAME} enable_single_controller=${use_pathways} \
+    use_pathways=${use_pathways} \
     checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
-    chips_per_vm=4 \
     rollout_tensor_parallelism=4 \
+    remat_policy=full \
+    max_target_length=512 \
+    weight_dtype=bfloat16 dtype=bfloat16 \
+    hbm_utilization_vllm=0.5 opt_type=sgd \
+    rl.reshard_chunk_size=32 \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
-    vllm_additional_config='{"maxtext_config": {"model_name": "gemma3-4b", "log_config": "false"}}' \
-    hbm_utilization_vllm=0.5
+    vllm_additional_config='{"maxtext_config": {"model_name": "gemma3-4b", "weight_dtype": "bfloat16", "dtype": "bfloat16", "log_config": "false"}}'
 
 
 # Step 3: Run inference on the checkpoint generated from the previous run
@@ -63,5 +71,7 @@ python3 -m maxtext.inference.vllm_decode \
     load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/model_params \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
     hbm_utilization_vllm=0.5 \
+    weight_dtype=bfloat16 dtype=bfloat16 \
     prompt='Suggest some famous landmarks in London.' \
-    use_chat_template=True scan_layers=false enable_single_controller=${use_pathways}
+    use_chat_template=True scan_layers=true enable_single_controller=${use_pathways} \
+    ici_tensor_parallelism=4

tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh

@@ -16,6 +16,16 @@
 
 set -ex
 
+export VLLM_WORKER_MULTIPROC_METHOD='spawn'
+export MODEL_IMPL_TYPE='flax_nnx'
+export GRPC_ENABLE_FORK_SUPPORT='0'
+export JAX_RANDOM_WEIGHTS='1'
+export VLLM_ENABLE_V1_MULTIPROCESSING='0'
+export SKIP_JAX_PRECOMPILE='1'
+export NEW_MODEL_DESIGN='0'
+export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION='upb'
+export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY='PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'
+
 run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
 use_pathways=${2:-false}
 MODEL_NAME='gemma4-26b'
@@ -30,7 +40,8 @@ python3 -m maxtext.inference.vllm_decode \
     model_name=${MODEL_NAME} \
     load_parameters_path=${UNSCANNED_CKPT_PATH} \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
-    hbm_utilization_vllm=0.85 \
+    hbm_utilization_vllm=0.55 \
+    weight_dtype=bfloat16 dtype=bfloat16 \
     prompt="Suggest some famous landmarks in London." \
     use_chat_template=True \
     scan_layers=false \
@@ -67,7 +78,8 @@ python3 -m maxtext.inference.vllm_decode \
     model_name=${MODEL_NAME} \
     load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/model_params \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
-    hbm_utilization_vllm=0.85 \
+    hbm_utilization_vllm=0.55 \
+    weight_dtype=bfloat16 dtype=bfloat16 \
     prompt='Suggest some famous landmarks in London.' \
     use_chat_template=True \
     scan_layers=false \

tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh

@@ -16,6 +16,16 @@
 
 set -ex
 
+export VLLM_WORKER_MULTIPROC_METHOD='spawn'
+export MODEL_IMPL_TYPE='flax_nnx'
+export GRPC_ENABLE_FORK_SUPPORT='0'
+export JAX_RANDOM_WEIGHTS='1'
+export VLLM_ENABLE_V1_MULTIPROCESSING='0'
+export SKIP_JAX_PRECOMPILE='1'
+export NEW_MODEL_DESIGN='0'
+export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION='upb'
+export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY='PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'
+
 run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
 use_pathways=${2:-false}
 MODEL_NAME='llama3.1-70b'
@@ -31,7 +41,8 @@ python3 -m maxtext.inference.vllm_decode \
     load_parameters_path=${UNSCANNED_CKPT_PATH} \
     tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
-    hbm_utilization_vllm=0.85 \
+    hbm_utilization_vllm=0.55 \
+    weight_dtype=bfloat16 dtype=bfloat16 \
     prompt='Suggest some famous landmarks in London.' \
     use_chat_template=True scan_layers=false enable_single_controller=${use_pathways} \
     ici_tensor_parallelism=8
@@ -44,11 +55,16 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \
     num_batches=2 batch_size=1 num_test_batches=2 \
     model_name=${MODEL_NAME} tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \
     enable_single_controller=${use_pathways} \
+    use_pathways=${use_pathways} \
+    max_target_length=512 \
     remat_policy=full \
+    weight_dtype=bfloat16 dtype=bfloat16 \
+    hbm_utilization_vllm=0.55 opt_type=sgd \
+    rl.reshard_chunk_size=32 \
     checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
-    rollout_tensor_parallelism=4 \
+    rollout_tensor_parallelism=8 \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
-    vllm_additional_config='{"maxtext_config": {"model_name": "llama3.1-70b", "log_config": "false"}}'
+    vllm_additional_config='{"maxtext_config": {"model_name": "llama3.1-70b", "weight_dtype": "bfloat16", "dtype": "bfloat16", "log_config": "false"}}'
 
 
 # Step 3: Run inference on the checkpoint generated from the previous run
@@ -57,7 +73,8 @@ python3 -m maxtext.inference.vllm_decode \
     load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/model_params \
     tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
-    hbm_utilization_vllm=0.6 \
+    hbm_utilization_vllm=0.55 \
+    weight_dtype=bfloat16 dtype=bfloat16 \
     prompt='Suggest some famous landmarks in London.' \
     use_chat_template=True scan_layers=true enable_single_controller=${use_pathways} \
     ici_tensor_parallelism=8