Merge pull request #2588 from AI-Hypercomputer:end_to_end_fix

PiperOrigin-RevId: 827750218

Author: Google-ML-Automation

end_to_end/tpu/gemma/2b/test_gemma.sh

@@ -27,7 +27,7 @@ python3 -m MaxText.convert_gemma_chkpt --base_model_path ${CHKPT_BUCKET}/${MODEL
 # Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data
 export DATASET_PATH=gs://maxtext-dataset
 # Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to a GCS bucket that you own, this bucket will store all the files generated by MaxText during a run
-export BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs
+export BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/gemma-2b
 # We define `CONVERTED_CHECKPOINT` to refer to the checkpoint subdirectory. This way it is easier to use this path in the `train.py` and `decode.py` commands
 export CONVERTED_CHECKPOINT=${MODEL_BUCKET}/${MODEL_VARIATION}/${idx}/0/items
 export RUN_NAME=unscanned_chkpt_${idx}

end_to_end/tpu/gemma2/2b/test_gemma2.sh

@@ -29,7 +29,7 @@ python3 -m MaxText.convert_gemma2_chkpt --base_model_path ${CHKPT_BUCKET}/${MODE
 # Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data
 export DATASET_PATH=gs://maxtext-dataset
 # Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to a GCS bucket that you own, this bucket will store all the files generated by MaxText during a run
-export BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs
+export BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/gemma2-2b
 # We define `CONVERTED_CHECKPOINT` to refer to the checkpoint subdirectory. This way it is easier to use this path in the `train.py` and `decode.py` commands
 export CONVERTED_CHECKPOINT=${MODEL_BUCKET}/${MODEL_VARIATION}/${idx}/0/items
 export RUN_NAME=unscanned_chkpt_${idx}

end_to_end/tpu/gemma3/4b/test_gemma3.sh

@@ -31,7 +31,7 @@ python3 -m MaxText.convert_gemma3_chkpt --base_model_path ${CHKPT_BUCKET}/${MODE
 # Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data
 export DATASET_PATH=gs://maxtext-dataset
 # Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to a GCS bucket that you own, this bucket will store all the files generated by MaxText during a run
-export BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs
+export BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/gemma3-4b
 # We define `CONVERTED_CHECKPOINT` to refer to the checkpoint subdirectory. This way it is easier to use this path in the `train` and `decode` commands
 export CONVERTED_CHECKPOINT=${MODEL_BUCKET}/${MODEL_VARIATION}/${idx}/0/items
 export RUN_NAME=unscanned_chkpt_${idx}

end_to_end/tpu/test_convergence_1b_params.sh

@@ -50,9 +50,9 @@ then
     # We use a local copy of tokenizer from https://huggingface.co/meta-llama/Llama-2-7b-hf
     # Alternatively, you can set tokenizer_path="meta-llama/Llama-2-7b-hf" and hf_access_token="<your-token>" after gaining access through HF website.
     gsutil cp -r gs://maxtext-dataset/hf/llama2-tokenizer "${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/MaxText/assets}}"
-    CMD_DATA=" hf_path=parquet tokenizer_path=${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/MaxText/assets}}/llama2-tokenizer"\
-        "hf_train_files=$DATASET_PATH/hf/c4/c4-train-*.parquet" \
-        "hf_eval_files=$DATASET_PATH/hf/c4/c4-validation-*.parquet "
+    CMD_DATA=" hf_path=parquet tokenizer_path=${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/MaxText/assets}}/llama2-tokenizer \
+        hf_train_files=$DATASET_PATH/hf/c4/c4-train-*.parquet \
+        hf_eval_files=$DATASET_PATH/hf/c4/c4-validation-*.parquet "
 fi
 
 TRAIN_CMD="python3 -m MaxText.train ${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/MaxText}/configs/base.yml \