with Michael Isaev², Nic McDonald², Jie Ren³, Vijay Thakkar² ⁴, Haicheng Wu², and Humphrey Shi¹.

¹ SHI Labs @ Georgia Tech, ² NVIDIA, ³ KAUST, ⁴ HPC Garage @ Georgia Tech.

Introduction

Much of AI today is continuing to scale up due to massively parallel systems. General purpose GPUs (GPGPUs) are made up of many parallel processors, each with their own local memory, execution engines, and register files. However, we are long past the time that single GPUs were enough to scale up, and are now in the age of much larger systems, such as NVIDIA’s GB200 NVL72, which pack tens of GPUs together, with an amazingly high-bandwidth network of switches enabling fast peer-to-peer communication. To end users, the GB200 NVL72 system will be “One giant CUDA GPU”.

For scaling up to the order of trillions of parameters, and utilizing the massive compute power provided by such systems, model parallelism techniques are the way forward, and among them, Tensor Parallelism (TP) has been the de-facto standard for large-scale training and inference in NVLink-based networks. However, communication can be a pain point for developers, researchers, and other end users. Not only can communication become a bottleneck and lead to under-utilization of compute resources, it would also require extra implementation effort.

In this post, we introduce a new implementation of Tensor Parallel GEMMs, which we call Distributed GEMM. This approach makes use of software and hardware features provided by the CUDA runtime and driver APIs alone to pipeline and overlap communication and computation, and makes very few assumptions about the specifics of the underlying GEMM kernel. More specifically, Distributed GEMMs use CUDA Graphs to allow the GPU Copy Engine (CE) to handle communication, leaving the Streaming Multiprocessors (SMs) and Tensor Cores (TCs) free to execute fast and performant GEMMs, unburdened by additional communication kernels and communication-related instructions. This approach can provide concurrency without breaking the CUDA programming model, and minimizes kernel launch latency gaps not just by using graphs, but also Programmatic Dependent Launch (PDL), a feature new to the Hopper architecture.

Distributed GEMM is implemented entirely using NVIDIA’s open source linear algebra framework, CUTLASS, and additional calls to the runtime and driver APIs. Distributed GEMMs can therefore extend the incredibly fast and performance-optimized kernels from NVIDIA to their tensor parallel form, with only a few lines of code added, and no additional dependencies!

Background

When you train or query a model on multiple GPUs, the simplest choice is data parallelism; that is to replicate weights across GPUs, and feed them different mini-batches of data. If you’ve ever used PyTorch’s Distributed Data Parallelism, as the name suggests, you’ve used data parallelism. However, data parallelism isn’t always the best approach, given that it forces users to have a minimum batch size, and in the classical case assumes the entire model can fit on a single GPU.

Model parallelism techniques on the other hand distribute weights across multiple processors, and this is especially important in larger models like LLMs, where memory capacity is a concern. Among those are Pipeline Parallelism (PP), Expert Parallelism (EP), Fully-Sharded Data Parallelism (FSDP) [3], and Tensor Parallelism (TP) [1,2]. Each of these techniques have their own advantages and limitations. In this post, we focus specifically on Tensor Parallelism, as it achieves lower memory capacity requirements and better GPU utilization when implemented in low-latency high-bandwidth networks such as NVLink. For that reason, it became the de-facto standard for many LLM implementations across various frameworks. That is especially the case in most LLM training infrastructure today, and has been since the technique was first introduced in Megatron-LM [1,2]. We’ll quickly go over a few basics in Tensor Parallelism.

The figure below from Korthikanti et al. [2] illustrates how a Transformer layer is divided across multiple processors. Since both activation and weight tensors are divided across processors, certain operations produce partial results, which will require communication.

For simplicity, we’ll exclude attention for now, which today uses Context Parallelism, an approach compatible with fused attention. In the MLP block, there are two back-to-back linear layers, separated by a GELU activation (which is typically fused into the preceding kernel). The communication and computation patterns fall into two categories here: All Gather + GEMM, and GEMM + Reduce Scatter. In both cases, each device “owns” a slice of the weights and activations. In the typical All Gather + GEMM scenario, activation slices which are distributed across devices are gathered locally on each device, and then the GEMM can produce a slice of the output. In the typical GEMM + Reduce Scatter scenario, each device produces a partial result of the output, which requires a reduction (summation) and scattering of slices of the output.

