fix(e2e): Fix dataset path for multimodal SFT, RL parallelism, and LLaMA to_hf settings

src/maxtext/checkpoint_conversion/to_huggingface.py

@@ -66,11 +66,12 @@
     scan_layers=True
 """
 
-import jax
-import jax.numpy as jnp
+import gc
 import os
-from typing import Sequence
 import time
+from typing import Sequence
+import jax
+import jax.numpy as jnp
 
 from transformers import AutoTokenizer, AutoProcessor
 
@@ -403,15 +404,21 @@ def _transform_weights_to_full_model(config, filtered_map_keys, state_dict, para
   processed_params_list = []
   lora_scaling = config.lora.lora_alpha / config.lora.lora_rank if config.lora.lora_rank > 0 else 1.0
   for key in MemoryMonitorTqdm(filtered_map_keys, leave=True):
-    weight = [state_dict[subkey] for subkey in key] if isinstance(key, tuple) else state_dict.get(key)
-    if weight is not None and not isinstance(key, tuple):
+    if isinstance(key, tuple):
+      weight = [state_dict.pop(subkey, None) for subkey in key]
+      delta = None
+    else:
       delta = _get_lora_delta(key, state_dict, lora_scaling)
-      if delta is not None:
+      weight = state_dict.pop(key, None)
+      if weight is not None and delta is not None:
         if delta.shape != weight.shape and delta.size == weight.size:
           delta = delta.reshape(weight.shape)
         weight = (jnp.asarray(weight, dtype=jnp.float32) + delta).astype(weight.dtype)
     if weight is not None:
       processed_params_list.extend(process_maxtext_param(key, weight, param_map, hook_fn_map, shape_map, config))
+      del weight
+      del delta
+      gc.collect()
   return dict(processed_params_list)
 
 

src/maxtext/checkpoint_conversion/utils/utils.py

@@ -479,8 +479,12 @@ def save_safetensor_file(
       state_dict = state_dict["model.safetensors"]
 
     if output_dir_final.startswith("gs://"):
+      local_path = os.path.join(local_dir_to_save_to, file_name)
+      numpy_save_file(state_dict, local_path, metadata={"format": "pt"})
+      state_dict.clear()
+      gc.collect()
       cloud_path = os.path.join(output_dir_final, file_name)
-      upload_state_dict_to_gcs(state_dict=state_dict, gs_bucket_path=cloud_path)
+      upload_file_to_gcs(local_file=local_path, gs_bucket_path=cloud_path, remove_local_file_after_upload=True)
     elif output_dir_final.startswith("hf://"):
       max_logging.log(f"  Serializing {file_name} to memory for Hugging Face Hub upload...")
       serialized_content = save_flax_to_bytes(state_dict, metadata={"format": "pt"})
@@ -556,21 +560,39 @@ def save_weight_files(
     # 'shards' is actually the single state_dict here
     save_safetensor_file(shards, local_dir_to_save_to, output_dir_final, SAFE_TENSORS_WEIGHTS_FILE)
   else:
-    # Save sharded weights in parallel
-    with ThreadPoolExecutor(max_workers=parallel_threads) as executor:
-      shard_items = list(shards.items())
-      futures = [
-          executor.submit(
-              save_safetensor_file,
-              shard_dict,
-              local_dir_to_save_to,
-              output_dir_final,
-              shard_name,
-          )
-          for shard_name, shard_dict in shard_items
-      ]
-      for future in futures:
-        future.result()
+    if output_dir_final.startswith("gs://"):
+
+      def _save_and_upload(shard_name, shard_dict):
+        local_path = os.path.join(local_dir_to_save_to, shard_name)
+        numpy_save_file(shard_dict, local_path, metadata={"format": "pt"})
+        shard_dict.clear()
+        del shard_dict
+        gc.collect()
+        cloud_path = os.path.join(output_dir_final, shard_name)
+        upload_file_to_gcs(local_file=local_path, gs_bucket_path=cloud_path, remove_local_file_after_upload=True)
+
+      with ThreadPoolExecutor(max_workers=parallel_threads) as executor:
+        futures = []
+        for shard_name in list(shards.keys()):
+          shard_dict = shards.pop(shard_name)
+          if "model.safetensors" in shard_dict and isinstance(shard_dict["model.safetensors"], dict):
+            shard_dict = shard_dict["model.safetensors"]
+          shard_dict = {k: v for k, v in shard_dict.items() if v is not None}
+          futures.append(executor.submit(_save_and_upload, shard_name, shard_dict))
+
+        for future in futures:
+          future.result()
+    else:
+      for shard_name in list(shards.keys()):
+        shard_dict = shards.pop(shard_name)
+        save_safetensor_file(
+            shard_dict,
+            local_dir_to_save_to,
+            output_dir_final,
+            shard_name,
+        )
+        shard_dict.clear()
+        gc.collect()
 
     # Save index file
     save_index_file(
@@ -669,6 +691,8 @@ def save_model_files(
     # The actual file saving within save_weight_files is guarded.
     # Unwrap nested dict if needed
     shards, index = shard_checkpoint(weight_arrays)
+    weight_arrays.clear()
+    gc.collect()
     save_weight_files(shards, index, current_save_path, output_dir, parallel_threads, remove_local_copy)
 
   if jax.process_index() == 0:

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

@@ -25,6 +25,9 @@ MODEL_NAME='gemma3-4b'
 BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}
 MULTIMODAL_SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned_multimodal/${run_id}/0/items
 
+# Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data
+export DATASET_PATH=${DATASET_PATH:-gs://maxtext-dataset}
+
 # Step 1: Install google-jetstream
 python3 -m pip install google-jetstream@https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip --no-deps
 
@@ -50,6 +53,9 @@ python3 -m maxtext.inference.decode \
     skip_jax_distributed_system=True
 
 # Step 3: Run SFT on the MaxText checkpoint on ChartQA dataset
+if [ -n "${run_id}" ]; then
+  gcloud storage rm --recursive "${BASE_OUTPUT_DIRECTORY}/multimodal/sft/${run_id}" || true
+fi
 python -m maxtext.trainers.post_train.sft.train_sft_native "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/post_train/sft-vision-chartqa.yml \
     run_name=${run_id} \
     model_name=${MODEL_NAME} \
@@ -60,7 +66,8 @@ python -m maxtext.trainers.post_train.sft.train_sft_native "${MAXTEXT_CONFIGS_DI
     scan_layers=true \
     async_checkpointing=False \
     attention=\'dot_product\' \
-    dataset_type=hf hf_path=parquet \
+    dataset_type=hf \
+    hf_path=parquet \
     hf_train_files=${DATASET_PATH}/hf/chartqa/train-* \
     base_output_directory=${BASE_OUTPUT_DIRECTORY}/multimodal/sft \
     load_parameters_path=${MULTIMODAL_SCANNED_CKPT_PATH} \
@@ -69,6 +76,7 @@ python -m maxtext.trainers.post_train.sft.train_sft_native "${MAXTEXT_CONFIGS_DI
     sharding_tolerance=0.05 \
     checkpoint_storage_use_zarr3=False \
     checkpoint_storage_use_ocdbt=False \
+    checkpoint_storage_concurrent_gb=30 \
     enable_single_controller=${use_pathways} \
     grain_worker_count=0
 

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

@@ -32,20 +32,35 @@ python3 -m maxtext.inference.vllm_decode \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
     hbm_utilization_vllm=0.85 \
     prompt="Suggest some famous landmarks in London." \
-    use_chat_template=True scan_layers=false enable_single_controller=${use_pathways} \
-    prefuse_moe_weights=True ici_tensor_parallelism=8
+    use_chat_template=True \
+    scan_layers=false \
+    enable_single_controller=${use_pathways} \
+    prefuse_moe_weights=True \
+    ici_tensor_parallelism=2
 
 # 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 \
-    num_batches=2 batch_size=16 num_test_batches=2 \
-    model_name=${MODEL_NAME} enable_single_controller=${use_pathways} \
-    checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
-    rollout_tensor_parallelism=4 \
+    run_name=${run_id} \
+    rl.loss_algo='grpo' \
+    scan_layers=false \
+    num_batches=2 \
+    batch_size=4 \
+    train_micro_batch_size=4 \
+    ici_expert_parallelism=4 \
+    ici_fsdp_parallelism=2 \
+    ici_data_parallelism=1 \
+    num_test_batches=2 \
+    model_name=${MODEL_NAME} \
+    enable_single_controller=${use_pathways} \
+    checkpoint_storage_use_zarr3=False \
+    checkpoint_storage_use_ocdbt=False \
+    rollout_tensor_parallelism=2 \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
-    vllm_additional_config='{"maxtext_config": {"model_name": "gemma4-26b", "log_config": "false", "prefuse_moe_weights": "true"}}'
+    vllm_additional_config='{"maxtext_config": {"model_name": "gemma4-26b", "log_config": "false", "prefuse_moe_weights": "true"}}' \
+    remat_policy=full \
+    hbm_utilization_vllm=0.55
 
 # Step 3: Run inference on the checkpoint generated from the previous run
 python3 -m maxtext.inference.vllm_decode \
@@ -54,5 +69,8 @@ python3 -m maxtext.inference.vllm_decode \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
     hbm_utilization_vllm=0.85 \
     prompt='Suggest some famous landmarks in London.' \
-    use_chat_template=True scan_layers=false enable_single_controller=${use_pathways} \
-    prefuse_moe_weights=True ici_tensor_parallelism=8
+    use_chat_template=True \
+    scan_layers=false \
+    enable_single_controller=${use_pathways} \
+    prefuse_moe_weights=True \
+    ici_tensor_parallelism=2

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

@@ -21,8 +21,7 @@ use_pathways=${2:-false}
 MODEL_NAME='llama3.1-70b'
 
 # Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to the GCS paths where you have the scanned and unscanned checkpoints stored
-BASE_OUTPUT_DIRECTORY=${BASE_OUTPUT_DIRECTORY:-gs://runner-maxtext-logs/${MODEL_NAME}}
-OUTPUT_DIR=${GCS_OUTPUT:-${BASE_OUTPUT_DIRECTORY}}
+BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}
 UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items
 SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned/${run_id}/0/items
 
@@ -39,7 +38,7 @@ python3 -m maxtext.inference.vllm_decode \
 
 # Step 2: Run RL on the converted checkpoint
 python3 -m maxtext.trainers.post_train.rl.train_rl \
-    base_output_directory=${OUTPUT_DIR}/rl \
+    base_output_directory=${BASE_OUTPUT_DIRECTORY}/rl \
     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 \
@@ -55,7 +54,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \
 # Step 3: Run inference on the checkpoint generated from the previous run
 python3 -m maxtext.inference.vllm_decode \
     model_name=${MODEL_NAME} \
-    load_parameters_path=${OUTPUT_DIR}/rl/${run_id}/checkpoints/actor/2/model_params \
+    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 \

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

@@ -30,4 +30,5 @@ python3 -m maxtext.checkpoint_conversion.to_huggingface \
     use_multimodal=${USE_MULTIMODAL} \
     scan_layers=$SCAN_LAYERS \
     weight_dtype=bfloat16 \
-    --parallel_threads=2
+    checkpoint_storage_use_zarr3=False \
+    checkpoint_storage_use_ocdbt=False