From ebd3fc0eb399c5d40820fbcc5e86937765cd902a Mon Sep 17 00:00:00 2001 From: Jacky Fang Date: Wed, 19 Aug 2026 14:28:21 +0000 Subject: [PATCH] fix(e2e): fix end-to-end post-training pipelines and optimize checkpoint conversion memory - Gemma 4 26B RL: Fix HLO temporary memory exceedance on v5p-32 by setting 4-way expert parallelism (ici_expert_parallelism=4, ici_fsdp_parallelism=2, ici_data_parallelism=1, batch_size=4, train_micro_batch_size=4) and clean up script formatting. - Checkpoint Conversion to Hugging Face: - Prevent host RAM exhaustion during large 70B conversion by progressively popping/freeing MaxText weights during layer transformation. - Stream SafeTensors shards to GCS with bounded in-worker ThreadPoolExecutor serialization, capping ephemeral disk usage and avoiding memory spikes. - Gemma 3 Multimodal SFT: Configure checkpoint_storage_concurrent_gb=30 and cleanup temporary run directory to ensure smooth SFT execution. - LLaMA 3.1 70B RL: Set checkpoint_storage_use_zarr3=False and checkpoint_storage_use_ocdbt=False. --- .../checkpoint_conversion/to_huggingface.py | 19 +++++-- .../checkpoint_conversion/utils/utils.py | 56 +++++++++++++------ .../gemma3/4b/test_gemma3_multimodal_sft.sh | 10 +++- .../tpu/gemma4/26b/test_gemma4_rl.sh | 38 +++++++++---- .../tpu/llama3.1/70b/test_llama3.1_70b_rl.sh | 7 +-- .../llama3.1/70b/test_llama3.1_70b_to_hf.sh | 3 +- 6 files changed, 95 insertions(+), 38 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/to_huggingface.py b/src/maxtext/checkpoint_conversion/to_huggingface.py index 80a30516f1..cf586f0752 100644 --- a/src/maxtext/checkpoint_conversion/to_huggingface.py +++ b/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) diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py index 449b2d2194..abf6cca6df 100644 --- a/src/maxtext/checkpoint_conversion/utils/utils.py +++ b/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: diff --git a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh index 613c2043ec..72ca98c8ea 100644 --- a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh +++ b/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 diff --git a/tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh b/tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh index a2f1b7a1b7..ba2e294e57 100644 --- a/tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh +++ b/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 \ No newline at end of file + use_chat_template=True \ + scan_layers=false \ + enable_single_controller=${use_pathways} \ + prefuse_moe_weights=True \ + ici_tensor_parallelism=2 \ No newline at end of file diff --git a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh b/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh index 3c4ce752a4..e50fab88f0 100644 --- a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh +++ b/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 \ diff --git a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_to_hf.sh b/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_to_hf.sh index 9516ac84ae..af09f41dbb 100644 --- a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_to_hf.sh +++ b/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 \ No newline at end of file + checkpoint_storage_use_zarr3=False \ + checkpoint_storage_use_ocdbt=False \ No newline at end of file