Extreme PyTorch: Inside the Most Demanding ML Workloads—and the Open Challenges in Building AI Agents to Democratize Them

Extreme PyTorch: Inside the Most Demanding ML Workloads—and the Open Challenges in Building AI Agents to Democratize Them

Speakers: Soumith Chintala, Scientist-Engineer, Meta / NYU

Conference: MLSys 2025

YouTube: https://slideslive.com/39042963

Overview

Soumith Chintala, a distinguished Scientist-Engineer at Meta and NYU, presented a comprehensive talk at MLSys 2025, delving into the intricate world of PyTorch and its application to the most demanding machine learning workloads. The presentation, meticulously assembled with contributions from many Meta colleagues, aimed to demystify why PyTorch, despite its seemingly simple "Hello World" interface, has evolved into a multi-million-line, multi-gigabyte binary. Chintala articulated the core challenges and philosophical underpinnings driving PyTorch's development, particularly in the context of extreme-scale AI training, such as the Llama series of large language models.

The talk provided a unique insider's perspective on the engineering complexities involved in scaling PyTorch from single-GPU setups to clusters comprising tens of thousands of GPUs. Chintala meticulously detailed the critical dimensions of compute, memory, and communication, and how their precision, reliability, speed, and cost dictate the design choices in high-performance ML systems. Beyond the technical deep dive, a significant portion of the presentation focused on PyTorch's strategic vision: to foster a "combinatorial explosion" of AI innovation by ensuring composability of features, lowering the barrier to entry for ML scientists, and exploring the future of AI agents that can democratize complex HPC tasks. This talk is crucial for anyone involved in AI research, infrastructure, or deployment, offering insights into the present state and future trajectory of large-scale machine learning.

Background

▶ Watch: Talk introduction and agenda overview (0:00)

PyTorch, at its core, is envisioned by many as a simple framework for defining models and training loops. However, this "Hello World" simplicity belies the immense complexity required to execute such programs efficiently across a diverse and rapidly evolving hardware landscape. Chintala highlighted that the fundamental reason for PyTorch's substantial codebase (millions of lines, 2GB binary) is its necessity to abstract and optimize for heterogeneous hardware, including NVIDIA GPUs (e.g., Ampere, Blackwell with specialized matmul units), AMD GPUs, and even specialized accelerators like Cerebras. The underlying challenge, as Chintala succinctly put it, revolves around compute, memory, and communication—how much of each exists, where it is located, and how fast data can move between them.

Scaling these fundamental dimensions introduces further complexity. Whether mapping a program to a single GPU, multiple GPUs in a node, or a cluster of a million GPUs, engineers must contend with four critical sub-dimensions: precision, reliability, speed, and cost. At extreme scales, such as a 100,000-GPU cluster, communication becomes inherently unreliable due to hundreds of thousands of connections, and individual compute nodes (GPUs) are expected to fail frequently. This necessitates sophisticated fault tolerance mechanisms. PyTorch's widespread adoption means its influence extends from massive data centers training cutting-edge LLMs to mobile phones, self-driving cars, robots, and even, humorously, "microwaves" and veterinary AI cancer scans, each posing unique constraints on memory footprint, just-in-time compilation, and static memory allocation.

The landscape of AI innovation further complicates PyTorch's mission. From early ConvNets to Generative Adversarial Networks (GANs) and the ubiquitous Transformers, new architectures and training paradigms emerge constantly. Beyond simple forward-backward passes, complex workflows like reinforcement learning (RL) for agents interacting with environments (which might be browsers or even human annotators) introduce timing complexities and external dependencies. These "web server stack" problems, while not new to software engineering, are novel to the AI world and demand optimization for components beyond core model training, such as data loading and environment interaction. Chintala outlined a consistent taxonomy of AI innovation: improving model architectures, developing better optimizers, scaling with larger models and more data, defining clever objective functions, injecting better priors (e.g., AlphaFold), collecting better data, leveraging unsupervised/semi-supervised learning, and, critically, using hardware more efficiently or designing better hardware. PyTorch's evolution is directly tied to enabling and accelerating innovation across all these fronts.

Key Findings

▶ Watch: Understanding compute, memory, communication in GPUs (3:00)

The talk underscored several key findings and strategic directions for PyTorch, revealing its philosophy and technical advancements.

Firstly, PyTorch's core mission is to enable a combinatorial explosion of AI innovation. This means ensuring that diverse new ideas—such as CPU offloading, activation checkpointing, and novel neural network architectures—are simultaneously expressible and composable within the framework. PyTorch actively works to prevent a "framework lottery," where users are bottlenecked by the underlying platform, as seen with earlier frameworks like Caffe, which rigidly focused on ConvNets. The goal is to allow ML scientists to push boundaries without needing to become HPC engineers.

