Skip to content

Commit 9469054

Browse files
authored
Update maxtext user guide (#56)
1 parent 0cbe8ba commit 9469054

2 files changed

Lines changed: 7 additions & 4 deletions

File tree

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ Currently, there are two reference engine implementations available -- one for J
1616
### Jax
1717

1818
- Git: https://github.com/google/maxtext
19-
- README: https://github.com/google/JetStream/blob/main/jetstream/max_text/README.md
19+
- README: https://github.com/google/JetStream/blob/main/docs/online-inference-with-maxtext-engine.md
2020

2121
### Pytorch
2222

jetstream/max_text/README.md renamed to docs/online-inference-with-maxtext-engine.md

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,7 @@ export ICI_AUTOREGRESSIVE_PARALLELISM=-1
126126
export ICI_TENSOR_PARALLELISM=1
127127
export SCAN_LAYERS=false
128128
export WEIGHT_DTYPE=bfloat16
129-
export PER_DEVICE_BATCH_SIZE=6
129+
export PER_DEVICE_BATCH_SIZE=4
130130
```
131131

132132
#### Create Llama2-13b environment variables for server flags
@@ -191,7 +191,10 @@ Note: these flags are from [MaxText config](https://github.com/google/maxtext/bl
191191

192192
```bash
193193
cd ~
194-
python JetStream/jetstream/tools/requester.py
194+
# For Gemma model
195+
python JetStream/jetstream/tools/requester.py --tokenizer maxtext/assets/tokenizer.gemma
196+
# For Llama2 model
197+
python JetStream/jetstream/tools/requester.py --tokenizer maxtext/assets/tokenizer.llama2
195198
```
196199

197200
The output will be similar to the following:
@@ -250,7 +253,7 @@ wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/r
250253
# run benchmark with the downloaded dataset and the tokenizer in maxtext
251254
# You can control the qps by setting `--request-rate`, the default value is inf.
252255
python JetStream/benchmarks/benchmark_serving.py \
253-
--tokenizer /home/$USER/maxtext/assets/tokenizer.gemma \
256+
--tokenizer maxtext/assets/tokenizer.gemma \
254257
--num-prompts 1000 \
255258
--dataset sharegpt \
256259
--dataset-path ~/ShareGPT_V3_unfiltered_cleaned_split.json \

0 commit comments

Comments
 (0)