JetStream

JetStream은 XLA 장치의 LLM 추론을위한 처리량 및 메모리 최적화 엔진입니다.

에 대한

JetStream은 XLA 장치에서 LLM 추론을위한 처리량 및 메모리 최적화 엔진으로 TPU (및 향후 GPU -PRS 환영)로 시작합니다.

제트 스트림 엔진 구현

현재 JAX 모델과 Pytorch 모델 용으로 사용 가능한 두 가지 참조 엔진 구현이 있습니다.

Jax

• 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 Cloud TPU VM에서 MaxText와 온라인 추론 [readme]

• V5E Cloud TPU VM에서 Pytorch와의 온라인 추론 [readme]

• Jetstream과 함께 GKE의 TPU를 사용하여 Gemma를 제공하십시오

• 벤치 마크 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