Getting Started with Distributed Multi-GPU Libraries for Scalable AI and HPC | NVIDIA GTC 2025

Getting Started with Distributed Multi-GPU Libraries for Scalable AI and HPC

Speakers: Mads Kristensen, Sr. Software Engineer, NVIDIA

Conference: NVIDIA GTC

YouTube: https://www.youtube.com/watch?v=AjUUorlK-KU

Overview

In this insightful talk from NVIDIA GTC, Mads Kristensen, a Senior Software Engineer at NVIDIA, demystifies the landscape of multi-GPU programming for scalable AI and High-Performance Computing (HPC). The presentation addresses the growing necessity of leveraging multiple GPUs, driven primarily by the exponential increase in data sizes that often outpace the memory capacity of a single GPU, as well as the need to accelerate computation time. Kristensen provides a comprehensive overview of the various pathways developers can take to transition their existing sequential projects to a multi-GPU environment, highlighting the trade-offs between ease of use, generality, and fine-grained control.

The talk introduces a spectrum of programming models and libraries, categorizing them by their level of abstraction and how they manage data parallelism—explicitly or implicitly. Kristensen emphasizes that while the single-GPU ecosystem has matured into a highly optimized, interoperable environment, multi-GPU programming is still evolving, often requiring more manual data redistribution. He champions a pragmatic approach, encouraging developers to start with high-level, productive libraries and only descend to lower-level abstractions when performance bottlenecks or unsupported operations necessitate it. This article will delve into the technical details, practical implications, and key takeaways from Kristensen's expert guidance on navigating the complex world of distributed multi-GPU computing.

Background

▶ Watch: Introduction and motivation for multi-GPUs (0:00)

The journey into multi-GPU computing mirrors the early days of single-GPU acceleration with CUDA, albeit with its own set of challenges. Initially, CUDA enabled developers to offload specific bottlenecks from CPU to a single GPU, often involving frequent data transfers between host and device memory. Over time, the single-GPU ecosystem evolved dramatically, with a broad array of CUDA-accelerated libraries gaining the ability to communicate directly, often with zero-copy compatibility, bypassing the CPU almost entirely for complex workflows. This allowed for seamless data flow and optimized pipelines on a single GPU.

However, as Kristensen points out, the multi-GPU ecosystem is currently in an earlier stage of this evolution. While individual libraries might support multi-GPU operations, achieving seamless, zero-copy data transfer and inter-library communication across multiple devices often requires significant manual effort in data redistribution. This manual orchestration can be extremely expensive in terms of development time and computational overhead, presenting a major hurdle for scaling applications. The underlying problem is the lack of unified frameworks that allow distributed libraries to interoperate as smoothly as their single-GPU counterparts. The motivation for moving to multi-GPU is clear: problem sizes in AI and HPC are growing exponentially, exceeding the memory and computational limits of even the most powerful single GPUs. This necessitates distributed approaches, but the current state of tooling often forces developers back into a more explicit, lower-level programming paradigm than they might prefer.

Key Findings

▶ Watch: Overview of multi-GPU programming libraries and tools (4:00)

Mads Kristensen's talk highlights several key findings regarding the adoption and effective use of multi-GPU libraries for scalable AI and HPC:

1. The Multi-GPU Ecosystem is Evolving: Unlike the mature, interoperable single-GPU landscape, multi-GPU programming currently often requires manual data redistribution between libraries, making it more challenging but with a clear path towards unified frameworks.
2. Trade-off Between Generality/Productivity and Control: There's an inherent diagonal trade-off between how general and easy-to-use a multi-GPU library is versus how much explicit control it offers over parallelization and data distribution. Developers must choose based on their specific needs and expertise.
3. Two Core Parallelism Models:

• Explicit Data Parallelism: Frameworks expose parallelism to the user (e.g., SPMD, MPI-style). Programmers manage data distribution and communication. This offers high control and compatibility with existing MPI ecosystems but has a steep learning curve.
• Implicit Data Parallelism: Libraries infer parallelism from sequential-like code. Programmers write code as if it were sequential, and the framework handles distribution and optimization. This offers high productivity and abstracts away parallel programming complexities but can make performance reasoning difficult.

