JetStream es un motor optimizado para el rendimiento y la memoria para la inferencia LLM en los dispositivos XLA, comenzando con TPUS (y GPU en el futuro - PRS bienvenido).
Actualmente, hay dos implementaciones de motores de referencia disponibles: una para modelos JAX y otra para modelos Pytorch.
Git: https://github.com/google/mextext
Readme: https://github.com/google/jetstream/blob/main/docs/online-inference-with-maxtext-ingine.md
Git: https://github.com/google/jetstream-pytorch
Readme: https://github.com/google/jetstream-pytorch/blob/main/readme.md
Inferencia en línea con Maxtext en V5E Cloud TPU VM [ReadMe]
Inferencia en línea con Pytorch en V5E Cloud TPU VM [ReadMe]
Sirva Gemma usando TPUS en GKE con Jetstream
Benchmark Jetstream Server
Observabilidad en el servidor JetStream
Perfil en el servidor JetStream
Configuración local independiente de JetStream
make install-deps
Use los siguientes comandos para ejecutar un servidor localmente:
# 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