In either case, we can break down the computation and communication further into stages in order to pipeline some or potentially all communication behind computation.

Aside from saving on activation memory, with pipelined tensor parallelism, we can also hide the exposed communication step under certain conditions. Instead of performing collective communications, we can use point-to-point communication and allow each processor to start working on its local buffers immediately in the first iteration. This eliminates the exposed communication step, as illustrated in the figure below.

This means that the wall time will nearly only reflect computation time, as all communication will be hidden behind it. However, note that single-stage TP can still be a more performant choice in more latency-bound cases, such as small-batch short-sequence decoding. To make matters worse, the local GEMM stages in such cases can also easily wind up memory-bandwidth-bound.

The concept is not by any means new and has long existed in HPC, and was adopted in AI applications relatively late, by Google Research [5], Microsoft Research [6], ByteDance [7], and PyTorch [8].

The difference with our new approach is not in functionality, but rather in implementation. Instead of relying on communication primitives in libraries such as NCCL, or fusing communication steps into a specific kernel, we simply use host-side CuTe layout logic, peer-to-peer access, and existing CUDA graph APIs to implement Distributed GEMMs, adding very few instructions to the GEMM kernels themselves. Aside from being more extensible and being much less prone to register spills, this approach also does not limit itself to a specific kernel, making performance optimization through profiling much more convenient.

Approach

CUTLASS is a framework for writing performant kernels targeting NVIDIA’s Tensor Cores, since their introduction with the Volta architecture in 2017. The more recent introduction of the Hopper architecture, which included a Tensor Memory Accelerator (TMA), led to a change in the programming model that led to the CUTLASS 3 API, powered by an expressive framework for tensor layout algebra, CuTe. However, most compute-heavy kernels within CUTLASS (GEMMs and convolution, both of which are matrix-matrix multiplications and thus target Tensor Cores) are designed to run on a single GPU, and can deliver very high performance.

We approach the implementation of Distributed GEMMs from a perspective that makes the following observations:

1. In most compute-bound problems, we prefer to pipeline the local GEMMs and overlap communication and computation.
2. Kernels written for modern architectures like Hopper heavily rely on SM occupancy (persistent kernels), which means communication primitives should try and avoid occupying SMs entirely.
3. Given the many choices of kernel parameters that the Hopper programming model allows for, computation kernels cannot realistically be rewritten entirely for distributed settings.
4. Kernel launch latency (the gap between multiple pipeline stages of the compute kernel) can be an extra overhead that should be avoided.

To that end, we introduce Distributed GEMMs, which aside from pipelining local GEMMs, take a direct-access approach, where communication is primarily done by peer devices directly accessing each other’s buffers to rotate partial inputs or partial results around the network. This, when carefully handled, will allow the Copy Engine (CE) to handle data movement, while the SMs are free to execute their single wave of persistent threadblocks uninterrupted.

This simple approach allows us to transform any existing GEMM kernel into a Distributed GEMM kernel by prepending only a few instructions to the kernel code, which are mainly agnostic to the underlying structure and programming model of the kernel.

If you’re already familiar with CUTLASS, or more specifically CUTLASS 3, you already know that creating a GEMM kernel is a matter of instantiating one (directly in CUTLASS 2, or using Collectives in CUTLASS 3), and wrapping it with the universal GEMM API:

using CollectiveEpilogue = typename CollectiveBuilder<
 ArchTag, OperatorClass,
 TileShape, ClusterShape,
 EpilogueTileType,
 ElementAccumulator, ElementCompute,
 ElementC, LayoutC, AlignmentC,
 ElementD, LayoutD, AlignmentD,
 EpilogueSchedule
  >::CollectiveOp;

using CollectiveMainloop = typename CollectiveBuilder<
 ArchTag, OperatorClass,
 ElementA, LayoutA, AlignmentA,
 ElementB, LayoutB, AlignmentB,
 ElementAccumulator,
 TileShape, ClusterShape,
 StageCount,
 KernelSchedule
  >::CollectiveOp;

using GemmKernel = GemmUniversal<
  ProblemShapeProfile,
  CollectiveMainloop,
  CollectiveEpilogue>;

using Gemm = GemmUniversalAdapter<GemmKernel>;

