Skip to content

llama2 recipe: use py module, correctnesss fixes #75

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

llama2 recipe: use py module, correctnesss fixes #75

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions inference/trillium/JetStream-Maxtext/Llama2-7B/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ gsutil cp -r llama/llama-2-7b/* ${CHKPT_BUCKET}

# Checkpoint conversion
cd maxtext
bash ../JetStream/jetstream/tools/maxtext/model_ckpt_conversion.sh llama2 7b ${CHKPT_BUCKET} ${MAXTEXT_BUCKET_SCANNED} ${MAXTEXT_BUCKET_UNSCANNED}
bash ../JetStream/jetstream/tools/maxtext/model_ckpt_conversion.sh llama2 7b ${CHKPT_BUCKET} ${MAXTEXT_BUCKET_SCANNED} ${MAXTEXT_BUCKET_UNSCANNED} False

# The path to the unscanned checkpoint should be set by the script, but set it explicitly if it hasn't
# For example export UNSCANNED_CKPT_PATH=gs://${MAXTEXT_BUCKET_UNSCANNED}/llama2-7b_unscanned_chkpt_2024-08-23-23-17/checkpoints/0/items
Expand All @@ -71,7 +71,7 @@ export WEIGHT_DTYPE=bfloat16
export PER_DEVICE_BATCH_SIZE=11

cd ~/maxtext
python MaxText/maxengine_server.py \
python -m MaxText.maxengine_server \
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be good to checkout a specific MaxText commit so that future MaxText changes don't break the recipes. Example here

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

guess we want to run git checkout tpu-recipes-v0.1.4 before python -m MaxText....

MaxText/configs/base.yml \
tokenizer_path=${TOKENIZER_PATH} \
load_parameters_path=${LOAD_PARAMETERS_PATH} \
Expand All @@ -89,6 +89,7 @@ python MaxText/maxengine_server.py \
In terminal tab 2, run the benchmark:
```bash
source venv-maxtext/bin/activate
git lfs pull # To pull down the test dataset

python JetStream/benchmarks/benchmark_serving.py \
--tokenizer ~/maxtext/assets/tokenizer.llama2 \
Expand All @@ -101,7 +102,7 @@ python JetStream/benchmarks/benchmark_serving.py \
--dataset openorca
```

After the benchmark finishes, you should see something like
After the benchmark finishes, you should see something like
```bash
Successful requests: 995
Benchmark duration: 305.366344 s
Expand Down
7 changes: 4 additions & 3 deletions inference/v5e/JetStream-Maxtext/Llama2-7B/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ gsutil cp -r llama/llama-2-7b ${CHKPT_BUCKET}

# Checkpoint conversion
cd maxtext
bash ../JetStream/jetstream/tools/maxtext/model_ckpt_conversion.sh llama2 7b ${CHKPT_BUCKET} ${MAXTEXT_BUCKET_SCANNED} ${MAXTEXT_BUCKET_UNSCANNED}
bash ../JetStream/jetstream/tools/maxtext/model_ckpt_conversion.sh llama2 7b ${CHKPT_BUCKET} ${MAXTEXT_BUCKET_SCANNED} ${MAXTEXT_BUCKET_UNSCANNED} False

# The path to the unscanned checkpoint should be set by the script, but set it explicitly if it hasn't
# For example export UNSCANNED_CKPT_PATH=gs://${MAXTEXT_BUCKET_UNSCANNED}/llama2-7b_unscanned_chkpt_2024-08-23-23-17/checkpoints/0/items
Expand All @@ -71,7 +71,7 @@ export WEIGHT_DTYPE=bfloat16
export PER_DEVICE_BATCH_SIZE=11

cd ~/maxtext
python MaxText/maxengine_server.py \
python -m MaxText.maxengine_server \
MaxText/configs/base.yml \
tokenizer_path=${TOKENIZER_PATH} \
load_parameters_path=${LOAD_PARAMETERS_PATH} \
Expand All @@ -89,6 +89,7 @@ python MaxText/maxengine_server.py \
In terminal tab 2, run the benchmark:
```bash
source venv-maxtext/bin/activate
git lfs pull # To pull down the test dataset

python JetStream/benchmarks/benchmark_serving.py \
--tokenizer ~/maxtext/assets/tokenizer.llama2 \
Expand All @@ -101,7 +102,7 @@ python JetStream/benchmarks/benchmark_serving.py \
--dataset openorca
```

After the benchmark finishes, you should see something like
After the benchmark finishes, you should see something like
```bash
Successful requests: 995
Benchmark duration: 305.366344 s
Expand Down