AI Hypercomputer Architecture

The AI Hypercomputer system has three key layers: 

• Performance-optimized hardware
• open software
• flexible consumption

AI Hypercomputer is an integrated system designed to efficiently scale and deploy AI applications.

Layer 1: High-performance AI hardware (TPUs/GPUs), fast networking, and optimized storage for demanding AI.

Layer 2: Open software (PyTorch, GKE, Kueue) simplifies AI workflows and boosts productivity.

Layer 3: Flexible consumption models (on-demand, spot, CUDs, reservations, DWS) for various AI workloads.

Layer 1: Performance-optimized, purpose-built infrastructure

This foundational layer provides the raw computational power required for demanding AI tasks. It encompasses accelerator resources, high-speed networking, and efficient storage infrastructure, all meticulously optimized for performance at scale.

At the heart of the infrastructure are the processing units designed for AI acceleration. This includes state-of-the-art hardware like Google Cloud TPUs (Tensor Processing Units) and powerful Cloud GPUs (graphics processing units), complemented by CPUs where needed.

These components provide the core processing capabilities for complex model training and inference.

Handling the massive datasets common in AI requires specialized storage solutions. AI Hypercomputer supports diverse options, including block, file, and object storage. For example, Hyperdisk ML provides block storage optimized specifically for AI inference/serving, significantly reducing model load times and improving cost efficiency.

Additionally, caching capabilities in Cloud Storage FUSE and Parallelstore enhance throughput and reduce latency during both training and inference. This optimized storage is critical across the entire AI data pipeline—from preparation and training through to inference, delivery, and data protection. A fully managed, high-performance parallel file system like Google Cloud Managed Lustre offers multi-petabyte scalability and significant throughput. They are optimized for AI and HPC applications.

Connecting these powerful accelerators efficiently is paramount. Google Cloud employs advanced network technologies, such as the Jupiter network fabric and optical circuit switching (OCS), to create highly scalable data center networks. This infrastructure delivers petabit-scale bandwidth, crucial for large distributed training tasks. For example, A3 Mega VMs leverage this Jupiter fabric.

Furthermore, specialized VPC infrastructure optimized for direct GPU-to-GPU connectivity offers up to 3.2 Tbps capacity per VM and supports RDMA for ultra-low latency data transfer, vital for Generative AI and HPC workloads.

Layer 2: Open software

Building upon its robust hardware, AI Hypercomputer utilizes an open software stack. This layer features optimized versions of popular frameworks, libraries, compilers, reference projects, orchestration tools and operating systems, all working together seamlessly. The core goal is to simplify access to the powerful underlying TPUs and GPUs, significantly enhancing performance, boosting productivity, and improving the overall usability for AI tasks.

Orchestration

To deploy and manage a large number of accelerators as a single unit, you can also use 

Cluster Director for Google Kubernetes Engine, or 

Cluster Director for Slurm, or 

directly through Compute Engine APIs. 

GKE is an excellent choice to handle orchestration needs, chosen for its foundation in open-source Kubernetes, which offers portability, customizability, performance, and scalability.

GKE within AI Hypercomputer includes specialized features tuned for AI. For example, gen-AI-aware scaling, container preloading for faster startups, optimized data access via tools like Google Cloud Storage Fuse with caching and Parallelstore, and efficient job management using the Kueue queuing system

ML Frameworks:

For model development, the stack provides optimized support for key open-source machine learning frameworks like 

Jax, PyTorch, and Keras, leveraging the powerful XLA compiler for peak performance on Google's hardware. To streamline workflows, it includes helpful libraries like Optimum TPU and PyTorch/XLA that simplify running models efficiently, alongside specialized tools like JetStream for LLM inference, MaxText as a reference for dense LLMs, and MaxDiffusion for diffusion models.

Support for distributed frameworks like Ray on GKE and cluster blueprints for repeatable deployments further ensures a comprehensive environment for building and scaling cutting-edge AI applications.

Layer 3: Flexible consumption

Google Cloud provides a versatile range of consumption models for its AI Hypercomputer system, designed to optimize cost, flexibility, and resource availability for diverse AI/ML workloads.