JetStream

JetStreamは、XLAデバイス上のLLM推論用のスループットおよびメモリ最適化エンジンです。

について

JetStreamは、XLAデバイスのLLM推論用のスループットとメモリの最適化されたエンジンであり、TPU（および将来のGPU -PRS歓迎）から始まります。

JetStreamエンジンの実装

現在、利用可能な2つの参照エンジンの実装があります。1つはJAXモデル用、もう1つはPytorchモデル用です。

ジャックス

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

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

Pytorch

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

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

ドキュメント

• V5EクラウドTPU VM [README]のMAXTEXTを使用したオンライン推論

• V5EクラウドTPU VM [ReadMe]のPytorchを使用したオンライン推論

• ジェットストリームでGKEでTPUを使用してジェマを提供します

• ベンチマークジェットストリームサーバー

• JetStream Serverの観察可能性

• 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