From cf65dd6346b418fe740bd2068437a7ee89f2292c Mon Sep 17 00:00:00 2001
From: Zijun Zhou <zijunzhou@google.com>
Date: Fri, 29 Mar 2024 18:37:40 +0000
Subject: [PATCH 1/4] add model ckpt conversion scripts for JetStream MaxText
 Serving

---
 README.md                                     |  9 ++-
 benchmarks/eval_accuracy.py                   | 14 ++++
 jetstream/core/utils/async_multifuture.py     | 14 ++++
 jetstream/tests/engine/test_token_utils.py    | 14 ++++
 jetstream/{core => }/tools/load_tester.py     |  0
 .../tools/maxtext/model_ckpt_conversion.sh    | 81 +++++++++++++++++++
 .../maxtext/model_ckpt_finetune_with_aqt.sh   | 73 +++++++++++++++++
 jetstream/{core => }/tools/requester.py       |  0
 8 files changed, 203 insertions(+), 2 deletions(-)
 rename jetstream/{core => }/tools/load_tester.py (100%)
 create mode 100644 jetstream/tools/maxtext/model_ckpt_conversion.sh
 create mode 100644 jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh
 rename jetstream/{core => }/tools/requester.py (100%)

diff --git a/README.md b/README.md
index d98d42be..fce0b18a 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,11 @@ JetStream is a fast library for LLM inference and serving on TPUs.
 
 ## Getting Started
 
+### Setup
+```
+pip install -r requirements.txt
+```
+
 ### Run local server & Testing
 
 Use the following commands to run a server locally:
@@ -14,10 +19,10 @@ Use the following commands to run a server locally:
 python -m jetstream.core.implementations.mock.server
 
 # Test local mock server
-python -m jetstream.core.tools.requester
+python -m jetstream.tools.requester
 
 # Load test local mock server