1. High-Level Libraries for Productivity: Libraries like cuNumeric (NumPy on multi-GPU) and Dask-cuDF (Pandas on multi-GPU) offer significant productivity gains by allowing users to accelerate existing sequential code with minimal changes, leveraging lazy evaluation and task graphs to infer and optimize distributed operations.
2. Common Pitfalls in Implicit Parallelism: Developers using implicit frameworks must be aware of patterns that can lead to inefficient distributed execution, such as Python loops with indexing, advanced indexing with generated indices, and unintentionally offloading large distributed data to single-GPU/CPU libraries.
3. Escape Hatches for Customization: Even with implicit frameworks, there are mechanisms (e.g., custom tasks in Legate, map_partitions in Dask) to drop down to lower-level control for unsupported operations or performance-critical sections, allowing for a hybrid approach.
4. Recommended Strategy: Start with the highest level of abstraction (implicit data parallelism) to maximize productivity and minimize parallel programming complexity. Only if performance falls short or specific operations are unsupported should developers descend to lower-level or custom explicit code.

Technical Deep Dive

▶ Watch: Categorizing parallel libraries: explicit vs. implicit (6:00)

Kristensen categorizes multi-GPU libraries along two dimensions: how domain-specific vs. general they are, and their productivity/ease of use. The ideal, highly general and easy-to-use library, is a continuous goal. The talk focuses on the more productive, higher-level abstraction libraries, specifically cuNumeric and Dask-cuDF, while also touching upon lower-level Legate and raw communication libraries. A crucial distinction is drawn between explicit and implicit data parallelism.

Explicit Data Parallelism

In the explicit model, the framework exposes parallelism directly to the programmer, typically adopting a Single Program, Multiple Data (SPMD) style. Each process, often mapped to a single GPU, runs the same code. The programmer is responsible for orchestrating data distribution, managing communication patterns (e.g., to avoid deadlocks), and ensuring data compatibility between processes. This model is powerful for its control and compatibility with the vast MPI (Message Passing Interface) ecosystem, especially when using device-aware MPI which can directly send and receive GPU buffers. However, it demands expertise in parallel programming and MPI.

Kristensen illustrates this with examples:

