JetStream

The experimental JAX inference framework is designed for prototyping new LLM inference ideas. It focuses on high throughput and memory optimization using JAX.

Key performance features include:

• Paged Attention
• Chunked Prefill and Piggybacking Decode
• Collective Matmul

Key framework features include:

• Pythonic model builder
• JAX manual sharding
• Interface for different hardware supports
• On-the-fly Hugging Face (HF) model conversion and deployment

JetStream core is a subpackage providing a server and library designed for continuous batching serving. It is optimized for throughput and memory when performing LLM inference on XLA devices.

Key feature:

• Interleaved mode: Provides continuous batching to optimize inference performance. This mode uses JAX directly on a single-host TPU.

JetStream provides two primary reference engine implementations for LLM inference on XLA devices:

Jax Implementation

Optimized for Jax models. It is closely integrated with the MaxText repository.

Pytorch Implementation

Optimized for Pytorch models. It is managed via the jetstream-pytorch repository.

To improve benchmark performance, you can warm up the JetStream server before the main run. There are two supported modes:

• full: Warms up the server with all input requests.
• sampled: Warms up the server with a sampling of input requests across different bucket sizes of input lengths.

Example using full warmup mode:

python JetStream/benchmarks/benchmark_serving.py   \
--tokenizer ~/maxtext/assets/tokenizer.llama2  \
--warmup-mode full   \
--save-result   \
--save-request-outputs   \
--request-outputs-file-path outputs.json   \
--num-prompts 1000   \
--max-output-length 1024   \
--dataset openorca

When running benchmarks, you can capture the model's predictions and automatically evaluate them using ROUGE metrics.

• Use --save-request-outputs to save predictions to a file.
• Use --run-eval true to trigger automatic evaluation after the benchmark completes. If --save-result is also used, the evaluation scores will be saved.

Example:

python benchmark_serving.py \
--tokenizer /home/{username}/maxtext/assets/tokenizer \
--num-prompts 10  \
--dataset sharegpt \
--dataset-path ~/data/ShareGPT_V3_unfiltered_cleaned_split.json \
--max-output-length 1024  \
--save-request-outputs \
--run-eval true

Install JAX with TPU support using the official JAX releases. You can verify the installation by checking if JAX can detect the TPU devices.

# Install JAX for TPU
pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

# Verify TPU access
python -c "import jax; print(jax.device_count())"

You can run a local mock server environment to test JetStream components without requiring full XLA hardware. Use the following sequence of commands to start the server, send requests, and perform load testing.

1. Start the server: Runs the mock implementation.
2. Test the server: Uses the requester tool to send individual requests.
3. Load test the server: Uses the load tester to simulate high traffic.

# Start a server
python -m jetstream.core.implementations.mock.server

# Test local mock server
python -m jetstream.tools.requester

# Load test local mock server
python -m jetstream.tools.load_tester

JetStream Server uses a Prometheus client server to collect metrics from the orchestrator and engines. Metrics are disabled by default. To enable them, you must set the PROMETHEUS_PORT environment variable and pass it to the server entrypoint using the prometheus_port argument. Once enabled, metrics are accessible via HTTP requests to the server's address (e.g., 0.0.0.0:9000).

# Set the port environment variable
export PROMETHEUS_PORT=9090

# Run the MaxText server with the prometheus_port argument
python3 -m MaxText.maxengine_server \
  MaxText/configs/base.yml \
  prometheus_port=${PROMETHEUS_PORT} \
  # ... other required arguments ...

After starting the server, use the provided scripts to run performance, accuracy, and audit benchmarks.

# Run performance benchmarks
cd ~/JetStream/benchmarks/mlperf/scripts
bash ./generate_server_performance_run.sh

# Run accuracy benchmarks
cd Google/code/llama2-70b/tpu_v5e_8_jetstream_maxtext/scripts/
bash ./generate_server_accuracy_run.sh

# Run audit benchmarks
bash ./generate_server_audit_run.sh