-python -m jetstream.core.tools.load_tester
+python -m jetstream.tools.load_tester
 
 ```
 
diff --git a/benchmarks/eval_accuracy.py b/benchmarks/eval_accuracy.py
index a46ba78f..667dbb54 100644
--- a/benchmarks/eval_accuracy.py
+++ b/benchmarks/eval_accuracy.py
@@ -1,3 +1,17 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import argparse
 import nltk
 import evaluate
diff --git a/jetstream/core/utils/async_multifuture.py b/jetstream/core/utils/async_multifuture.py
index 8c844656..0ee5418d 100644
--- a/jetstream/core/utils/async_multifuture.py
+++ b/jetstream/core/utils/async_multifuture.py
@@ -1,3 +1,17 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import asyncio
 from concurrent import futures
 import threading
diff --git a/jetstream/tests/engine/test_token_utils.py b/jetstream/tests/engine/test_token_utils.py
index 033d3b72..eef78891 100644
--- a/jetstream/tests/engine/test_token_utils.py
+++ b/jetstream/tests/engine/test_token_utils.py
@@ -1,3 +1,17 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import os
 import unittest
 from typing import List
diff --git a/jetstream/core/tools/load_tester.py b/jetstream/tools/load_tester.py
similarity index 100%
rename from jetstream/core/tools/load_tester.py
rename to jetstream/tools/load_tester.py
diff --git a/jetstream/tools/maxtext/model_ckpt_conversion.sh b/jetstream/tools/maxtext/model_ckpt_conversion.sh
new file mode 100644
index 00000000..5e89d35c
--- /dev/null
+++ b/jetstream/tools/maxtext/model_ckpt_conversion.sh
@@ -0,0 +1,81 @@
+#!/bin/bash
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+set -ex
+
+idx=$(date +%Y-%m-%d-%H-%M)
+# Modify the `MODEL` and `MODEL_VARIATION` based on the model you use.
+export MODEL=$1
+export MODEL_VARIATION=$2
+export MODEL_NAME=${MODEL}-${MODEL_VARIATION}
+
+# After downloading checkpoints, copy them to GCS bucket at $CHKPT_BUCKET \
+# Please use seperate GCS paths for uploading model weights from kaggle ($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}
+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.
+export BASE_OUTPUT_DIRECTORY=gs://${USER}-runner-maxtext-logs
+
+# Point `DATASET_PATH` to the GCS bucket where you have your training data
+export DATASET_PATH=gs://${USER}-maxtext-dataset
+
+export BUCKET_LOCATION=US
+
+# Create three GCS buckets for the demo.
+gcloud storage buckets create ${MODEL_BUCKET} --location=${BUCKET_LOCATION} || true
+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}
+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"
+else
+    # We install torch CPU because the checkpoint conversion script MaxText/llama_or_mistral_ckpt.py does not need a TPU/GPU
+    pip install torch --index-url https://download.pytorch.org/whl/cpu
+    CONVERT_CKPT_SCRIPT="llama_or_mistral_ckpt.py"
+fi
+
+JAX_PLATFORMS=cpu python MaxText/${CONVERT_CKPT_SCRIPT} \
+  --base-model-path ${CHKPT_BUCKET} \
+  --maxtext-model-path ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx} \
+  --model-size ${MODEL_VARIATION}
+echo "Writen MaxText compatible checkpoint to ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}"
+
+# We define `CONVERTED_CHECKPOINT` to refer to the checkpoint subdirectory.
+export CONVERTED_CHECKPOINT=${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}/0/items
+
+# Covert MaxText compatible checkpoints to unscanned checkpoints.
+# Note that the `CONVERTED_CHECKPOINT` is in a `scanned` format which is great for training but for efficient decoding performance we want the checkpoint in an `unscanned` format.
+export RUN_NAME=${MODEL_NAME}_unscanned_chkpt_${idx}
+
+JAX_PLATFORMS=cpu python MaxText/generate_param_only_checkpoint.py \
+MaxText/configs/base.yml \
+base_output_directory=${BASE_OUTPUT_DIRECTORY} \
+load_parameters_path=${CONVERTED_CHECKPOINT} \
+run_name=${RUN_NAME} \
+model_name=${MODEL_NAME} \
+force_unroll=true
+echo "Writen MaxText unscanned checkpoint to ${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints"
+
+# We will use the unscanned checkpoints by passing `UNSCANNED_CKPT_PATH` into `LOAD_PARAMETERS_PATH` in the following sections.
+export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints/0/items
diff --git a/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh b/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh
new file mode 100644
index 00000000..e196c7a5
--- /dev/null
+++ b/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh
@@ -0,0 +1,73 @@
+#!/bin/bash
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+set -ex
+
+idx=$(date +%Y-%m-%d-%H-%M)
+# Modify the `MODEL` and `MODEL_VARIATION` based on the model you use.
+export MODEL=$1
+export MODEL_VARIATION=$2
+export MODEL_NAME=${MODEL}-${MODEL_VARIATION}
+
+# After downloading checkpoints, copy them to GCS bucket at $CHKPT_BUCKET \
+# Please use seperate GCS paths for uploading model weights from kaggle ($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}
+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.
+export BASE_OUTPUT_DIRECTORY=gs://${USER}-runner-maxtext-logs
+
+# Point `DATASET_PATH` to the GCS bucket where you have your training data
+export DATASET_PATH=gs://${USER}-maxtext-dataset
+
+# We define `CONVERTED_CHECKPOINT` to refer to the checkpoint subdirectory.
+export CONVERTED_CHECKPOINT=${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}/0/items
+
+# Fine tune the converted model checkpoints with AQT.
+export RUN_NAME=finetune_aqt_${idx}
+
+python3 MaxText/train.py \
+MaxText/configs/base.yml \
+run_name=${RUN_NAME} \
+base_output_directory=${BASE_OUTPUT_DIRECTORY} \
+dataset_path=${DATASET_PATH} \
+steps=501 \
+enable_checkpointing=True \
+load_parameters_path=${CONVERTED_CHECKPOINT} \
+model_name=${MODEL_NAME} \
+per_device_batch_size=1 \
+quantization=int8 \
+checkpoint_period=100
+
+# We will convert the `AQT_CKPT` to unscanned checkpoint in the next step.
+export AQT_CKPT=${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints/100/items
+
+# Covert MaxText compatible AQT-fine-tuned checkpoints to unscanned checkpoints.
+# Note that the `AQT_CKPT` is in a `scanned` format which is great for training but for efficient decoding performance we want the checkpoint in an `unscanned` format.
+export RUN_NAME=${MODEL_NAME}_unscanned_chkpt_${idx}
+
+JAX_PLATFORMS=cpu python MaxText/generate_param_only_checkpoint.py \
+MaxText/configs/base.yml \
+base_output_directory=${BASE_OUTPUT_DIRECTORY} \
+load_parameters_path=${AQT_CKPT} \
+run_name=${RUN_NAME} \
+model_name=${MODEL_NAME} \
+force_unroll=true
+echo "Writen MaxText unscanned checkpoint to ${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints"
+
+# We will use the unscanned checkpoints by passing `UNSCANNED_CKPT_PATH` into `LOAD_PARAMETERS_PATH` in the following sections.
+export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints/0/items
diff --git a/jetstream/core/tools/requester.py b/jetstream/tools/requester.py
similarity index 100%
rename from jetstream/core/tools/requester.py
rename to jetstream/tools/requester.py

From 18997f3d67f1900282f584be23d5163c049acb56 Mon Sep 17 00:00:00 2001
From: Zijun Zhou <zijunzhou@google.com>
Date: Fri, 29 Mar 2024 20:43:17 +0000
Subject: [PATCH 2/4] Add description and device requirements

---
 jetstream/tools/maxtext/model_ckpt_conversion.sh        | 6 ++++++
 jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh | 6 ++++++
 2 files changed, 12 insertions(+)

diff --git a/jetstream/tools/maxtext/model_ckpt_conversion.sh b/jetstream/tools/maxtext/model_ckpt_conversion.sh
index 5e89d35c..d381d52b 100644
--- a/jetstream/tools/maxtext/model_ckpt_conversion.sh
+++ b/jetstream/tools/maxtext/model_ckpt_conversion.sh
@@ -14,6 +14,12 @@
 # limitations under the License.
 
 
+# This script will do the following:
+# - Create GCS buckets to store model artifacts for the JetStream Maxtext Inference demo.
+# - Convert the downloaded checkpoints to MaxText compatible checkpoints.
+# - Convert the MaxText compatible checkpoints to unscanned checkpoints for inference.
+# Device requirements:
+# - Both checkpoints conversion only requires CPU (with JAX CPU mode).
 set -ex
 
 idx=$(date +%Y-%m-%d-%H-%M)
diff --git a/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh b/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh
index e196c7a5..ffaa96fb 100644
--- a/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh
+++ b/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh
@@ -14,6 +14,12 @@
 # limitations under the License.
 
 
+# This script will do the following:
+# - Finetuning the MaxText compatible checkpoint (converted from original checkpoints) with AQT
+# - Convert the AQT-finetuned checkpoints to unscanned checkpoints for inference
+# TPU device requirements:
+# - For llama2-7b, it requires at least a v5e-8 TPU VM.
+# - For llama2-13B/70b, it requires a v4-128 TPU VM.
 set -ex
 
 idx=$(date +%Y-%m-%d-%H-%M)

From bf74b0da4ed063716dd532feb407314c37e1b96b Mon Sep 17 00:00:00 2001
From: Zijun Zhou <zijunzhou@google.com>
Date: Fri, 29 Mar 2024 20:46:22 +0000
Subject: [PATCH 3/4] address nit

---
 jetstream/tools/maxtext/model_ckpt_conversion.sh        | 2 +-
 jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/jetstream/tools/maxtext/model_ckpt_conversion.sh b/jetstream/tools/maxtext/model_ckpt_conversion.sh
index d381d52b..91f44edb 100644
--- a/jetstream/tools/maxtext/model_ckpt_conversion.sh
+++ b/jetstream/tools/maxtext/model_ckpt_conversion.sh
@@ -29,7 +29,7 @@ export MODEL_VARIATION=$2
 export MODEL_NAME=${MODEL}-${MODEL_VARIATION}
 
 # After downloading checkpoints, copy them to GCS bucket at $CHKPT_BUCKET \
-# Please use seperate GCS paths for uploading model weights from kaggle ($CHKPT_BUCKET) and MaxText compatible weights ($MODEL_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}
 export MODEL_BUCKET=gs://${USER}-maxtext
diff --git a/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh b/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh
index ffaa96fb..555c0a60 100644
--- a/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh
+++ b/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh
@@ -29,7 +29,7 @@ export MODEL_VARIATION=$2
 export MODEL_NAME=${MODEL}-${MODEL_VARIATION}
 
 # After downloading checkpoints, copy them to GCS bucket at $CHKPT_BUCKET \
-# Please use seperate GCS paths for uploading model weights from kaggle ($CHKPT_BUCKET) and MaxText compatible weights ($MODEL_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}
 export MODEL_BUCKET=gs://${USER}-maxtext

From 68ed383ad642e44fc38a91d2620ba04d53f7e647 Mon Sep 17 00:00:00 2001
From: Zijun Zhou <zijunzhou@google.com>
Date: Fri, 29 Mar 2024 21:22:10 +0000
Subject: [PATCH 4/4] verification

---
 .../tools/maxtext/model_ckpt_conversion.sh    | 19 +++++++++++--------
 .../maxtext/model_ckpt_finetune_with_aqt.sh   |  5 ++++-
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/jetstream/tools/maxtext/model_ckpt_conversion.sh b/jetstream/tools/maxtext/model_ckpt_conversion.sh
index 91f44edb..cffa4b1c 100644
--- a/jetstream/tools/maxtext/model_ckpt_conversion.sh
+++ b/jetstream/tools/maxtext/model_ckpt_conversion.sh
@@ -50,22 +50,25 @@ gcloud storage buckets create ${DATASET_PATH} --location=${BUCKET_LOCATION} || t
 # 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}
-gsutil -m cp -r $3 ${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"
+    JAX_PLATFORMS=cpu python MaxText/${CONVERT_CKPT_SCRIPT} \
+    --base_model_path ${CHKPT_BUCKET} \
+    --maxtext_model_path ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx} \
+    --model_size ${MODEL_VARIATION}
 else
     # We install torch CPU because the checkpoint conversion script MaxText/llama_or_mistral_ckpt.py does not need a TPU/GPU
     pip install torch --index-url https://download.pytorch.org/whl/cpu
     CONVERT_CKPT_SCRIPT="llama_or_mistral_ckpt.py"
+    JAX_PLATFORMS=cpu python MaxText/${CONVERT_CKPT_SCRIPT} \
+    --base-model-path ${CHKPT_BUCKET} \
+    --maxtext-model-path ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx} \
+    --model-size ${MODEL_VARIATION}
 fi
-
-JAX_PLATFORMS=cpu python MaxText/${CONVERT_CKPT_SCRIPT} \
-  --base-model-path ${CHKPT_BUCKET} \
-  --maxtext-model-path ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx} \
-  --model-size ${MODEL_VARIATION}
-echo "Writen MaxText compatible checkpoint to ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}"
+echo "Written MaxText compatible checkpoint to ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}"
 
 # We define `CONVERTED_CHECKPOINT` to refer to the checkpoint subdirectory.
 export CONVERTED_CHECKPOINT=${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}/0/items
@@ -81,7 +84,7 @@ load_parameters_path=${CONVERTED_CHECKPOINT} \
 run_name=${RUN_NAME} \
 model_name=${MODEL_NAME} \
 force_unroll=true
-echo "Writen MaxText unscanned checkpoint to ${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints"
+echo "Written MaxText unscanned checkpoint to ${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints"
 
 # We will use the unscanned checkpoints by passing `UNSCANNED_CKPT_PATH` into `LOAD_PARAMETERS_PATH` in the following sections.
 export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints/0/items
diff --git a/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh b/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh
index 555c0a60..359bb0b7 100644
--- a/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh
+++ b/jetstream/tools/maxtext/model_ckpt_finetune_with_aqt.sh
@@ -40,6 +40,9 @@ export BASE_OUTPUT_DIRECTORY=gs://${USER}-runner-maxtext-logs
 # Point `DATASET_PATH` to the GCS bucket where you have your training data
 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/
+
 # We define `CONVERTED_CHECKPOINT` to refer to the checkpoint subdirectory.
 export CONVERTED_CHECKPOINT=${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}/0/items
 
@@ -73,7 +76,7 @@ load_parameters_path=${AQT_CKPT} \
 run_name=${RUN_NAME} \
 model_name=${MODEL_NAME} \
 force_unroll=true
-echo "Writen MaxText unscanned checkpoint to ${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints"
+echo "Written MaxText unscanned checkpoint to ${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints"
 
 # We will use the unscanned checkpoints by passing `UNSCANNED_CKPT_PATH` into `LOAD_PARAMETERS_PATH` in the following sections.
 export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints/0/items