And in this case, all you would need to do is declare a Gemm type operator, instantiate with your operand pointers and other runtime parameters, and run.

But what if you were interested in distributing this among multiple GPUs (let’s say 8 for instance) on an NVLink-based network?
With Distributed GEMM it’s very simple. First you would need to define your tensor parallelism strategy:

using DistSchedule = AllGather1D_TilingCD_RotatingA<_8>;

In this case, we’re doing an AllGather + GEMM operation which rotates the left hand side GEMM operand (A), with single-dimensional tiling. In the context of a Tensor Parallel transformer, this just means operand A will be your activations (what was previously running with sequence parallelism.)

Here’s a list of schedules implemented so far:

  AllGather1D_TilingCD_RotatingA
  AllGather1D_TilingCD_RotatingB
  ReduceScatter1D_TilingA_RotatingC
  ReduceScatter1D_TilingB_RotatingC

Distributed GEMM implements these schedules mainly by providing an interface through which we can define sharding, tiling, and buffering patterns, all through CuTe primitives.
After that, all we really need is to instantiate our GEMM kernel with a Distributed GEMM wrapper:

using DistGemmKernel = DistributedGemmKernelWrapper<GemmKernel, DistSchedule>;

using DistGemm = DistributedGemmUniversalAdapter<DistGemmKernel>;

And that is really most of it! The rest comes down to defining, initializing, and running the distributed GEMM, the same way you would a single-GPU GEMM. There’s obviously some additional logic that you would need, because now we’re setting up a CUDA context and stream on multiple devices instead of a single one, and similarly breaking down one problem into 8. To run a typical single-process single-GPU GEMM in CUTLASS you would write:

Gemm gemm;
assert(gemm.can_implement(args));
gemm.initialize(args, workspace_ptr, stream);
gemm.run(stream);

where args is the operator arguments instance, workspace_ptr is a global memory pointer to scratch space required by the operator, and stream is the cuda stream.

Now with Distributed GEMM, to run a single-process multi-GPU GEMM in CUTLASS you would write:

DistGemm dist_gemm_arr[TP];

for (int i = 0; i < 8; ++i) {
  assert(dist_gemm_arr[i].can_implement(args_arr[i]));
}

for (int i = 0; i < 8; ++i) {
  cudaDeviceSet(i);
  dist_gemm_arr[i].initialize(
    args_arr,
    workspace_ptr_arr,
    exclusive_workspace_ptr_arr,
    i,
    stream_arr[i]);
}

for (int i = 0; i < 8; ++i) {
  cudaDeviceSet(i);
  dist_gemm_arr[i].run(stream_arr[i]);
}

where arguments, workspace pointers, streams, and operator instances are all arrays. There’s also an additional exclusive workspace, which is reserved for synchronization primitives (e.g. barriers) that must last across runs. The space allocated for these must not overlap with any other scratch space in the memory pool, but this is a detail that concerns memory management in frameworks and software integrating this operation.

For more details, you can refer to example 65 in CUTLASS.

Performance

Our setup environment is a single DGX box with 8 H100s with NVL18 (any-to-any), therefore we ran everything with TP=8.

We used problem sizes from fc1 and fc2 linear layers in Llama 70B and 405B, which are listed below:

For our baseline, we consider a few options.

First is a single-stage Distributed GEMM which would have fully exposed communication, as illustrated in Figure 3. However, given that this schedule is not yet implemented in the Distributed GEMM framework, we only run a single GEMM excluding the communication. Since we would shard across output columns (N) for AG + GEMM, and across the contracting dimension (K) for GEMM + RS, the problem shapes for this baseline are as follows:

Our second baseline is the same GEMM shape that Distributed GEMM would run at each stage. Since we cut the GEMMs into TP stages in order to do point-to-point communication and attempt to fully hide communication behind computation, the shapes would be further divided by TP along tokens (M):

We run these two baselines directly through the CUTLASS profiler, which runs a number of different kernels that can support these problem shapes, through which we can figure out the best GEMM kernel for these shapes, and its performance level.

Since Distributed GEMM is not yet supported in the CUTLASS profiler, we chose the same kernel configuration for it as the second option. While this is a useful proxy, DistGEMMs must still be profiled holistically in order to get the best performance.

