Update ckpt conversion scripts (#25)

Author: JoeZijunZhou

jetstream/tools/maxtext/model_ckpt_conversion.sh

@@ -31,7 +31,8 @@ export MODEL_NAME=${MODEL}-${MODEL_VARIATION}
 # After downloading checkpoints, copy them to GCS bucket at $CHKPT_BUCKET \
 # Please use seperate GCS paths for uploading open source model weights ($CHKPT_BUCKET) and MaxText compatible weights ($MODEL_BUCKET).
 # Point these variables to a GCS bucket that you created.
-export CHKPT_BUCKET=gs://${USER}-maxtext/chkpt/${MODEL}/${MODEL_VARIATION}
+# An example of CHKPT_BUCKET could be: gs://${USER}-maxtext/chkpt/${MODEL}/${MODEL_VARIATION}
+export CHKPT_BUCKET=$3
 export MODEL_BUCKET=gs://${USER}-maxtext
 
 # Point `BASE_OUTPUT_DIRECTORY` to a GCS bucket that you created, this bucket will store all the files generated by MaxText during a run.
@@ -47,11 +48,6 @@ gcloud storage buckets create ${MODEL_BUCKET} --location=${BUCKET_LOCATION} || t
 gcloud storage buckets create ${BASE_OUTPUT_DIRECTORY} --location=${BUCKET_LOCATION} || true
 gcloud storage buckets create ${DATASET_PATH} --location=${BUCKET_LOCATION} || true
 
-# Copy the downloaded checkpoints to `CHKPT_BUCKET`.
-# Gemma example: gsutil -m cp -r 7b ${CHKPT_BUCKET}
-# Llama2 example: gsutil -m cp -r llama-2-7b ${CHKPT_BUCKET}
-sudo gsutil -m cp -r $3 ${CHKPT_BUCKET}
-
 # Covert model checkpoints to MaxText compatible checkpoints.
 if [ "$MODEL" == "gemma" ]; then
     CONVERT_CKPT_SCRIPT="convert_gemma_chkpt.py"

jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh

@@ -31,7 +31,8 @@ export MODEL_NAME=${MODEL}-${MODEL_VARIATION}
 # After downloading checkpoints, copy them to GCS bucket at $CHKPT_BUCKET \
 # Please use seperate GCS paths for uploading open source model weights ($CHKPT_BUCKET) and MaxText compatible weights ($MODEL_BUCKET).
 # Point these variables to a GCS bucket that you created.
-export CHKPT_BUCKET=gs://${USER}-maxtext/chkpt/${MODEL}/${MODEL_VARIATION}
+# An example of CHKPT_BUCKET could be: gs://${USER}-maxtext/chkpt/${MODEL}/${MODEL_VARIATION}
+export CHKPT_BUCKET=$3
 export MODEL_BUCKET=gs://${USER}-maxtext
 
 # Point `BASE_OUTPUT_DIRECTORY` to a GCS bucket that you created, this bucket will store all the files generated by MaxText during a run.
@@ -41,7 +42,7 @@ export BASE_OUTPUT_DIRECTORY=gs://${USER}-runner-maxtext-logs
 export DATASET_PATH=gs://${USER}-maxtext-dataset
 
 # Prepare C4 dataset for fine tuning: https://github.com/allenai/allennlp/discussions/5056
-sudo gsutil -u $3 -m cp 'gs://allennlp-tensorflow-datasets/c4/en/3.0.1/*' ${DATASET_PATH}/c4/en/3.0.1/
+sudo gsutil -u $4 -m cp 'gs://allennlp-tensorflow-datasets/c4/en/3.0.1/*' ${DATASET_PATH}/c4/en/3.0.1/
 
 # We define `CONVERTED_CHECKPOINT` to refer to the checkpoint subdirectory.
 export CONVERTED_CHECKPOINT=${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}/0/items