Hardware-Aware Training and Inference for Large-Scale AI

Hardware-Aware Training and Inference for Large-Scale AI

Speakers: Animashree Anandkumar, Professor, Caltech

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

In an era where large-scale AI models are continually pushing the boundaries of computational resources, Professor Animashree Anandkumar's talk at MLSys 2025 presents a compelling vision for the future of machine learning systems: co-designing algorithms and hardware. Moving beyond the traditional paradigm of optimizing systems for fixed machine learning architectures, Anandkumar advocates for a fundamental reimagining of ML algorithms themselves to better suit the constraints and opportunities presented by modern hardware. This approach is not merely about incremental improvements through pruning or quantization but about architecting algorithms from first principles with hardware efficiency in mind.

The talk delves into innovative techniques that drastically reduce memory and bandwidth requirements during both training and inference, often yielding surprising benefits in accuracy and robustness. From extreme quantization methods like Logarithmic Number Systems (LNS) and SignSGD to sophisticated low-rank compression techniques such as Galore, Anandkumar demonstrates how principled algorithmic changes can unlock unprecedented efficiency. Furthermore, she extends this hardware-aware philosophy to the emerging field of Physical AI, highlighting the unique challenges posed by multi-dimensional, high-resolution physical data and showcasing how Neural Operators can revolutionize scientific discovery and engineering design by offering orders of magnitude speedups over traditional simulations. This talk is a call to action for the ML and systems communities to collaborate more deeply, fostering a new era of innovation where algorithmic elegance meets hardware practicality.

Background

▶ Watch: Co-designing ML algorithms and systems for efficiency (2:00)

The rapid growth of machine learning, particularly with the advent of large language models (LLMs) and diffusion models, has led to an insatiable demand for computational resources. Historically, the machine learning community often operated under the assumption that hardware improvements would keep pace with algorithmic advancements. However, as Professor Anandkumar highlights, this era of "abundance" has shifted to one of "scarcity." Moore's Law, which predicted the exponential growth in transistor density, has slowed significantly, and more critically, a "memory law" is also facing physical limitations, particularly concerning High Bandwidth Memory (HBM) capacity and throughput. Building on her insights from the advisory board of SK Hynix, a leader in HBM, Anandkumar underscores that simply relying on hardware to get "better, faster, more capacity, cheaper" is no longer a trivial expectation.

This bottleneck manifests primarily as memory and bandwidth constraints. Modern LLMs and increasingly complex models for physical AI require massive amounts of memory for weights, activations, and optimizer states, and substantial bandwidth for data transfer, especially in distributed settings. Traditional machine learning algorithms, such as stochastic gradient descent (SGD), were conceived decades ago without any consideration for hardware limitations, focusing instead on theoretical convergence and generalization properties. While efforts like pruning and quantization have provided some relief, they are often applied as post-hoc optimizations, leading to potential accuracy degradation or stability issues. The core problem, therefore, is a disconnect: algorithms designed for an idealized computational environment are being run on increasingly constrained physical hardware. This necessitates a fundamental re-evaluation, pushing researchers to co-design ML algorithms and systems from the ground up. The emergent field of Physical AI, dealing with 3D/4D data, multi-physics simulations, and extremely high-resolution inputs (terabytes per data sample), further amplifies these memory and bandwidth challenges, making hardware-aware algorithmic design not just an advantage, but a necessity.

Key Findings

▶ Watch: Exploring applications beyond language models: physical AI (4:40)

Professor Anandkumar's talk presents several key findings that collectively argue for a paradigm shift towards hardware-aware ML algorithm design:

