PyTorch/XLA integration with JetStream (https://github.com/google/JetStream) for LLM inference"
☆79Dec 18, 2025Updated 3 months ago
Alternatives and similar repositories for jetstream-pytorch
Users that are interested in jetstream-pytorch are comparing it to the libraries listed below. We may earn a commission when you buy through links labeled 'Ad' on this page.
Sorting:
- JetStream is a throughput and memory optimized engine for LLM inference on XLA devices, starting with TPUs (and GPUs in future -- PRs wel…☆420Jan 5, 2026Updated 3 months ago
- Google TPU optimizations for transformers models☆136Jan 23, 2026Updated 2 months ago
- Testing framework for Deep Learning models (Tensorflow and PyTorch) on Google Cloud hardware accelerators (TPU and GPU)☆65Mar 11, 2026Updated 3 weeks ago
- ☆19Nov 5, 2025Updated 5 months ago
- torchax is a PyTorch frontend for JAX. It gives JAX the ability to author JAX programs using familiar PyTorch syntax. It also provides JA…☆204Apr 1, 2026Updated last week
- GPU virtual machines on DigitalOcean Gradient AI • AdGet to production fast with high-performance AMD and NVIDIA GPUs you can spin up in seconds. The definition of operational simplicity.
- ☆17Feb 18, 2026Updated last month
- a Jax quantization library☆109Updated this week
- Latent Large Language Models☆19Aug 24, 2024Updated last year
- ☆19Oct 6, 2023Updated 2 years ago
- MLIR-based partitioning system☆177Apr 3, 2026Updated last week
- ☆22Apr 2, 2026Updated last week
- JaxPP is a library for JAX that enables flexible MPMD pipeline parallelism for large-scale LLM training☆71Updated this week
- ☆16Apr 10, 2022Updated 4 years ago
- A simple, performant and scalable Jax LLM!☆2,201Updated this week
- Managed hosting for WordPress and PHP on Cloudways • AdManaged hosting with the flexibility to host WordPress, Magento, Laravel, or PHP apps, on multiple cloud providers. Cloudways by DigitalOcean.
- KernelBench v2: Can LLMs Write GPU Kernels? - Benchmark with Torch -> Triton (and more!) problems☆23Jul 4, 2025Updated 9 months ago
- torchprime is a reference model implementation for PyTorch on TPU.☆47Mar 3, 2026Updated last month
- Fast and easy distributed model training examples.☆12Nov 26, 2024Updated last year
- PyTorch distributed training acceleration framework☆54Aug 13, 2025Updated 7 months ago
- Example of applying CUDA graphs to LLaMA-v2☆11Aug 25, 2023Updated 2 years ago
- ☆28Jun 3, 2024Updated last year
- JAX Implementations of Descript Audio Codec and EnCodec☆35Mar 30, 2025Updated last year
- Implementation of Direct Preference Optimization☆17Jul 17, 2023Updated 2 years ago
- JAX backend for SGL☆263Apr 3, 2026Updated last week
- Managed Database hosting by DigitalOcean • AdPostgreSQL, MySQL, MongoDB, Kafka, Valkey, and OpenSearch available. Automatically scale up storage and focus on building your apps.
- ☆89Updated this week
- Minimal yet performant LLM examples in pure JAX☆246Updated this week
- ☆28Mar 17, 2026Updated 3 weeks ago
- ☆27Oct 26, 2024Updated last year
- An IR for efficiently simulating distributed ML computation.☆33Jan 13, 2024Updated 2 years ago
- Legible, Scalable, Reproducible Foundation Models with Named Tensors and Jax☆703Jan 26, 2026Updated 2 months ago
- ☆151Updated this week
- ☆17Updated this week
- Official repository for "BLEUBERI: BLEU is a surprisingly effective reward for instruction following"☆32Jun 5, 2025Updated 10 months ago
- Managed Database hosting by DigitalOcean • AdPostgreSQL, MySQL, MongoDB, Kafka, Valkey, and OpenSearch available. Automatically scale up storage and focus on building your apps.
- See details in https://github.com/pytorch/xla/blob/r1.12/torch_xla/distributed/fsdp/README.md☆25Dec 22, 2022Updated 3 years ago
- NVIDIA Resiliency Extension is a python package for framework developers and users to implement fault-tolerant features. It improves the …☆283Updated this week
- 한국어 언어 모델 학습을 위한 프로젝트(Flax, Pytorch with Huggingface Accelerate)☆32Sep 13, 2023Updated 2 years ago
- ☆34May 14, 2025Updated 10 months ago
- Training NVIDIA NeMo Megatron Large Language Model (LLM) using NeMo Framework on Google Kubernetes Engine☆16Apr 28, 2025Updated 11 months ago
- ML from scratch in Jax☆12Aug 20, 2025Updated 7 months ago
- Script for bundling Common Voice (https://commonvoice.mozilla.org/) clips by language☆11Apr 13, 2023Updated 2 years ago