Secondly, PyTorch unequivocally prioritizes the ML scientist experience over the HPC scientist's ideal. This is evident in its compiler design, where non-negotiable constraints include support for dynamic shapes, partial graphs, mutation, side effects, and custom operations. While these features complicate compiler development, they are crucial for the flexibility ML scientists demand. This user-centric approach has even led hardware vendors, initially resistant, to adapt their backends to PyTorch's open-ended requirements.

Thirdly, PyTorch is actively leveraging AI to bridge the gap between ML scientists and HPC scientists. The vision is to create a world where an AI co-pilot can assist scientists in optimizing their workloads, making the distinction between ML and HPC skills less relevant. This aligns with PyTorch's principle of "not being married to a technology," but rather to the user experience, embracing whatever tools or methods best serve its users.

In terms of technical advancements for Extreme PyTorch, several innovations were highlighted:

• Performance Optimizations: TorchCompile for just-in-time compilation, Triton for custom kernel development, FlexAttention as a hand-optimized kernel, and continuous bread-and-butter GPU optimizations.
• Precision: Industry-wide efforts towards FP8 for training and FP4 for inference, with PyTorch baking this innovation into torchao, a faster-moving library for cutting-edge quantization and precision schemes.
• Reliability: The introduction of torchft (fault tolerance) to make training more resilient to failures at scale, and significant improvements in distributed checkpointing (torchcheckpoint), achieving a 19x speedup and an additional 5x reduction in overhead for zero-overhead checkpointing.
• Mobile/Edge Deployment: ExecuTorch, a new mobile runtime, pushes the limits of efficiency while maintaining PyTorch's generality, with a footprint of a few kilobytes.
• Frontend Leverage: PyTorch's frontend is so sticky that many backends, including those from TinyGrad, Modular, and various hardware vendors, are developing PyTorch backends. New DSLs like Helion are being developed to simplify GPU kernel writing.

For Llama training at scale, Meta's experience provided concrete findings:

• Custom Infrastructure: Training Llama 3 on 24,000 GPUs (and Llama 4 on 100,000+ GPUs) required building new data centers, custom RDMA over Ethernet network fabric, and a custom scheduler called MAAS to maximize effective training time by accelerating job initialization.
• Parallelism Challenges: Implementing 3D/4D parallelism, pipeline parallelism, and data parallelism across buildings necessitates complex sharding strategies and careful "Excel sheet exercises" due to varying bandwidths.
• Training Stability: Debugging failures (GPU, network, software bugs) is a "nightmare" at scale. Tools like the NCCL flight recorder (baked into PyTorch) are essential for post-mortem analysis.
• Future Paradigms: The SPMD programming model is breaking down at extreme scales, especially with cross-data center and pipeline parallelism. PyTorch is exploring a Pathways-style single controller model with device meshes to address this.

Finally, PyTorch is actively fostering a community and tools for democratizing HPC performance using AI. Initiatives like KernelBench, KernelBot, and KernelLLM aim to crowdsource and use LLMs to generate optimized GPU kernels, creating leaderboards and competitions to accelerate this process.

Technical Deep Dive

▶ Watch: Core challenge: compute, memory, communication hierarchy (4:00)

PyTorch's ability to operate from microcontrollers to million-GPU clusters stems from its sophisticated handling of the compute, memory, and communication hierarchy. This abstraction layer allows PyTorch to map high-level Python programs to diverse hardware, such as NVIDIA's Blackwell with its dedicated matmul units or AMD's different GPU architectures. The framework continuously optimizes how workloads interact with various memory levels (registers, shared memory, HBM) and communication channels (PCIe, NVLink, network interconnects).

At scale, PyTorch addresses the four critical dimensions of precision, reliability, speed, and cost. For precision, the focus is on enabling lower-precision training (e.g., FP8) and inference (e.g., FP4) to reduce energy consumption and memory footprint. This innovation is often too rapid for PyTorch's core, so it's incubated in libraries like torchao. Speed is tackled through TorchCompile, a JIT compiler that optimizes graph execution, and specialized kernels like Triton and FlexAttention. Triton allows users to write highly optimized GPU kernels in Python, while Helion, an experimental DSL, further simplifies this by eliminating boilerplate for common patterns like tiled loops.

Reliability is paramount in large-scale, long-running training jobs where hardware failures are inevitable. TorchFT (Fault Tolerance) is a PyTorch package designed to make training more resilient, allowing jobs to continue despite GPU or communication failures. Complementing this, distributed checkpointing (torchcheckpoint) has seen significant improvements, reducing overhead by 19x and enabling "zero-overhead" checkpoints with an additional 5x reduction. This is crucial for recovering from faults without tearing down and restarting entire jobs that can run for weeks. For debugging, the NCCL flight recorder is integrated into PyTorch, providing post-mortem logs of communication events, which are invaluable for diagnosing issues like network congestion or software bugs.