• Algorithmic Reimagination for Hardware Efficiency: The central finding is that by fundamentally rethinking ML algorithms, significant hardware resource reductions (memory, bandwidth, energy) can be achieved without sacrificing accuracy, and in some cases, even improving it. This challenges the conventional wisdom that "lossy" compression or low-precision methods inherently degrade performance.
• Extreme Quantization is Viable for Training: Techniques like Logarithmic Number Systems (LNS) for weight/activation representation and SignSGD for gradient communication demonstrate that quantizing data to extremely low bit-widths (e.g., single bit for gradients, exponent-only for numbers) can yield massive energy (90% reduction for LNS) and bandwidth (32x reduction for SignSGD) savings while maintaining, or even exceeding, full-precision accuracy. This is attributed to designing optimization algorithms (e.g., multiplicative updates for LNS, mirror descent for SignSGD) that are inherently compatible with these low-precision representations.
• Low-Rank Optimization for Pre-training: The Galore (Gradient Low-Rank Projection) method proves that compressing gradients and optimizer states into a low-rank subspace can drastically reduce memory requirements (80% reduction in Galore 2) during large-scale pre-training. Unlike fine-tuning specific tasks (e.g., LoRA), Galore preserves the full capacity of the model, enabling the training of large foundation models on significantly less memory-intensive hardware (e.g., a 7B LLaMA model on a single consumer-grade RTX 4090).
• Efficient Inference for Long Contexts: Novel strategies like HeadInfer (head-by-head KV cache offloading) and MLP sequence breaking address memory bottlenecks in long-context LLMs, making them runnable on resource-constrained devices without performance degradation. Furthermore, pruning redundant tokens in Chain-of-Thought reasoning can save memory and surprisingly improve accuracy by reducing "confusion."
• Physical AI Revolutionized by Neural Operators: Neural Operators offer a powerful framework for tackling complex scientific and engineering problems. They achieve orders of magnitude speedups (e.g., 10,000x for weather forecasting, 1,000,000x for fusion plasma simulation) and superior accuracy over traditional physics-based simulations. This is possible by learning directly from data, leveraging resolution-invariant architectures (e.g., using Fourier transforms), and operating effectively at lower precisions (FP16, FP8) where HPC traditionally demands FP64.
• Physics-Informed AI for Real-World Impact: The success in areas like hurricane prediction, fusion reactor control, and medical device design demonstrates that AI models, when designed to be physics-informed and verifiable, can move beyond mere "hallucination of ideas" to provide physically valid and impactful solutions, drastically reducing R&D costs and accelerating scientific discovery.

Technical Deep Dive

▶ Watch: Physical AI's challenges: massive data, memory, bandwidth (5:40)

The core thesis of Professor Anandkumar's talk is the necessity of co-designing machine learning algorithms and systems. This involves a radical departure from merely optimizing existing ML methods on the system side. Instead, the focus shifts to reimagining the algorithms themselves to be inherently hardware-efficient, often by embracing what might seem like "lossy" operations from a traditional systems perspective. The speaker emphasizes that "noise can be a friend of learning," suggesting that carefully introduced "lossiness" can lead to surprising gains.

Inference-Time Memory Optimizations for LLMs

For large language models, particularly those with long contexts or complex reasoning chains, memory and bandwidth bottlenecks are prominent at inference time.

• HeadInfer (Head-by-head KV Cache Offloading): The KV cache for long contexts can quickly exhaust GPU memory. HeadInfer addresses this by intelligently offloading parts of the KV cache to the CPU. Instead of naively offloading entire layers, it performs a more granular, head-by-head offloading within the attention layers. By inferring the importance of different attention heads, the system can decide which heads to offload, significantly reducing GPU memory requirements. This allows long-context inference on consumer-grade GPUs like the NVIDIA RTX 4090, democratizing access to large models.
• MLP Sequence Breaking: In many LLM architectures, the Multi-Layer Perceptron (MLP) layers can create peak memory bottlenecks due to large intermediate activations. The technique of breaking down the processing of these MLP layers into "mini-sequences" helps manage these peaks. This simple yet effective idea allows for fitting much longer sequences into memory while maintaining the same throughput, as the bottleneck was primarily memory capacity, not computational speed. Both HeadInfer and MLP sequence breaking are lossless methods, meaning they don't compromise the model's accuracy.
• Chain-of-Thought (CoT) Pruning: Reasoning models using CoT generate a sequence of intermediate thoughts, often leading to a rapidly expanding and redundant KV cache. This method focuses on identifying and pruning redundant or unimportant tokens within the reasoning chain. By keeping only relevant information, memory usage is significantly reduced. Surprisingly, this pruning can also lead to slightly improved accuracy (e.g., 105% relative accuracy), as it helps the model avoid confusion from irrelevant or repetitive thoughts.

Training-Time Efficiency through Low Precision and Compression

Training large models is even more resource-intensive than inference, making hardware-aware optimizations crucial.

