TechPrufer
Get started
Google
Google logoGoogle · Feature trail

JAX Model Evaluation Metrics

12 posts · Jul 2024 – Sep 2026

What this trail covers

The JAX ecosystem continues to evolve with advancements in model building and training, specifically focusing on running Ray workloads on TPUs. This includes improvements in serving large models with Ray Serve and vLLM, efficient data feeding with Ray Data's iter_jax_batches, and distributed training with JaxTrainer. Official TPU Docker images and enhanced dashboard metrics for TPU utilization are now available. Metrax is introduced as a high-performance library for efficient and robust model ev. This post details the acceleration of spatio-temporal attention for video diffusion models on TPUs by implementing sparse attention kernels using JAX and Pallas Splash Attention. It describes the transition from theoretical sparsity to practical end-to-end speedups by optimizing tile traversal strategies, avoiding intra-tile masking, and aligning sparse masks with execution tiles to reduce latency and improve efficiency on TPU v6e chips.

All Google trails
  1. Accelerating Spatio-Temporal Attention for Video Diffusion on TPUs- Google Developers BlogThis post details the acceleration of spatio-temporal attention for video diffusion models on TPUs by implementing sparse attention kernels using JAX and Pallas Splash Attention. It describes the transition from theoretical sparsity to practical end-to-end speedups by optimizing tile traversal strategies, avoiding intra-tile masking, and aligning sparse masks with execution tiles to reduce latency and improve efficiency on TPU v6e chips.
  2. Reproducing OLMo 3 7B Pre-training in MaxText: case study of large scale training on TPUs- Google Developers BlogThis post details the reproduction of OLMo 3 7B pre-training using MaxText on Google Cloud TPUs. Key contributions include: PyTorch to JAX model conversion with logit-parity checks, verification of a data-loader bug using held-out evaluations, demonstration of checkpoint-and-resume reliability over multi-week runs and after host failures, successful mid-flight resizing of training jobs, seamless switching of TPU generations (Ironwood to v5p), performance optimizations achieving 44.5% MFU on Ironwood, and co-design of attention mechanisms for improved TPU utilization and speed. The post also outlines the necessary MaxText components built for this reproduction, including the model architecture, optimizer, data pipeline, and checkpoint conversion tools.
  3. Decoding cosmic signals with deep learning and Keras- Google Developers BlogThis post details the application of deep learning, specifically Keras, to astroparticle physics for event reconstruction. It explains how neural networks can be trained on simulated detector signals to reconstruct shower properties with unprecedented precision, overcoming limitations of traditional methods that rely on compressed quantities like charge and arrival time. The post highlights the use of Keras for this task and mentions its application to data from observatories like Pierre Auger and IceCube, contrasting it with traditional methods like fitting NKG functions and quantifying waveform spikiness for mass estimation.
  4. Enterprise-Grade Precision for Long-Context Multimodal Embedding Inference on Cloud TPU- Google Developers BlogThis post details the engineering optimizations made to the vLLM framework for serving embedding models on Google Cloud TPUs. Key contributions include: 1. A unified, hardware-safe vocabulary padding strategy for exact tensor alignment during All-Gather execution in Tensor Parallelism. 2. Attribute promotion within the unquantization pipeline to ensure compatibility with vLLM's TPU lazy-loader, eliminating initialization failures. 3. Sharding-aware pre-warming to lock JAX/XLA compilation caches prior to inference, stabilizing production pipelines. 4. A hybrid StepPool architecture with metadata migration to _CachedRequestState_ to handle long contexts and prevent state loss across steps. The post also provides a Python code example for initializing Qwen3-Embedding-8B on TPU and presents rigorous numerical parity evaluations using cosine similarity against XPU baselines, demonstrating enterprise-grade precision.
  5. HeyGen x Google Cloud: Bringing Avatar IV to TPUs- Google Developers BlogThis post details the optimization of HeyGen's Avatar IV diffusion model pipeline for Google Cloud TPUs (v6e). Key contributions include: 1. Porting the PyTorch model to JAX using torchax, allowing unmodified model code to run on TPUs. 2. Implementing custom Pallas kernels for TPU-specific optimizations, particularly for attention variants. 3. Optimizing the pipeline by "hiding the collective" operations in sequence parallelism through head grouping, reducing their critical path impact. 4. Deleting the sparse attention mask by aligning kernel block sizes to frame boundaries, eliminating partial blocks and associated overhead. 5. Unchaining the softmax serial dependency by replacing the running max with a precomputed upper bound derived from norms, allowing blocks to stream through the inner loop. These optimizations resulted in a 1.86x speedup per generated video chunk, comparable performance to 8x H100 GPUs, and improved cost efficiency.
  7. Run Ray on TPU, Part 2: Ray AI libraries- Google Developers BlogThis post details how to run Ray workloads (Serve, Data, Train) on Google TPUs, building on the foundational concepts of slices and placement groups from Part 1. It highlights the critical `topology` field in Ray Serve for multi-host model deployment, the `iter_jax_batches()` API in Ray Data for efficient JAX array batching, and the `JaxTrainer` for distributed JAX training on TPUs. It also introduces official TPU Docker images and enhanced Ray Dashboard metrics for TPU utilization. Key failure modes discussed include incorrect topology leading to unresolvable worker communication across slices.
  8. Run Ray on TPU, Part 1: The foundations- Google Developers BlogThis post introduces the official support for running Ray workloads on Google Cloud TPUs. It details how Google Kubernetes Engine (GKE) provisions and labels TPU slices, and how Ray Core utilizes these labels to reserve entire slices atomically using the `slice_placement_group()` function. The post explains the concept of TPU slices and topologies, and how this integration allows existing Ray code, including the task-and-actor model, JaxTrainer, and Ray Serve, to run on TPUs without significant code changes. It highlights the shift from experimental support to first-class accelerator status for TPUs in Ray.
  9. Systems Engineering Playbook: Optimizing Qwen 3.5-397B MoE on Ironwood (TPU7x)- Google Developers BlogThis post details the systems engineering optimizations applied to deploy and serve the Qwen 3.5-397B MoE model on Ironwood (TPU v7x). It introduces a modular, model-agnostic optimization strategy using reusable JAX/Pallas kernels and hardware-aware cost models to address the scalability challenge of optimizing new model families. Key contributions include: 1. Applying this strategy to Qwen 3.5's novel components (Gated DeltaNet, Attention Data Parallelism) to achieve ~3.1x decode-heavy and ~4.7x prefill-heavy performance uplift. 2. Addressing sharding constraints imposed by Qwen 3.5's architecture (2 KV heads in GQA, 512 experts in MoE) by co-designing a hybrid sharding scheme: 8-way Attention Batch Sharding (DP=8) combined with 8-way Expert Parallelism (EP=8) in MoE layers, overcoming limitations of traditional Tensor Parallelism. 3. Detailing the benchmark setup with asymmetric workloads (prefill-heavy, decode-heavy) across concurrency tiers and the mixed-engine orchestration on a single-host cluster with 4 Ironwood chips. 4. Integrating these optimizations into open-source serving frameworks (vLLM, SGLang) for production-ready migration.
  10. We terminated a TPU mid-training and it recovered in seconds: Introduction to elastic training with MaxText- Google Developers BlogThis post introduces and demonstrates elastic training for large model training on Cloud TPUs using JAX, MaxText, and Pathways. It details how hardware failures (e.g., a TPU worker dying) are detected as Python exceptions by Pathways, caught by an `elastic_retry` decorator from `pathways-utils`, and then handled by Orbax for safe restoration from the last viable checkpoint. This process allows the training job to recover in-place without a full restart, significantly reducing downtime. The post explains the underlying components (TPU chips, slices, host VMs, JAX, MaxText, Pathways, Orbax), the difference between elastic training and suspend-resume for preemptible instances, and provides a small-scale example to reproduce the behavior. It highlights that the primary benefit is avoiding the teardown and rescheduling of the entire workload, saving hundreds of seconds compared to a full restart.
  11. DiffusionGemma: The Developer Guide- Google Developers BlogThis post introduces DiffusionGemma, an experimental non-autoregressive text generation model built on the Gemma 4 backbone. It details DiffusionGemma's architecture, which utilizes Uniform State Diffusion and Block Autoregressive Diffusion to achieve compute-bound parallel generation, bypassing memory-bandwidth limitations and enabling faster token generation (up to 1000+ tokens/sec on H100). Key features include bidirectional attention for simultaneous text block evaluation and self-correction, and developer-friendly sizes (26B MoE, activating 3.8B parameters) for deployment within 18GB VRAM. The post showcases its application in solving Sudoku puzzles, demonstrating its ability to handle constrained problems through parallel denoising and error correction via re-noising. It also outlines serving DiffusionGemma using vLLM with a specific command-line example and provides resources for downloading weights, integration, fine-tuning, and deployment.
  12. Supercharging LLM inference on Google TPUs: Achieving 3X speedups with diffusion-style speculative decoding- Google Developers BlogThis post details the integration of DFlash, a diffusion-style speculative decoding method, into the vLLM TPU inference framework. It describes the technical challenges and solutions for porting DFlash from GPU/PyTorch to the Google TPU/JAX AI Stack, including the design of a dual-cache architecture for attention, intelligent context management with power-of-2 padding, and re-engineering the proposer to synchronize with the true accepted token count to bridge the metadata gap in TPU inference. Benchmarks show an average 3.13x speedup on TPU v5p, with peak speedups reaching nearly 6x, and a 2.29x end-to-end serving speedup compared to EAGLE-3. The 'K-Flat' breakthrough highlights that on TPUs, the cost of verifying many tokens is nearly identical to verifying few, shifting the bottleneck to draft quality.
  13. MaxText Expands Post-Training Capabilities: Introducing SFT and RL on Single-Host TPUs- Google Developers BlogIntroduces Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL) capabilities within MaxText, specifically for single-host TPU configurations. Details GRPO and GSPO RL algorithms, highlighting their memory efficiency and stability improvements. Provides code examples for launching SFT and RL runs. Mentions integration with Hugging Face datasets and checkpoints, and the use of Tunix and vLLM for optimized execution.
  14. Learn how to easily finetune FunctionGemma, a small language model, using the JAX-based Tunix library on Google TPUs for fast and cost-effective agent development.- Google Developers BlogThis post introduces Tunix, a JAX-based library for LLM post-training, and demonstrates its use for supervised fine-tuning of FunctionGemma on Google TPUs. It details the process of loading model weights, applying LoRA adapters using Qwix, defining a custom dataset for completion-only loss, and initiating training with `optax.adamw`. The post highlights efficient TPU utilization and significant accuracy gains after one epoch of training, and shows how to merge LoRA adapters for downstream deployment. It also references the extended JAX AI Stack and the use of LiteRT for on-device deployment.
  15. Building production AI on Google Cloud TPUs with JAX- Google Developers BlogThis post introduces the JAX AI Stack, a comprehensive platform for building production AI on Google Cloud TPUs. It details the core libraries (JAX, Flax, Optax, Orbax) and extended ecosystem components (XLA, Pathways, Pallas, Tokamax, Qwix, Grain, MaxText, MaxDiffusion, Tunix, vLLM serving) that enable scalable, performant, and resilient machine learning workflows. It highlights the architectural philosophy of modularity and performance, and provides examples of its use by companies like Kakao, Lightricks, and Escalante. The post also directs readers to a technical report and a new central hub (jaxstack.ai) for further exploration.
  16. Introducing Metrax: performant, efficient, and robust model evaluation metrics in JAX- Google Developers BlogIntroduces Metrax, a new high-performance library for model evaluation metrics in JAX. Metrax provides predefined metrics for classification, regression, recommendation, vision, audio, and language models, and supports parallel computation of 'at K' metrics (e.g., PrecisionAtK, RecallAtK, NDCGAtK) using JAX's `vmap` and `jit` for efficiency. It integrates with the JAX AI Stack and is already used in production at Google. The post includes code examples for computing metrics and merging batch states, and highlights community contributions.
  17. Building High-Performance Data Pipelines with Grain and ArrayRecord- Google Developers BlogIntroduces Grain, a flexible data loading library for JAX, and ArrayRecord, a highly efficient file format, to build high-performance data pipelines for large-scale machine learning. Grain ensures efficient data loading and preprocessing with multiprocessing and stateful, checkpointable iterators for reproducibility. ArrayRecord, built on Riegeli, offers efficient random access via a metadata index and parallel I/O through data chunks, enabling true global shuffling and significantly higher read throughput compared to sequential formats like TFRecord. The post also details methods for converting TFRecord datasets to ArrayRecord using `tfds build` or Apache Beam for scalable conversion.
  18. Introducing Tunix: A JAX-Native Library for LLM Post-Training- Google Developers BlogIntroduces Tunix, a new JAX-native library for LLM post-training. Tunix provides production-ready trainers for Supervised Fine-Tuning (SFT) with parameter-efficient methods (LoRA, QLoRA), Direct Preference Optimization (DPO), and Reinforcement Learning (RL) algorithms (PPO, GRPO, GSPO). It also includes a Knowledge Distillation trainer. Tunix is designed for performance on TPUs, especially with MaxText, and offers a customizable 'white-box' design. Quantitative results show a ~12% relative improvement in pass@1 answer accuracy on the GSM8K benchmark using Tunix's GRPO implementation with Gemma 2 2B-IT. Tunix is available on PyPI and is intended to support agentic AI development.
  19. Beyond backpropagation: JAX's symbolic power unlocks new frontiers in scientific computing- Google Developers BlogThis post details how researchers are leveraging JAX's composable transformations and Taylor mode automatic differentiation (AD) to solve high-order Partial Differential Equations (PDEs) efficiently. It highlights the limitations of standard deep learning frameworks (which rely on repeated backpropagation) for computing high-order derivatives and contrasts them with JAX's Taylor mode AD, which computes them in a single forward pass. The post introduces the Stochastic Taylor Derivative Estimator (STDE) method, which achieved over 1000x speed-up and 30x memory reduction for solving a 1-million-dimensional PDE. It also mentions the generalization of JAX to handle infinite-dimensional vectors (functions in Hilbert space) for calculating variational derivatives, a capability not found in other frameworks. The post emphasizes JAX's role as a foundational library for differentiable programming in scientific discovery.
  20. Stop “vibe testing” your LLMs. It's time for real evals.- Google Developers BlogThis post introduces Stax, an experimental developer tool designed to streamline the LLM evaluation lifecycle. It addresses the challenges of non-deterministic AI model outputs and the limitations of traditional unit tests by providing a framework for rigorous evaluation. Key features include data import/creation, pre-built autoraters for common criteria (coherence, factuality, conciseness), and the ability to build custom autoraters tailored to specific use cases (e.g., brand voice, PII adherence). The post emphasizes the importance of custom evals over general benchmarks for differentiating AI applications and codifying unique goals.
  21. Train a GPT2 model with JAX on TPU for free- Google Developers BlogThis post provides a hands-on tutorial for building and pretraining a GPT-2 model using JAX on Google TPUs. It details the use of the JAX ecosystem (Flax, Orbax, Optax) for model definition, including the implementation of a Transformer block with NNX. It explains how to leverage SPMD for parallel computation across TPU cores and implement data parallelism. The post also covers setting up a hardware mesh, partitioning model parameters with `nnx.with_partitioning_`, defining JIT-compiled loss and training step functions, and using an AdamW optimizer with a cosine decay schedule. It demonstrates data partitioning using `jax.device_put` and mentions the use of Weights and Biases for tracking training progress.
  22. A roboticist's journey with JAX: Finding efficiency in optimal control and simulation- Google Developers BlogThis post details how a robotics researcher is leveraging JAX for computational efficiency in optimal control and simulation. It highlights the use of JAX's `vmap` for parallelization and `scan` for trajectory simulation, achieving significant speedups. The post also discusses integrating model-based control with learning-based representations using JAX's composability, automatic differentiation (`grad`), and functional paradigm. The researcher developed `LQRax`, a JAX-native LQR solver supporting GPU acceleration, `vmap`, `scan`, and `grad`, which complements existing JAX robotics tools like Brax and MJX.
  23. Stanford’s Marin foundation model: The first fully open model developed using JAX- Google Developers BlogThis post details the Marin project's use of JAX and the custom Levanter framework to achieve end-to-end reproducibility for foundation models. It highlights specific engineering challenges and solutions: eliminating Python interpreter overhead with `@jax.jit` and fused operations, using `jax.value_and_grad` and gradient checkpointing for efficiency, and employing Pallas Splash Attention. It also covers managing large-scale parallelism with SPMD and Haliax for named tensors, building resilient compute clusters with Google Cloud TPU Multislice and Ray, and ensuring bitwise determinism through JAX's PRNGs and Tensorstore for data loading. The training process itself, including hardware migrations and adaptive data mixtures, is presented as a demonstration of the stack's robustness.
  24. Using KerasHub for easy end-to-end machine learning workflows with Hugging Face- Google Developers BlogIntroduces KerasHub's capability to load and utilize model checkpoints from Hugging Face Hub across different ML frameworks (JAX, PyTorch, TensorFlow). Details the integration with SafeTensors format and provides examples for loading Mistral, Llama 3.1, and Gemma models using specific KerasHub model classes (MistralCausalLM, Llama3CausalLM, Gemma3CausalLM). Highlights the flexibility of mixing model architectures with weights from various sources.
  25. Get started with Gemma on Ray on Vertex AI- Google Developers BlogThis post details the process of fine-tuning Gemma models on Vertex AI using Ray. It covers setting up a Google Cloud project, enabling necessary APIs, creating a Cloud Storage bucket and Artifact Registry repository, and building a custom Docker image for the Ray cluster. It then explains how to create a Ray cluster on Vertex AI using the custom image and the Vertex AI SDK. Finally, it outlines the steps for submitting a fine-tuning job using Ray Train and the Ray Jobs API, including preparing Python scripts for training and distributed execution, and monitoring the job status. The post provides concrete code examples for each step.