After profiling each problem shape (all assuming FP16), the following kernel configuration was among the best-performing ones for all four cases: Ping-ponging warp-specialized kernel with TMA epilogue, 128x256x64 tile shape, and 1x2x1 cluster size.

The table below reports best achievable performance according to the CUTLASS profiler when running: 1. A single-GPU GEMM, 2. Single-stage GEMM, and 3. Pipelined GEMM. All three exclude communication and are run on a single GPU.

The final column reports the performance of a fully functional DistGEMM from the example in CUTLASS, with the kernel configuration above.

All instances were profiled for approximately 1 second, with approximately 1 second of warm-up, in order to measure sustainable performance.

DistGEMM achieves 91%, 99%, 85%, and 93% of the best performance achievable by a single-stage GEMM (excluding communication) in the four cases.

It also achieves 92%, 93%, 89%, and 91% of the best performance achievable by an 8-stage GEMM (also excluding communication) in the four cases.

While we see that local GEMM kernels (those without communication) typically reach around 80–85% of peak performance, DistGEMM appears to reach only 70–80% of peak performance. We believe that this is mainly due to overheads from synchronization and communication, which could potentially be minimized with a DistGEMM profiler.

Based on our estimates of collective communications in state-of-the-art Tensor Parallelism solutions, DistGEMM could achieve up to 15% speedup in the 405B cases, and up to 24% speedup in the 70B cases. DistGEMM is also faster than the All Reduce setting (All Reduce + GEMM + GEMM instead of All Gather + GEMM + GEMM + Reduce Scatter) could achieve up to 8% and 12% speedup in the 405B and 70B cases. However, note that these are still operation-level comparisons, and not from running an entire LLM or a layer of an LLM.

Conclusion

In conclusion, we discussed the role of Tensor Parallelism in today’s AI infrastructure, and its potential with cutting-edge systems such as the GB200 NVL72. While Distributed GEMM is not yet integrated into any major framework or library, we hope that its simplicity in design, extensibility, and potential to fully overlap communication and computation, leading to massive performance gains in certain communication-bound scenarios, will lead to its adoption in deep learning frameworks and linear algebra libraries.

What we’ll be working on in the near term will involve improving the existing implementation, extending it to support more schedules, and setting up the basic infrastructure required for profiling and measuring its performance.

If you have any questions or suggestions, feel free to join the discussion on our pull request.

Acknowledgements

We’d like to thank Cris Cecka, Amr Elmeleegy, Alan Kaatz, Jack Kosaian, dePaul Miller, Brian Pharris, Richard Vuduc, Bing Xu, and Xiuxia Zhang for their feedback.

References

1. Shoeybi, Mohammad, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. “Megatron-lm: Training multi-billion parameter language models using model parallelism.” arXiv preprint arXiv:1909.08053 (2019).
2. Korthikanti, Vijay Anand, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. “Reducing activation recomputation in large transformer models.” Proceedings of Machine Learning and Systems 5 (2023).
3. Zhao, Yanli, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright et al. “PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel.” Proceedings of the VLDB Endowment 16, no. 12 (2023): 3848–3860.
4. Vaswani, Ashish, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. “Attention is all you need.” Advances in Neural Information Processing Systems (2017).
5. Wang, Shibo, Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ilbeyi, Blake Hechtman, Dehao Chen et al. “Overlap communication with dependent computation via decomposition in large deep learning models.” In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (2022).
6. Jangda, Abhinav, Jun Huang, Guodong Liu, Amir Hossein Nodehi Sabet, Saeed Maleki, Youshan Miao, Madanlal Musuvathi, Todd Mytkowicz, and Olli Saarikivi. “Breaking the computation and communication abstraction barrier in distributed machine learning workloads.” In Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (2022).
7. Chang, Liwen, Wenlei Bao, Qi Hou, Chengquan Jiang, Ningxin Zheng, Yinmin Zhong, Xuanrun Zhang et al. “FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion.” arXiv preprint arXiv:2406.06858 (2024).
8. Wang, Yifu, and Horace He, Less Wright, Luca Wehrstedt, Tianyu Liu, Wanchao Liang. “[Distributed w/ TorchTitan] Introducing Async Tensor Parallelism in PyTorch.” (2024).

Distributed GEMM was originally published in SHI Labs on Medium, where people are continuing the conversation by highlighting and responding to this story.