• MPI for Python with CuPy: A Python implementation of a distributed matrix multiplication (Cannon's algorithm) using mpi4py and CuPy (a NumPy-like library for GPUs). The programmer manually handles rank management, data decomposition, and communication primitives (comm.Send, comm.Recv). This requires explicit knowledge of how data blocks are shuffled and processed.
• cuBLAS-MP: A C-based example showing distributed BLAS (Basic Linear Algebra Subprograms) operations on multiple GPUs. Here, the library provides a higher-level abstraction for the matrix multiplication itself, but the programmer still describes the data distribution and grid of GPUs to the library.
• PyTorch Distributed: While more abstracted than raw MPI, PyTorch's distributed tensors still operate within a grid-like or mesh concept, where the user generates tensors and the framework handles the distributed matrix multiplication. It offers varying levels of abstraction within the explicit paradigm.

Implicit Data Parallelism

The implicit model aims for high productivity by abstracting away the complexities of parallel programming. The user writes code as if it were sequential, and the library or framework infers data parallelism from these operations. This approach leverages lazy evaluation, where operations are tracked to build a task graph rather than executed immediately. Computation is only triggered when a specific result (e.g., a scalar value printed to the screen) is required. This gives the framework significant freedom to optimize data distribution, communication, and computation scheduling. The main downsides are that reasoning about performance can be challenging (as the underlying distributed operations are hidden), and certain sequential-like patterns can lead to unexpected performance degradation in a distributed context.

Key examples discussed:

• cuNumeric: This project aims to accelerate NumPy to run efficiently across multiple GPUs. By simply importing cunumeric instead of numpy, existing sequential NumPy code can often run in a distributed multi-GPU environment with minimal changes. cuNumeric infers data parallelism, handles duplication control, and manages data distribution. It is built on the Legate ecosystem, using Legate as its backend. This is a significant architectural choice, as it means cuNumeric arrays are compatible (often via zero-copy) with other Legate-backed libraries (like Legate dataframes or ML libraries), enabling cross-library optimization and fusion of operations.
• Dask-cuDF: This library provides a GPU-accelerated equivalent of Pandas dataframes, distributed across multiple GPUs. Dask acts as the distributed scheduler, and cuDF provides the GPU-accelerated Pandas-like primitives. Users can configure Dask to use cudf as its backend. Similar to cuNumeric, operations are lazy, and a compute() call triggers the execution of a task graph composed of cuDF operations. Dask-cuDF also features advanced memory management, including memory spilling (automatically moving data from GPU to host memory or disk) to handle datasets larger than aggregate GPU memory. It's part of the broader RAPIDS ecosystem.

Pitfalls and Escape Hatches in Implicit Parallelism

Kristensen highlights common pitfalls when using implicit frameworks:

• Avoid Python Loops with Indexing: Just as in single-GPU NumPy, direct Python loops iterating over distributed arrays with indexing are highly inefficient. Instead, vectorized cuNumeric operations or slicing should be used to allow the framework to optimize.
• Advanced Indexing: Using arrays of indices (e.g., arr[indices]) can be problematic. While valid on a single GPU, generating an array of indices from a distributed array can lead to significant communication overhead when those indices are then used for subsequent indexing, as the indices themselves might be distributed. Boolean masks are often a more efficient alternative.
• Unintentional Offloading: Calling a single-GPU/CPU library (e.g., scipy.linalg.cholesky) on a large cuNumeric array will force the entire distributed array to be gathered onto a single GPU (or even CPU memory), leading to out-of-memory errors or severe performance degradation. The solution is to use the distributed equivalent if available (e.g., cunumeric.linalg.cholesky).

When faced with unsupported operations or performance issues, implicit frameworks offer "escape hatches":

• Local Interpretation: A distributed array can be explicitly converted to a local array (e.g., arr.to_cupy_array()). This allows for local custom operations but will gather data to a single GPU, becoming a bottleneck for large datasets.
• Custom Tasks: For truly high-performance custom logic, users can define custom tasks using the framework's backend API. For cuNumeric, this means writing Legate tasks, which define how inputs and outputs are handled, and can run across multiple GPUs with explicit control over resource allocation and memory semantics. For Dask-cuDF, this involves using map_partitions to apply a custom function to each partition of a distributed dataframe, allowing for custom cuDF operations. This requires deeper knowledge of the backend but provides maximum flexibility and performance.

Coherent Shared Memory (Grace Blackwell)

During the Q&A, the topic of Grace Blackwell and coherent shared memory arises. While these architectures provide unified, coherent memory spaces where CPU and GPU can directly access each other's memory, Kristensen cautions that this doesn't automatically translate to optimal performance. While it simplifies programming by eliminating explicit memcpy operations, achieving high performance still requires providing hints to the runtime system regarding memory access patterns (e.g., prefetching) to avoid performance degradation, especially when accessing remote memory. It's a stepping stone for ease of use but not a complete automation of performance optimization.

Experimental Setup & Results

▶ Watch: Understanding explicit data parallelism: SPMD style (6:17)

The talk primarily focuses on programming models, library capabilities, and best practices rather than presenting formal experimental setups or detailed benchmark results. Mads Kristensen provides motivational examples and "not so fair comparisons" to illustrate the potential performance gains and productivity benefits of multi-GPU acceleration.

For cuNumeric, he shows a comparison of a conjugate gradient solver:

• A NumPy single-CPU execution versus an 8-GPU cuNumeric execution.
• While not a fair apples-to-apples comparison due to the hardware difference, the point is that such significant speedups can be achieved with minimal code changes (often just changing an import statement). This highlights the productivity gain rather than a direct performance benchmark against an optimized parallel CPU implementation.

For Dask-cuDF, an example involves processing US housing prices and a more complex use case:

• Nemo Curator, a pre-processing pipeline for Large Language Models (LLMs) that involves de-duplication of tokens and parsing input.
• By using Dask-cuDF, Nemo Curator was able to achieve "very good performance." Again, this is implicitly compared to CPU-based alternatives, emphasizing the acceleration without providing specific throughput numbers or a detailed breakdown of the experimental setup (e.g., dataset size, specific hardware, baseline performance).

It's important to note that the purpose of these examples is to motivate the use of these libraries and demonstrate the ease with which existing sequential code can be accelerated, rather than to provide rigorous scientific benchmarks. The talk acknowledges that detailed performance analysis and optimization often require dropping to lower-level abstractions or using specialized profiling tools, which both cuNumeric and Dask-cuDF offer.

Practical Implications

▶ Watch: Understanding implicit data parallelism: library handles it (7:49)

The insights from this talk have significant practical implications for practitioners, infrastructure teams, model builders, and deployers in the AI/ML and HPC domains.

1. Productivity First, Optimization Second: The overarching recommendation is to prioritize productivity by starting with high-level, implicit data parallel frameworks like cuNumeric and Dask-cuDF. This allows model builders and domain experts to leverage multi-GPU power without becoming parallel programming experts, significantly accelerating development cycles. Only when performance bottlenecks are identified or specific operations are unsupported should teams invest in lower-level, more explicit code.
2. Strategic Library Choice:

• For existing MPI/SPMD codebases or experts: The explicit data parallelism route, potentially using device-aware MPI with libraries like CuPy or PyTorch Distributed, is a natural fit. It offers maximum control and can integrate seamlessly with existing distributed infrastructure.
• For NumPy/Pandas users: cuNumeric and Dask-cuDF offer a direct path to multi-GPU acceleration with minimal code refactoring, making them ideal for data scientists and ML engineers familiar with these popular Python libraries.

1. Understanding Implicit Frameworks' Nuances: While productive, implicit frameworks require users to be aware of certain "pitfalls" (e.g., Python loops for indexing, advanced indexing, offloading to single-GPU/CPU libraries). Practitioners need to learn the "idiomatic" way to write code for these distributed libraries to avoid performance traps. This might involve adopting vectorized operations, using boolean masks, and leveraging the library's native distributed functions.
2. The Power of Ecosystems (Legate, RAPIDS): The talk highlights the value of integrated ecosystems like Legate (for cuNumeric) and RAPIDS (for Dask-cuDF). These provide common backends and primitives, enabling zero-copy data transfer and potential cross-library optimization, which is crucial for building complex, multi-stage data pipelines efficiently.
3. Hybrid Approaches are Key: No single abstraction level is perfect. The ability to "escape" to lower-level custom tasks (e.g., Legate tasks, Dask map_partitions) provides a crucial safety net. Infrastructure teams can develop highly optimized custom kernels for specific bottlenecks, which can then be seamlessly integrated into the high-level workflows of model builders. This facilitates a division of labor between domain experts and parallel programming specialists.
4. Profiling and Monitoring are Essential: For implicit frameworks where parallelism is hidden, robust profiling and monitoring tools are critical. These tools allow users to verify that their jobs are indeed utilizing multiple GPUs efficiently, identify communication bottlenecks, and understand how data is distributed—addressing concerns about "taking a day to run because they don't know it's running on a single GPU."
5. Coherent Memory: A Step, Not a Solution: While hardware advancements like Grace Blackwell with coherent shared memory simplify memory management between CPU and GPU, they do not fully automate performance optimization for distributed systems. Developers still need to consider data locality and provide hints (like prefetching) for optimal performance, especially when dealing with remote memory access patterns. This means that while the programming model might become easier, the need for performance tuning expertise remains.

Key Takeaways

• Multi-GPU is Essential: Data growth and computational demands necessitate multi-GPU solutions for AI and HPC, pushing beyond single-GPU limits.
• Choose Your Abstraction Level: Decide between explicit (high control, complex) and implicit (high productivity, abstracted) data parallelism based on project needs and team expertise.
• Start High, Go Low When Needed: Begin with high-level, implicit frameworks like cuNumeric or Dask-cuDF for rapid development, and only drop to lower-level custom tasks when encountering performance bottlenecks or unsupported operations.
• Understand Implicit Framework Pitfalls: Be aware of common patterns (e.g., loops with indexing, advanced indexing, single-GPU offloads) that can lead to inefficient distributed execution in implicit models.
• Leverage Ecosystems: Frameworks built on common backends (e.g., Legate for cuNumeric, RAPIDS for Dask-cuDF) enable efficient, zero-copy data transfer and optimization across different libraries.
• Profiling is Crucial: Use profiling tools to monitor GPU utilization and data distribution when using implicit frameworks to ensure efficient multi-GPU execution.

About the Speaker(s)

Mads Kristensen is a Senior Software Engineer at NVIDIA. His work focuses on developing and advancing distributed multi-GPU libraries for scalable AI and HPC applications. He possesses deep expertise in CUDA, parallel programming models, and the intricate challenges of optimizing data flow and computation across multiple GPUs. Kristensen is a key contributor to NVIDIA's efforts in building unified and highly productive software ecosystems, such as Legate, to democratize access to high-performance computing on GPU clusters.