The challenges of Llama training at Meta exemplify PyTorch's extreme usage. Training Llama 3 on 24,000 GPUs and Llama 4 on 100,000+ GPUs required a bespoke infrastructure. This included custom data centers designed for high-density GPU deployment, an open network fabric utilizing RDMA over Ethernet (a significant engineering feat compared to InfiniBand), and a custom scheduler called MAAS. MAAS is critical for minimizing job initialization time by pre-discovering network topologies and optimizing communication algorithms, which can otherwise consume substantial time in large-scale distributed jobs.

Parallelism strategies for LLM training are highly complex, involving 3D/4D parallelism, pipeline parallelism, and data parallelism. For Llama, these strategies had to be adapted to even span across different buildings within data centers, where effective bandwidth can be reduced. This necessitates meticulous "Excel sheet exercises" by GPU engineers to shard models and data optimally.

Chintala also highlighted the limitations of the current SPMD (Single Program, Multiple Data) programming model at extreme scales. While effective for many distributed workloads, SPMD struggles with the complexities of cross-data center communication and advanced sharding concepts like pipeline parallelism, where different ranks might execute vastly different control flows. To address this, PyTorch is exploring a Pathways-style single controller programming model (similar to Google's Pathways or Ray) with device meshes. This model allows a single controller to manage hundreds of thousands of GPUs, enabling more flexible and performant orchestration of large, heterogeneous tensor graphs and communications. An early prototype shows the ability to control massive GPU clusters from a Jupyter notebook, offering unprecedented control and interactivity.

Finally, for edge and mobile deployments, ExecuTorch provides an entirely new runtime specifically designed for efficiency, with a minimal footprint (a few kilobytes) while supporting the generality of PyTorch programs. This allows PyTorch to ship to a vast array of embedded devices.

Experimental Setup & Results

▶ Watch: Four dimensions of scaling to millions of GPUs (5:00)

The talk provided concrete examples of PyTorch's capabilities and the results achieved in demanding ML environments.

Llama Training:

• Models: Llama 3 and Llama 4 (future focus).
• Scale: Llama 3 was trained on 24,000 GPUs. Llama 4 is planned or being trained on a 100,000+ GPU cluster.
• Infrastructure: Custom data centers optimized for tens of thousands of GPUs. Open network fabric based on RDMA over Ethernet, a custom scheduler called MAAS.
• Job Duration: These jobs run for weeks on end on thousands of GPUs, highlighting the critical need for reliability and efficient recovery.
• Checkpointing Improvements: PyTorch's distributed checkpointing (via torchcheckpoint) reduced overhead by 19x. An additional 5x reduction was achieved for "zero-overhead" checkpointing, significantly improving job recovery and efficiency.
• Debugging: The NCCL flight recorder is a key debugging tool, baked into PyTorch, allowing for post-mortem analysis of communication issues.
• Parallelism: Extensive use of 3D/4D parallelism, adapted for challenges like pipeline parallelism across different physical buildings with varying bandwidths.

PyTorch Core Development & Testing:

• "Boiled Ocean Testing": PyTorch invests millions of dollars annually into an extensive CI testing fleet, comprising both NVIDIA and AMD GPUs, to ensure correctness and psychological safety for users across diverse hardware. This massive testing infrastructure is crucial for guaranteeing reliable results given the framework's complexity and broad usage.

Emerging Tools and Initiatives:

• Helion: An early, experimental DSL for easier GPU kernel writing, eliminating boilerplate often found in Triton. It's available on GitHub.
• KernelBot/KernelLLM: These initiatives are designed to democratize HPC performance. KernelBot surfaces interesting kernel optimization problems from state-of-the-art neural networks (e.g., softmax, add, matmul) with relevant shapes collected from various sources. It features an online bot and a leaderboard for community submissions. KernelLLM is early work fine-tuning Llama models to improve kernel generation, showing results comparable to much larger models like DeepSeek-Coder-V1. A competition with AMD (corrected from NVIDIA in the transcript) offered a $100,000 prize for faster kernels, spurring significant community contributions, particularly for AMD kernels.

These results and ongoing initiatives demonstrate PyTorch's continuous effort to push the boundaries of scale, efficiency, and accessibility in machine learning.

Practical Implications

▶ Watch: Mapping complex programs to large GPU clusters (6:30)

The insights from this talk have profound implications for various stakeholders in the AI ecosystem:

For Practitioners (ML Scientists, Model Builders):

PyTorch's commitment to flexibility – supporting dynamic shapes, partial graphs, mutation, and custom ops – means ML scientists can iterate rapidly and build complex models without being constrained by the framework. This fosters innovation and reduces the engineering overhead often associated with highly optimized, rigid systems. Tools like Helion offer a lower barrier to entry for writing custom GPU kernels, allowing practitioners to squeeze out more performance without needing deep HPC expertise. The existence of torchao means cutting-edge precision techniques (FP8, FP4) are accessible sooner, enabling more efficient training and deployment.

For Infrastructure Teams and Deployers:

The challenges highlighted in Llama training underscore the necessity of robust and fault-tolerant infrastructure. For long-running, extreme-scale jobs, TorchFT and highly optimized distributed checkpointing (19x speedup) are critical for maximizing effective training time and minimizing costly restarts. Infrastructure teams must design for inherent unreliability at scale, including GPU failures and network congestion. The move towards RDMA over Ethernet for custom network fabrics indicates a trend towards specialized, high-performance networking solutions. For deployment, ExecuTorch provides a pathway for deploying PyTorch models to resource-constrained environments like mobile phones and embedded systems, expanding the reach of AI applications.

For HPC Engineers:

While PyTorch prioritizes ML scientists, the talk acknowledges that HPC engineers are crucial for building the underlying optimizations. However, they must contend with PyTorch's unique constraints (dynamicism, side effects) that diverge from traditional compiler design. This necessitates a different approach to optimization, where the goal is to enable flexibility rather than pure, static performance. Tools like FlagGems allow HPC researchers to easily override PyTorch ops for performance research, bridging the gap between low-level optimization and high-level framework integration. The call for a Pathways-style single controller model signals a future where HPC engineers will need to design new programming paradigms to manage heterogeneous, massive clusters effectively.

Tradeoffs and Limitations:

The talk implicitly reveals several tradeoffs. The flexibility for ML scientists comes at the cost of increased complexity for HPC engineers building the PyTorch backend. While PyTorch strives for composability, managing an ever-growing array of innovations can lead to maintenance challenges. Debugging at extreme scale remains a "nightmare," with many effective tools being infrastructure-specific and not easily open-sourced. Furthermore, the talk emphasizes that physical limitations, particularly power availability, are becoming a fundamental constraint on building larger AI clusters, even more so than financial investment. This points to a future where energy efficiency will be a paramount design consideration for both hardware and software. The current SPMD model is reaching its limits for advanced parallelism strategies like pipeline parallelism across disparate compute resources, indicating a need for fundamental shifts in distributed programming models.

Key Takeaways

• ML Scientist-Centric Design: PyTorch's core philosophy prioritizes the flexibility and iterative workflow of ML scientists, even if it introduces complexity for underlying HPC engineering (e.g., dynamic shapes, partial graphs in compilers).
• Extreme Scale Demands New Engineering: Training models like Llama 3 (24,000 GPUs) and Llama 4 (100,000+ GPUs) requires custom data centers, specialized network fabrics (RDMA over Ethernet), advanced schedulers (MAAS), and sophisticated fault tolerance (TorchFT, distributed checkpointing).
• Reliability is Paramount: At scale, hardware failures and communication issues are inevitable. PyTorch addresses this with tools like TorchFT and dramatically improved distributed checkpointing (19x faster), alongside debugging aids like the NCCL flight recorder.
• Evolving Programming Models: The SPMD model is breaking down for extreme, heterogeneous, and cross-data center parallelism. PyTorch is exploring Pathways-style single controller models with device meshes to manage future massive AI clusters more effectively.
• Democratizing HPC with AI: PyTorch is actively investing in AI agents and tools (KernelBot, KernelLLM, Helion) to lower the barrier for ML scientists to achieve high-performance computing, enabling them to optimize GPU kernels and systems without deep HPC expertise.
• PyTorch as a "YOLO" Framework: Not married to specific technologies, PyTorch's primary commitment is to user experience, embracing whatever innovations best serve the AI community, from cutting-edge precision (torchao) to efficient edge runtimes (ExecuTorch).

About the Speaker(s)

Soumith Chintala is a prominent figure in the machine learning community, serving as a Scientist-Engineer at Meta and holding an affiliation with NYU. He is widely recognized as one of the co-founders of PyTorch, the leading open-source machine learning framework. His work at Meta involves not only the core development of PyTorch but also a broad range of infrastructural initiatives that support large-scale AI research and deployment. Chintala's expertise spans the theoretical underpinnings of machine learning and the practical engineering challenges of building and scaling AI systems, making him a key contributor to the advancement of the field.