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

[RexBearIU]

Description

Fixes end-to-end TPU post-training test suite configurations:

1. Gemma-3-4B Multimodal SFT (tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh):

   • Added export DATASET_PATH=${DATASET_PATH:-gs://maxtext-dataset} to point directly to GCS ChartQA parquet shards matching test_qwen3_multimodal_sft.sh.

2. Gemma-4-26B RL (tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh):

   • Updated vLLM decode parallelism to ici_tensor_parallelism=4 ici_expert_parallelism=2 and hbm_utilization_vllm=0.6 for MoE sharding.
   • Updated GRPO RL training to train_micro_batch_size=4 rollout_micro_batch_size=8 ici_expert_parallelism=8 rollout_tensor_parallelism=4 to fit the 8 trainer / 8 sampler device partition on v5p-32.
   • Added vllm_additional_config to pass model name during Step 3 checkpoint decode.

3. LLaMA-3.1-70B to HuggingFace (tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_to_hf.sh):

   • Set checkpoint_storage_concurrent_gb=80 and --parallel_threads=1 to optimize host RAM overhead during shard writing.

Tests

• Verified test_gemma3_multimodal_sft.sh on TPU v5p-8.
• Verified test_gemma4_rl.sh on TPU v5p-32.
• Verified test_llama3.1_70b_to_hf.sh SafeTensors unrolling on TPU v5p-8.

Checklist

• Pre-commit / linter checks pass (codespell, pylint, pyink).
• All test scripts execute cleanly with standard parameters.

[gemini-code-assist]

Code Review

This pull request updates several end-to-end TPU test scripts (Gemma 3, Gemma 4, and LLaMA 3.1) to adjust configurations, parallelism settings, batch sizes, and cleanup steps. Key feedback includes a recommendation to guard a recursive gcloud storage rm command in test_gemma3_multimodal_sft.sh to prevent accidental deletion of parent directories if ${run_id} is empty, and the removal of an unused vllm_additional_config parameter in test_gemma4_rl.sh.

[gemini-code-assist]

Running gcloud storage rm --recursive on a path that could potentially have an empty ${run_id} is dangerous, as it could result in the accidental deletion of the entire parent directory (${BASE_OUTPUT_DIRECTORY}/multimodal/sft/). It is highly recommended to guard this command with a check ensuring ${run_id} is non-empty.

Suggested change

	gcloud storage rm --recursive ${BASE_OUTPUT_DIRECTORY}/multimodal/sft/${run_id} || true
	if [ -n "${run_id}" ]; then gcloud storage rm --recursive ${BASE_OUTPUT_DIRECTORY}/multimodal/sft/${run_id} || true; fi

[gemini-code-assist]

The vllm_additional_config parameter is passed to vllm_decode.py here, but vllm_decode.py does not actually parse or use config.vllm_additional_config in either decode_with_vllm or decode_with_tunix. Passing this parameter is a no-op and should be removed to avoid confusion.

Suggested change

	vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
	hbm_utilization_vllm=0.85 \
	vllm_additional_config='{"maxtext_config": {"model_name": "gemma4-26b", "log_config": "false", "prefuse_moe_weights": "true"}}' \
	vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!