Google Cloud Next '24 Blog Series: Scaling Generative AI - How Google’s JAX and OpenXLA Are Powering the Next Wave of AI Models

At Google Cloud Next '24, Robert Crow, Product Manager at Google’s Frameworks team, delivered a deep dive into the ML frameworks driving generative AI development. The talk unpacked how JAX, OpenXLA, and emerging AI infrastructure tools are making massive-scale model training faster, cheaper, and more efficient.

With AI models like GPT-4, Gemini, and PaLM 2 pushing computational demands to unprecedented levels, the need for performance-optimized, scalable ML frameworks has never been more critical. Let’s break down the key insights from this session.

The Compute Explosion: Why Generative AI Needs New Frameworks

The Transformer revolution of 2017 set the stage for exponential growth in AI model complexity. Today, we’re looking at models with trillions of parameters, requiring tens of billions of petaFLOPs and hundreds of millions of dollars per training run.

Mustafa Suleyman, co-founder of DeepMind, predicts a 10x increase in AI compute demand within 12 months—and a 100x increase within three years. This means efficiency isn’t just important—it’s a survival factor for AI research and deployment.

JAX: The Backbone of Google's AI Research

JAX is the foundational ML framework behind Google’s AI models, including Gemini, AlphaFold, and PaLM 2. It provides automatic differentiation, vectorization, and just-in-time (JIT) compilation, enabling researchers to train and deploy massive models with unprecedented speed and efficiency.

Why JAX Matters:

• Automatic Differentiation & XLA Acceleration: Optimized for TPUs, GPUs, and CPUs, delivering near supercomputer-level performance.

• Composable & Modular Design: Works seamlessly with Flax, Optax, and Orbax for easier model development.

• Scalability: Powers training jobs using 50,000+ TPUs, ensuring linear scaling efficiency.

What the Experts Say:

Crow highlighted how JAX is the go-to framework for AI researchers, offering a blazingly fast and scalable alternative to PyTorch and TensorFlow. He demonstrated JIT compilation’s speed-up, reducing model execution times from 2 milliseconds to 150 microseconds—a 10x improvement with a single function call.

OpenXLA: The Secret Weapon for AI Acceleration

While JAX enables efficient ML model development, OpenXLA is the compiler stack that supercharges execution. Built as an open-source collaboration with NVIDIA, PyTorch, Hugging Face, and Intel, OpenXLA provides an optimized ML runtime for any accelerator.

Why OpenXLA is Game-Changing:

• Supercharges ML Workloads: Reduces compute overhead by intelligently optimizing linear algebra operations.

• Multi-Accelerator Support: Runs on TPUs, GPUs, and custom AI hardware.

• Automatic Kernel Optimizations: Delivers hardware-specific optimizations without manual tuning.

Crow emphasized OpenXLA’s ability to squeeze out peak performance from AI accelerators, providing an ML runtime that rivals high-performance supercomputing environments.

Scaling AI Models: JAX for Training and Deployment

Training AI at scale requires parallelism strategies that distribute workloads efficiently. JAX provides multiple approaches, ensuring seamless model scaling.

Key Strategies:

• Data Parallelism: Duplicates the model across GPUs/TPUs and shards data across accelerators.

• Model Parallelism: Splits model weights across multiple devices, ideal for training trillion-parameter models.

• Pipeline Parallelism: Executes different parts of a model simultaneously across accelerators, improving training speed.

Performance Gains:

JAX enables training jobs with 50,944 TPUs, delivering near-ideal scaling efficiency up to 8,000+ chips. The framework ensures AI teams spend less time tuning hardware and more time iterating on models.

The Rise of AI-Specific Hardware Optimization

Beyond software optimizations, JAX enables advanced hardware tuning, making AI models even faster and more cost-efficient.

Key Hardware Innovations:

• Triton for Custom Kernel Optimization: Allows AI engineers to write custom GPU kernels, maximizing hardware performance.

• Palace for JAX-Optimized Kernels: Provides native JAX syntax for writing high-performance GPU/TPU kernels.

• Quantization (AQT): Reduces memory and compute costs by training models at INT8 precision while maintaining accuracy.

Beyond Training: Deploying AI with JAX and JetStream

Once models are trained, they need efficient serving and inference solutions. Enter JetStream, a new AI inference engine designed for throughput-optimized LLM deployment.

JetStream Features:

• Continuous Batching & Weight Optimization: Reduces inference latency for real-time AI applications.

• Quantization & KV Cache Optimizations: Improves model efficiency without sacrificing accuracy.

• Scalable Across Cloud and On-Prem: Works seamlessly with Google Cloud, Hugging Face, and local AI clusters.

Crow emphasized how JetStream leverages JAX for high-speed inference, enabling LLMs to serve responses in milliseconds rather than seconds.

Final Thoughts: The Future of AI Runs on JAX & OpenXLA

As AI models grow larger and more complex, JAX and OpenXLA are becoming essential tools for AI researchers and enterprises. From Gemini to AlphaFold, these frameworks are proving to be the backbone of next-gen AI development.

The takeaway? The future of AI will be shaped by the frameworks that make large-scale compute efficient, scalable, and accessible.