AI Hypercomputer software updates: Faster training and inference, a new resource hub, and more

The potential of AI has never been greater, and infrastructure plays a foundational role in driving it forward. AI Hypercomputer is our supercomputing architecture based on performance-optimized hardware, open software, and flexible consumption models. Together, these offer exceptional performance and efficiency, resiliency at scale, and give you the flexibility to choose offerings at each layer to suit your needs.

Today, we’re announcing major updates to the AI Hypercomputer software layer for training and inference performance, improved resiliency at scale, as well as a centralized hub for AI Hypercomputer resources.

AI Hypercomputer resources on Github

AI Hypercomputer’s open software layer not only supports leading ML Frameworks and orchestration options, but also provides workload optimizations and reference implementations to improve the time-to-value for your specific use case. To make the innovations in our open software stack easily accessible to developers and practitioners, we are introducing AI Hypercomputer GitHub organization, a central place where you can discover reference implementations such as MaxText, MaxDiffusion, orchestration tools such as xpk (the Accelerated Processing Kit for cluster creation and workload management), and performance recipes for GPUs on Google Cloud. We’ll continue to add to this list and adapt these resources to a rapidly evolving landscape, and we invite you to contribute with us. 

MaxText now supports A3 Mega VMs

MaxText is a high-performance, highly scalable, open-source reference implementation for large language models (LLMs). You can now use performance-optimized LLM training examples for A3 Mega VMs, which are powered by NVIDIA H100 Tensor Core GPUs and offer a 2X improvement in GPU-to-GPU network bandwidth over A3 VMs. We worked closely with NVIDIA to optimize JAX and XLA, enabling the overlap of collective communication and computation on GPUs. Additionally, we added optimized model configurations and example scripts for GPUs with XLA flags enabled.

MaxText with A3 Mega VMs can deliver near linear scaling of training performance as you scale the number of VMs in the cluster, as demonstrated below by Llama2-70b pre-training.

Furthermore, you can use FP8 mixed-precision training on A3 Mega VMs to achieve additional acceleration and hardware utilization. We added FP8 support in MaxText via Accurate Quantized Training (AQT), the same quantization library that also powers INT8 mixed-precision training on Cloud TPUs. 

Our benchmarks on dense models indicate FP8 training with AQT can deliver up to 55% improvement in effective model flops utilization (EMFU) compared to bf16. You can find recipes and optimal training examples for A3 Mega here.

Reference implementations and kernels for MoEs

For most mixture of experts (MoE) use cases, it’s useful to have consistent resource utilization of a limited number of experts. However, for certain applications, the ability to use more experts to develop richer responses is more important. To give you this flexibility, we’ve now expanded MaxText to include both “capped” and “no-cap” MoE implementations so you can choose the best implementation for your model architecture. Capped MoE models offer predictable performance, while no-cap models dynamically allocate resources for optimal efficiency.

To further accelerate MoE training, we’ve open-sourced Pallas kernels, which are optimized for block-sparse matrix multiplication on Cloud TPUs (Pallas is an extension to JAX to allow fine-grained control over code generated for XLA devices such as GPUs and TPUs; Block-sparse matmul is currently available only for TPUs). These kernels can be used with both PyTorch and JAX, providing high-performance building blocks for training your MoE models.

Our benchmarks with no-cap MoE model (Mixtral-8x7b) indicate near linear scaling with fixed per device batch size (Figure-2). We also observed close to linear scaling when we increased the number of experts (Figure-3)  in the base configuration with the number of accelerators, indicative of performance on models with higher sparsity as well.

Monitoring large-scale training

Having large clusters of accelerators that are expected to work together as a unit on a training task can complicate MLOps. You may have questions such as “Did host transfer latencies spike for a reason?” or “Why is this one device in a segfault?” Yet, monitoring large-scale training jobs with the right metrics is imperative to maximizing your resource utilization and improving overall ML Goodput. To simplify this crucial part of your MLOps charter, we’ve introduced a reference monitoring recipe. This recipe helps you create a Cloud Monitoring dashboard within your Google Cloud project that shows useful statistical measures such as average or max CPU utilization and helps identify outliers in the setup so you can take corrective actions.

SparseCore on Cloud TPU v5p is now GA

Recommender models, as well as models that rely on embeddings, require high-performance random memory access to use those embeddings. SparseCore, the TPU’s hardware accelerator for embeddings, enables you to build more powerful and efficient recommendation systems. Each Cloud TPU v5p chip has four dedicated SparseCores delivering up to 2.5X performance improvement for DLRM-V2 over its predecessor.

Improving LLM inference performance

Finally, to improve LLM inference performance, we introduced KV cache quantization and ragged attention kernels in JetStream, an open-source throughput-and-memory-optimized engine for LLM inference. Together, these enhancements improve inference performance by up to 2X on Cloud TPU v5e.

Empowering your AI journey

From pushing the boundaries of model training and inference, to enhancing accessibility through a central resource repository, each component of the AI Hypercomputer is a building block for the next generation of AI. We envision a future where AI practitioners can seamlessly scale from concept to production, unburdened by infrastructure limitations. 

Explore the latest AI Hypercomputer resources, including the optimized recipes, accelerated processing kit, reference implementations and more.