• Low Precision Training with Logarithmic Number System (LNS): While hardware increasingly supports lower precision floating-point formats like FP8 and FP4, the talk proposes an even more radical approach: training directly in a Logarithmic Number System (LNS). Inspired by biological brains, LNS discards the mantissa and retains only the exponent bits. This extreme quantization (e.g., 8-bit LNS vs. FP32) offers a vast dynamic range with very few bits, making it efficient for hardware (e.g., 90% energy reduction). However, standard SGD or Adam optimizers fail with LNS. The key is to use multiplicative updates instead of additive ones, effectively operating directly on the exponents. This principled approach allows full training in LNS without accuracy loss, demonstrating that a deep understanding of the number system can unlock significant hardware gains.
• SignSGD (Single-Bit Gradient Communication): SignSGD (and its variant, Lion optimizer) represents another form of extreme quantization applied to gradients. Instead of communicating full-precision (e.g., FP32) gradients in distributed training, only the sign of the gradient (a single bit) is transmitted. While seemingly an immense loss of information, theoretically, SignSGD can be viewed as mirror descent on a non-Euclidean manifold, guaranteeing convergence to local optima similar to SGD. Practically, it proves more robust to heavy-tail noise and outliers. In a distributed setting, instead of averaging full-precision gradients, workers communicate their single-bit signs, and a majority vote determines the global sign. This reduces communication bandwidth by 32 times (from 32 bits to 1 bit per gradient component), making it highly beneficial in bandwidth-constrained environments like edge devices or low-bandwidth networks.
• Galore (Gradient Low-Rank Projection): For pre-training large models, the memory consumption of gradients and optimizer states (e.g., in Adam) is a major bottleneck. Unlike LoRA, which compresses weights for fine-tuning specific tasks, Galore focuses on compressing gradients and optimizer states during pre-training, maintaining the model's full learning capacity. The intuition is that gradients evolve slowly, meaning they can be effectively represented in a low-rank subspace that slowly adapts over time. Galore projects both gradients and optimizer states into such a low-rank subspace. The original work demonstrated the ability to pre-train a 7B LLaMA model on a single RTX 4090. Galore 2 further scales this, achieving an 80% reduction in memory for optimizer and gradient states during large-scale pre-training (500 billion tokens), matching the accuracy of 8-bit Adam baselines.

Physical AI and Neural Operators

The talk extends the co-design philosophy to Physical AI, which deals with modeling and simulating the physical world. This domain presents unique challenges due to multi-dimensional data (3D space + 1D time = 4D), massive resolutions (e.g., 1000x1000x1000 grid leading to billions of "context tokens"), and multi-physics interactions (e.g., heat, fluid dynamics, microphysics in weather).

• Neural Operators: These are a class of neural networks designed to learn mappings between function spaces, rather than finite-dimensional vectors. This key difference allows them to be resolution-invariant: they can be trained on data at one resolution and perform inference at any other resolution, even finer ones. They can also seamlessly integrate physical laws (e.g., Navier-Stokes equations) as loss functions, even at resolutions where observational data is sparse.
• Architecture for Physical AI: Unlike traditional scientific computing which relies on iterative, local stencil, FP64 CPU-based simulations, Neural Operators leverage global convolutions implemented efficiently via Fourier transforms. Fourier transforms offer an N log N computational complexity, which is superior to the N^2 complexity of attention mechanisms for very large sequence lengths common in physical data. This enables efficient multi-resolution training and inference.
• Precision and Stability: While traditional HPC demands FP64 for numerical stability, Neural Operators can achieve high accuracy with much lower precision, such as FP16 or even FP8. A simple yet effective technique discovered is using tan H as a pre-activation function, which improves stability and allows for robust training at lower precisions, negating the need for complex normalization layers or gradient clipping.
• Tensor Galore: Given that weights and data in Physical AI are multi-dimensional tensors, the concept of Galore is extended to Tensor Galore. This method uses tensor rank decompositions and can incorporate sparsity to compress multi-dimensional gradients and optimizer states even more effectively than matrix-rank approaches, which is crucial for the memory-hungry nature of these models.

In summary, the technical deep dive illustrates a consistent theme: by understanding the underlying mathematical and computational properties of both algorithms and hardware, it's possible to devise highly efficient and often surprisingly robust solutions that transcend the limitations of conventional ML system design.

Experimental Setup & Results

▶ Watch: Mathematical reasoning with verification systems like Lean (8:00)

The talk presents compelling results across various domains, showcasing the efficacy of hardware-aware algorithmic co-design:

Inference Optimizations for LLMs:

• HeadInfer (KV Cache Offloading): Demonstrated significant memory reduction, enabling long-context inference on a consumer-grade NVIDIA RTX 4090 GPU. This highlights its potential for democratizing access to large models.
• MLP Sequence Breaking: Achieved the ability to handle much longer sequences while maintaining the same throughput, proving its effectiveness in overcoming intermediate memory bottlenecks without performance penalties. The method is lossless.
• Chain-of-Thought Pruning: Led to substantial memory savings in the KV cache for reasoning models. Surprisingly, it also resulted in a slightly improved accuracy (e.g., 105% relative accuracy), suggesting that pruning redundant information can reduce "confusion" and enhance decision-making.

Training Optimizations:

• Logarithmic Number System (LNS) Training: Through specialized multiplicative updates, training directly in LNS achieved a 90% reduction in energy consumption without any loss in accuracy compared to full-precision training. This demonstrates a path to highly energy-efficient ML.
• SignSGD (Single-Bit Gradient Communication): In distributed training, SignSGD enabled a 32x reduction in communication bandwidth by transmitting only the sign of gradients. This was shown to work well in practice, maintaining accuracy while drastically cutting communication overhead, especially beneficial for low-bandwidth environments.
• Galore (Gradient Low-Rank Projection):
• The initial work showed that Galore was the first method to enable pre-training of a 7B LLaMA model on a single consumer-grade NVIDIA RTX 4090 GPU. This underscored its ability to fit large models into significantly smaller memory footprints during training.
• Galore 2, a scaled-up version, was trained on 500 billion tokens from scratch. It achieved an 80% reduction in memory requirements for optimizer and gradient states, while matching the accuracy of an 8-bit Adam baseline. This result is critical for making large-scale pre-training more accessible and cost-effective.

Physical AI with Neural Operators:

• ForecastNet (High-Resolution Weather Model): This AI-based model was demonstrated to be tens of thousands of times faster (e.g., 10,000x to 40,000x) than traditional physics-based weather forecasts. Crucially, it was also highly accurate, capturing complex atmospheric phenomena like atmospheric rivers and even predicting hurricane landfalls several days earlier than physics-based models, despite limited hurricane data in its training set.
• Nuclear Fusion Plasma Simulation: An AI model simulating plasma evolution in a Tokamak reactor achieved a million times faster simulation speed compared to traditional physics-based methods. This speed, combined with accuracy and the ability to integrate real-time camera data, is vital for early detection of plasma disruptions, a major challenge for making fusion practical.
• Medical Device Design (Catheter): AI was used to optimally design the shape of triangular grooves within a catheter. Experimental results showed a 100x reduction in bacterial contamination by creating turbulence that prevents bacteria from swimming against the fluid current. This highlights the power of physics-informed AI for accelerating R&D and designing novel solutions.
• Neural Operator Capabilities: The underlying Neural Operator framework was shown to enable resolution-invariant learning and efficient operation with lower precision (e.g., FP16, FP8) for complex fluid dynamics, a domain traditionally demanding FP64 in HPC. The use of Fourier transforms and techniques like tan H pre-activation contributed to this efficiency and stability.
• Tensor Galore: The extension of Galore to multi-dimensional tensors, combined with sparsity, is showing promise in significantly reducing memory requirements for training these inherently memory-hungry Physical AI models.

Collectively, these results provide strong evidence that algorithmic co-design with hardware in mind can lead to groundbreaking advancements in efficiency, accessibility, and capability across various AI domains.

Practical Implications

▶ Watch: AI models hitting the hardware wall: compute scarcity (10:40)

The implications of Professor Anandkumar's work are profound for practitioners, infrastructure teams, model builders, and deployers navigating the increasingly resource-constrained landscape of large-scale AI.

For practitioners and model builders, the ability to train and infer large models with significantly less hardware fundamentally changes the accessibility equation. Techniques like Galore and HeadInfer mean that state-of-the-art models are no longer exclusively the domain of institutions with multi-million dollar GPU clusters. A student with a consumer-grade RTX 4090 can now pre-train or run long-context inference for models that previously required enterprise-grade hardware. This democratization fosters broader research and innovation, allowing more individuals and smaller teams to experiment with and deploy powerful AI.

