JetStream

Jetstream是XLA設備上LLM推斷的吞吐量和內存優化引擎。

關於

JetStream是XLA設備上LLM推斷的吞吐量和內存優化引擎，從TPU開始（將來GPU和GPU-歡迎PRS）。

Jetstream Engine實施

當前，有兩個可用的參考引擎實現 - 一種用於JAX模型，另一種用於Pytorch型號。

JAX

• git：https：//github.com/google/maxtext

• README：https：//github.com/google/jetstream/blob/main/main/docs/online-inline-inline-with-maxtext-engine.md

Pytorch

• git：https：//github.com/google/jetstream-pytorch

• readme：https：//github.com/google/jetstream-pytorch/blob/main/main/readme.md

文件

• 在V5E Cloud TPU VM上使用Maxtext在線推斷[README]

• 在V5E Cloud TPU VM上與Pytorch在線推斷[readme]

• 使用tpus在gke上使用jetstream使用tpu

• 基準測試服務器

• Jetstream服務器中的可觀察性

• 在Jetstream服務器中進行分析

• Jetstream獨立本地設置

Jetstream獨立本地設置

入門

設定

make install-deps

運行本地服務器和測試

使用以下命令在本地運行服務器：

# 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

測試核心模塊

# Test JetStream core orchestrator
python -m unittest -v jetstream.tests.core.test_orchestrator

# Test JetStream core server library
python -m unittest -v jetstream.tests.core.test_server

# Test mock JetStream engine implementation
python -m unittest -v jetstream.tests.engine.test_mock_engine

# Test mock JetStream token utils
python -m unittest -v jetstream.tests.engine.test_token_utils
python -m unittest -v jetstream.tests.engine.test_utils