Infrastructure teams stand to gain immensely from reduced memory and bandwidth demands. Lower memory footprints translate directly into fewer GPUs needed per model, or the ability to run more models on existing hardware, leading to substantial cost savings in acquisition and operational expenses (power, cooling). The 32x bandwidth reduction offered by SignSGD, for instance, can alleviate network bottlenecks in distributed training, enabling faster convergence and more efficient utilization of compute clusters, particularly in environments with less robust interconnects than NVLink or InfiniBand. The 90% energy reduction from LNS training also contributes to a more sustainable and cost-effective AI ecosystem.

For deployers, the inference-time optimizations are critical. Running long-context LLMs on edge devices or less powerful cloud instances becomes more feasible. The ability to prune redundant tokens in Chain-of-Thought reasoning, which also improves accuracy, is a win-win, offering better performance with lower resource consumption. This directly impacts the cost of deploying AI agents and advanced reasoning systems, making them more commercially viable.

The work in Physical AI presents a paradigm shift for scientific computing and engineering. AI models can replace or augment traditional physics simulations, which are often CPU-bound, rely on FP64 precision, and are not differentiable. The orders of magnitude speedups (e.g., 10,000x for weather, 1,000,000x for fusion) mean that complex simulations that once took days or weeks can now run in minutes or seconds, enabling rapid prototyping, iterative design, and real-time control. The physics-informed nature of these models, combined with their differentiability, allows for inverse design problems (e.g., finding the optimal catheter shape) that were previously intractable. This promises to accelerate scientific discovery and engineering innovation, drastically reducing R&D costs by minimizing the need for expensive physical experiments.

However, the talk also highlights tradeoffs and limitations. Adopting these hardware-aware algorithms requires a willingness to challenge established practices. The "noise can be a friend of learning" philosophy means moving away from the "lossless at all costs" mentality prevalent in traditional systems engineering. Implementing LNS or SignSGD requires changes to optimizer logic and potentially custom hardware support (though FP4/FP8 are getting closer). The massive, multi-dimensional data of Physical AI still presents immense challenges for current transformer systems, necessitating specialized architectures and tensor-parallel approaches that are not yet fully mature for industrial scale. The community must continue to invest in both algorithmic innovation and system-level integration to fully realize the potential of this co-design philosophy.

Key Takeaways

• Co-design is Crucial: To overcome the limitations of slowing hardware advancements, a fundamental shift from optimizing systems for fixed ML algorithms to co-designing ML algorithms with hardware constraints is essential.
• Extreme Quantization Works: Radical low-precision methods like Logarithmic Number Systems (LNS) for numerical representation and SignSGD for gradient communication can achieve massive energy (90%) and bandwidth (32x) reductions during training without sacrificing accuracy, sometimes even improving it.
• Low-Rank Optimization for Pre-training: Techniques like Galore (Gradient Low-Rank Projection) enable efficient pre-training of large foundation models by significantly reducing memory for gradients and optimizer states (80% reduction), making them accessible on consumer-grade GPUs.
• Smart Inference for Accessibility: Innovations like HeadInfer (KV cache offloading) and MLP sequence breaking democratize long-context LLM inference, while Chain-of-Thought pruning improves efficiency and accuracy in reasoning models.
• Physical AI Revolution: Neural Operators offer orders of magnitude speedups (10,000x to 1,000,000x) and superior accuracy for scientific and engineering simulations (e.g., weather, fusion, medical device design) by leveraging resolution-invariant, physics-informed learning, even at lower precisions (FP16/FP8).
• Noise as a Feature: Embracing "lossy" operations, such as extreme quantization or gradient compression, can lead to more robust and efficient learning, challenging the traditional lossless paradigm in system design.

About the Speaker(s)

Animashree Anandkumar is a distinguished Professor at Caltech, where her research focuses on the intersection of machine learning algorithms and systems. She is a leading voice in the field of hardware-aware AI, advocating for the co-design of algorithms and hardware to push the boundaries of what's possible in large-scale AI. Prior to her role at Caltech, she spent time at AWS, where she played a significant role in the development of SageMaker, Amazon's machine learning platform. Her deep understanding of hardware constraints is further informed by her position on the advisory board of SK Hynix, a major manufacturer of High Bandwidth Memory (HBM). Anandkumar's work bridges theoretical advancements in machine learning with practical systems challenges, fostering interdisciplinary collaborations between theorists, systems engineers, and domain experts.