Architectures of Scale: A Comprehensive Analysis of Multi-GPU Memory Management and Communication Optimization for Distributed Deep Learning

Posted on by uplatzblog

Section 1: The Scalability Imperative in Modern Deep Learning

1.1 The Exponential Growth of Model Complexity

The field of artificial intelligence, particularly deep learning, has been characterized by a relentless pursuit of scale. In recent years, the prevailing trend has demonstrated a strong correlation between the size of neural network models and their performance on a wide array of tasks.1 In Natural Language Processing (NLP), this trend has been particularly pronounced, with models like BERT-large (0.3 billion parameters), GPT-2 (1.5 billion), Megatron-LM (8.3 billion), and T5 (11 billion) setting new benchmarks in language understanding and generation.1 This paradigm, where increasing parameter counts leads to significant accuracy gains, has propelled the development of models that now contain hundreds of billions, and even trillions, of parameters.4

This exponential growth in model size is not merely an academic exercise; it is a direct response to the increasing complexity of tasks that AI systems are expected to perform. Larger models possess a greater capacity to learn intricate patterns and absorb vast amounts of information from massive datasets, leading to superior performance and emergent capabilities not seen in their smaller predecessors.1 However, this scaling has introduced a formidable set of engineering challenges, pushing the boundaries of existing hardware and training methodologies.

 

1.2 The Twin Challenges: Memory and Communication

 

The effort to train these colossal models has exposed two fundamental bottlenecks in modern computing infrastructure: memory capacity and inter-processor communication.2

First, the memory bottleneck arises because the complete set of data required to train a model—including its parameters, the gradients computed during backpropagation, the states maintained by the optimizer, and the intermediate activations from the forward pass—can easily exceed the memory capacity of a single accelerator device, such as a GPU.2 For instance, a trillion-parameter model stored in standard 16-bit precision would require at least 2 terabytes of memory for the parameters alone, far beyond the capacity of any single commercially available GPU.1 Simply adding more devices does not inherently solve this problem if each device is required to hold a full copy of the model, a limitation inherent in traditional distributed training approaches.1

Second, the communication bottleneck emerges as a direct consequence of distributing the training workload across multiple devices. To maintain a coherent training process, these distributed workers must frequently synchronize their state, typically by exchanging gradients or model parameters. This communication overhead can become the primary performance-limiting factor, especially as the number of devices grows.7 If the communication protocol is not highly optimized for the underlying hardware topology, the time spent waiting for data to be exchanged between devices can eclipse the time spent on actual computation, thereby diminishing or even negating the benefits of parallelization.7

 

1.3 Report Objectives and Structure

 

Addressing these twin challenges has become a central focus of research and development in high-performance computing and machine learning systems. A sophisticated ecosystem of techniques has emerged, spanning memory optimization algorithms, specialized communication libraries, and integrated deep learning frameworks.

This report provides a comprehensive, expert-level analysis of the state-of-the-art solutions for multi-GPU memory management and communication optimization in the context of distributed deep learning. The objective is to deliver a deeply technical and cohesive overview that connects foundational concepts to advanced implementations and practical performance tuning. The report is structured as follows:

  • Section 2 establishes the foundational paradigms of distributed training, including data, pipeline, and tensor parallelism.
  • Section 3 delves into advanced memory management architectures, with a focus on Activation Checkpointing and the Zero Redundancy Optimizer (ZeRO) family of techniques.
  • Section 4 provides a detailed examination of the NVIDIA Collective Communications Library (NCCL), its core primitives, and the underlying Ring and Tree communication algorithms.
  • Section 5 serves as a practical guide to performance tuning, exploring the impact of hardware interconnects and software-level configurations on NCCL performance.
  • Section 6 offers a comparative analysis of how these techniques are implemented in leading frameworks, specifically PyTorch Fully Sharded Data Parallel (FSDP) and Microsoft DeepSpeed.
  • Section 7 outlines a systematic workflow for profiling and diagnosing performance bottlenecks in distributed training jobs using tools like the PyTorch Profiler and NVIDIA Nsight Systems.
  • Section 8 concludes with a synthesis of the key strategies and provides a forward-looking perspective on the future of scalable model training.

Through this structured exploration, this report aims to equip researchers and engineers with the knowledge necessary to navigate the complex landscape of large-scale distributed training, enabling them to build, optimize, and scale the next generation of deep learning models efficiently.

 

Section 2: Paradigms of Distributed Training

 

To overcome the limitations of a single accelerator, the training workload must be parallelized across multiple devices. The strategies for this parallelization can be broadly categorized into data parallelism, where the data is split, and model parallelism, where the model itself is split. Advanced training regimes often employ a hybrid of these approaches to maximize efficiency.8

 

2.1 Data Parallelism (DP): The Foundational Approach

 

Data parallelism is the most common and conceptually straightforward approach to distributed training.8 It is the preferred method when the model can fit into the memory of a single GPU, but the dataset is large, and the goal is to accelerate training by processing more data in parallel.8

 

Mechanism

 

The data parallelism workflow proceeds in several distinct steps for each training iteration 11:

  1. Model Replication: An identical copy of the neural network model, with the same initial weights, is loaded onto each participating worker device (e.g., GPU).12
  2. Data Sharding: The global batch of training data is partitioned into smaller, equal-sized mini-batches or “shards.” Each worker receives a unique shard of the data.11
  3. Parallel Forward and Backward Pass: Each worker independently performs a forward pass on its data shard to compute predictions and the loss. Subsequently, it performs a backward pass to compute the gradients of the loss with respect to its local copy of the model parameters. During these passes, there is no communication between workers.12
  4. Gradient Synchronization: This is the critical communication step. The gradients computed on each worker, which are different due to the different data shards, must be aggregated across all workers. This ensures that every worker can perform an identical update to its model parameters, keeping the model replicas synchronized.11

 

The Role of All-Reduce

 

The aggregation of gradients in synchronous data parallelism is almost universally accomplished using the All-Reduce collective communication primitive.7 The

All-Reduce operation performs two functions: first, it applies a reduction operation (typically a sum) to the input buffers from all participating workers; second, it distributes the final, reduced result back to all workers. Consequently, after the All-Reduce operation completes, every GPU holds an identical copy of the globally summed gradients, as if the entire batch had been processed on a single, massive device.7 This synchronized gradient is then used by each worker’s optimizer to update its local model weights, ensuring all model replicas remain identical for the next training step.12

 

Synchronous vs. Asynchronous DP

 

The process described above is known as synchronous data parallelism, where all workers must complete their gradient computation and participate in the All-Reduce before any model updates occur. This is the standard approach in modern deep learning frameworks due to its stability and predictable convergence properties.11 An alternative, asynchronous data parallelism, allows workers to update a central parameter server independently without waiting for others. While this can improve hardware utilization in heterogeneous environments, it can suffer from “stale gradients,” where a worker computes gradients based on an older version of the model parameters, potentially leading to less stable convergence.11

 

Limitations

 

The primary drawback of data parallelism is its inherent memory redundancy. Each of the N GPUs in the data-parallel group must store a full copy of the model parameters, gradients, and optimizer states.14 This means that the maximum model size that can be trained is still fundamentally limited by the memory capacity of a single GPU, regardless of how many GPUs are available. This limitation is the primary motivation for the development of model parallelism techniques.

 

2.2 Model Parallelism: Scaling Beyond Single-GPU Memory

 

When a model is too large to fit into the memory of a single GPU, data parallelism is no longer viable. In such cases, model parallelism becomes necessary. This paradigm involves partitioning the model itself—its layers, parameters, and computations—across multiple devices.8

 

2.2.1 Pipeline Parallelism (Inter-Layer Parallelism)

 

Pipeline parallelism partitions a model vertically, assigning sequential layers or “stages” of the model to different devices.8 The output activations from one stage are passed as input to the next stage, which resides on a different GPU.

A naive implementation of this approach is highly inefficient due to the “pipeline bubble.” In this scenario, only one GPU is active at any given time as a single batch of data flows sequentially through the stages. The other GPUs remain idle, waiting to receive activations from the previous stage or to pass their results to the next.9

To mitigate this inefficiency, modern pipeline parallelism implementations employ micro-batching. The global data batch is split into multiple smaller micro-batches, which are fed into the pipeline in a staggered fashion. This creates an “assembly line” effect, where multiple GPUs can be active simultaneously, each processing a different micro-batch for its assigned stage. While this significantly reduces the idle time, a bubble of inefficiency still exists at the beginning (“ramp-up”) and end (“ramp-down”) of each global batch, where the pipeline is not yet full.8

 

2.2.2 Tensor Parallelism (Intra-Layer Parallelism)

 

Tensor parallelism addresses the scenario where even a single layer of a model is too large to fit on one GPU. Instead of partitioning between layers, this technique partitions the tensors (i.e., the weight matrices and activations) within a single layer horizontally across multiple devices.9

The seminal work in this area is NVIDIA’s Megatron-LM, which developed an efficient and scalable method for applying tensor parallelism to Transformer models.2 In a standard Transformer MLP block, which consists of two linear layers, the weight matrix of the first linear layer is split column-wise across the GPUs, and the weight matrix of the second linear layer is split row-wise.

This partitioning scheme requires specific communication patterns to ensure the mathematical correctness of the computation. After the first (column-parallel) linear layer, an All-Gather operation is needed to collect the distributed output activations from all GPUs. After the second (row-parallel) linear layer, an All-Reduce operation is required to sum the partial results from each GPU to produce the final, correct output.17 This intricate dance of sharded computation and collective communication allows for the execution of massive layers that would be impossible on a single device.

 

2.2.3 Sequence Parallelism

 

Sequence parallelism is a more recent and advanced form of model parallelism that works in conjunction with tensor parallelism to further reduce memory consumption, specifically targeting the activation memory.14 It achieves this by partitioning the input data along the sequence dimension. In Transformer layers that are not tensor-parallel (like LayerNorm), each GPU computes its portion of the sequence independently. In tensor-parallel layers, an

All-to-All communication is used to distribute sequence chunks and attention heads, allowing each GPU to compute a subset of the attention mechanism before another All-to-All gathers the results.14 This technique is particularly effective for training models with very long sequence lengths.22

 

2.3 Hybrid Strategies for Extreme Scale

 

Training the largest state-of-the-art models requires combining these different forms of parallelism into a hybrid strategy, often referred to as 3D or 4D parallelism.8 A common and effective configuration leverages the hierarchical nature of modern supercomputing clusters:

  1. Tensor Parallelism (Intra-Node): Tensor parallelism is communication-intensive, requiring frequent, low-latency, high-bandwidth collectives (All-Reduce, All-Gather). It is therefore best suited for GPUs within a single server node that are connected by ultra-fast interconnects like NVIDIA NVLink and NVSwitch.9
  2. Pipeline Parallelism (Inter-Node): Pipeline parallelism involves less frequent but still substantial communication (passing activations between stages). This makes it suitable for communication between nodes over a high-speed network fabric like InfiniBand.18
  3. Data Parallelism (Global): Data parallelism is applied over the entire group of devices. A single replica of the model (which is itself parallelized across GPUs using tensor and pipeline parallelism) is placed on each data-parallel worker. This allows for scaling the global batch size and increasing overall training throughput.18

This hybrid approach allows each dimension of parallelism to be mapped to the hardware topology it is best suited for, enabling the training of models with trillions of parameters across thousands of GPUs.

Strategy What is Split? Memory Distribution Primary Communication Pattern GPU Utilization Profile Key Advantage Key Disadvantage
Data Parallelism Data Batch Replicated Model States All-Reduce (Gradients) High, with synchronous waits for the slowest worker. Simple to implement; scales throughput. Memory redundancy; model size limited by single GPU memory.
Pipeline Parallelism Model Layers Partitioned Layers Point-to-Point (Activations) Suffers from “pipeline bubble” (idle time). Enables models larger than a single GPU; less communication than TP. Bubble inefficiency; can be complex to balance stages.
Tensor Parallelism Tensors within Layers Partitioned Tensors All-Gather, All-Reduce High, but communication-heavy; sensitive to interconnect speed. Enables layers larger than a single GPU; very memory efficient. High communication volume; complex to implement.

 

Section 3: Advanced Memory Management Architectures

 

While model parallelism provides a path to scale beyond the memory of a single GPU, it does not address the memory redundancy inherent in data parallelism or the significant memory overhead from activations. To tackle these issues, a new class of memory optimization techniques has been developed, fundamentally changing how model states are stored and managed during training.

 

3.1 Deconstructing Memory Consumption

 

To understand these advanced techniques, it is essential to first dissect the memory footprint of a standard training process. For each trainable parameter in a model, the GPU must store several pieces of data 13:

  1. Model Parameters: These are the weights and biases of the model. In mixed-precision training, this typically involves a 16-bit floating-point (FP16 or BF16) copy used for computation.
  2. Gradients: After the backward pass, a 16-bit gradient is computed for each parameter.
  3. Optimizer States: This is often the largest component of memory usage. The widely used Adam optimizer, for example, maintains two states for each parameter: a 32-bit momentum and a 32-bit variance. Furthermore, for stable mixed-precision training, the optimizer often works on a 32-bit (FP32) copy of the parameters. In total, this amounts to 12 bytes of optimizer state per model parameter (4 for FP32 params + 4 for momentum + 4 for variance).25
  4. Activations: These are the intermediate outputs of each layer from the forward pass. They must be stored in memory because they are required for calculating gradients during the backward pass. The size of the activations scales with the batch size, sequence length, and the number of layers in the model.13

A critical realization is that for a model trained with Adam in mixed precision, the optimizer states alone consume three times more memory than the FP32 model parameters and six times more than the FP16 parameters used for computation. This disproportionate memory usage, combined with the large memory footprint of activations, makes these two components the primary targets for optimization. The following techniques, Activation Checkpointing and the Zero Redundancy Optimizer, were designed specifically to address these two dominant sources of memory consumption.

 

3.2 Activation Checkpointing (Gradient Checkpointing): Trading Compute for Memory

 

Activation checkpointing, also known as gradient checkpointing, is a technique that reduces the memory required for storing activations by trading it for additional computation.26

 

The Core Idea and Mechanism

 

In standard backpropagation, all activations generated during the forward pass are stored in GPU memory. This is because the chain rule requires these values to compute the gradients during the backward pass. For a deep network with n layers, the memory required to store these activations scales linearly with the depth of the network, i.e., O(n).26

Activation checkpointing breaks this linear dependency. Instead of saving all activations, it strategically saves only a small subset of them, referred to as “checkpoints.” The activations for the layers between these checkpoints are discarded during the forward pass to save memory. Then, during the backward pass, when the gradients for a non-checkpointed segment are needed, the activations for that segment are recomputed on-the-fly, starting from the nearest preceding checkpoint.27

 

Theoretical Improvement and Practical Implementation

 

This trade-off is highly favorable. For a simple feed-forward network, an optimal checkpointing strategy (e.g., checkpointing every n​ layers) reduces the memory cost for activations from O(n) to O(n​).26 The computational overhead incurred by this strategy is equivalent to performing one extra forward pass through the model, as each non-checkpointed layer is recomputed exactly once during the backward pass.27 This sublinear memory cost allows for the training of significantly deeper models or the use of larger batch sizes than would otherwise be possible.27 This technique has proven essential for training very large models and is readily available in major frameworks, such as through

torch.utils.checkpoint in PyTorch.2

 

3.3 The Zero Redundancy Optimizer (ZeRO): Eliminating Memory Redundancy

 

While activation checkpointing addresses the memory cost of activations, the Zero Redundancy Optimizer (ZeRO), developed by Microsoft as part of the DeepSpeed library, targets the memory redundancy of parameters, gradients, and optimizer states in data parallelism.4 Instead of replicating these model states on every GPU, ZeRO partitions them across the data-parallel group, ensuring that each GPU stores only a fraction of the total state.1 ZeRO is implemented in three incremental stages, each offering greater memory savings.

 

3.3.1 ZeRO Stage 1: Partitioning Optimizer States

 

ZeRO Stage 1 targets the largest source of memory redundancy in mixed-precision training: the optimizer states.

  • Mechanism: The 32-bit optimizer states (FP32 parameters, momentum, and variance) are partitioned evenly across the data-parallel GPUs. Each GPU is responsible for updating only its assigned partition of the parameters. During the optimizer step, after gradients have been all-reduced, each GPU updates its local partition. An All-Gather operation is then performed to distribute the updated FP16 parameters to all GPUs so they have a complete, synchronized model for the next forward pass.33
  • Memory Savings: This stage reduces the memory required for the optimizer states by a factor of Nd​, where Nd​ is the data-parallel degree. For mixed-precision Adam, this results in an overall memory reduction of approximately 4x compared to standard data parallelism.1 The communication volume remains identical to standard data parallelism.1

 

3.3.2 ZeRO Stage 2: Partitioning Gradients and Optimizer States

 

ZeRO Stage 2 builds upon Stage 1 by also partitioning the 16-bit gradients.

  • Mechanism: During the backward pass, instead of performing an All-Reduce operation to make the full gradient tensor available on all GPUs, a Reduce-Scatter operation is used. This operation both sums the gradients across all GPUs and distributes them in a partitioned manner, so each GPU receives only the shard of the final gradient corresponding to its partition of the optimizer states. This avoids the need for each GPU to ever store the full gradient tensor.33
  • Memory Savings: By eliminating the redundant storage of gradients, Stage 2 doubles the memory savings of Stage 1, resulting in an 8x reduction compared to standard data parallelism.1 The communication volume remains the same as standard data parallelism.1

 

3.3.3 ZeRO Stage 3: Full Model State Partitioning

 

ZeRO Stage 3 achieves the maximum level of memory savings by partitioning the 16-bit model parameters themselves.

  • Mechanism: Each GPU only holds a partition of the parameters at all times. During the forward and backward passes, just before a layer is executed, the full parameters for that specific layer are dynamically reconstructed on all GPUs via an All-Gather operation. Immediately after the computation is complete, the unsharded parameters are discarded, freeing the memory.5
  • Memory Savings: The memory required for model states is now partitioned across all data-parallel GPUs, leading to a memory reduction that scales linearly with the number of GPUs. This enables the training of models with trillions of parameters.1
  • Communication: This stage introduces a significant amount of communication, as All-Gather operations are required for each layer during both the forward and backward passes. This makes Stage 3 highly dependent on fast interconnects but offers the ultimate scalability in model size.5

 

3.4 Beyond GPU Memory: ZeRO-Offload and ZeRO-Infinity

 

To push the boundaries of model scale even further, DeepSpeed introduced extensions to ZeRO that leverage system memory beyond the GPU.

  • ZeRO-Offload: This technique offloads the partitioned optimizer states and gradients (from ZeRO-2) or parameters (from ZeRO-3) from the GPU memory to the host CPU’s main RAM.35 The parameter update step, which is computationally less intensive than the forward/backward passes, can also be offloaded and executed on the CPU. While communication over the PCIe bus is much slower than intra-GPU interconnects, ZeRO-Offload overlaps this communication with the GPU’s computation and uses highly optimized CPU Adam implementations to mitigate the performance impact. This allows for training models up to 13 billion parameters on a single GPU.36
  • ZeRO-Infinity: This is the ultimate extension of the offloading concept. It allows partitioned model states, including the parameters themselves, to be offloaded not just to CPU RAM but also to much larger, albeit slower, NVMe solid-state drives.5 This makes it theoretically possible to train models with trillions of parameters on a cluster with a modest amount of total GPU memory, effectively breaking the GPU memory wall.4
ZeRO Stage Optimizer States Gradients Parameters Key Communication Pattern Theoretical Memory Reduction
Stage 0 (DP) Replicated Replicated Replicated All-Reduce 1x
Stage 1 Partitioned Replicated Replicated All-Reduce (Gradients), All-Gather (Params) ~4x
Stage 2 Partitioned Partitioned Replicated Reduce-Scatter (Gradients), All-Gather (Params) ~8x
Stage 3 Partitioned Partitioned Partitioned All-Gather (Params), Reduce-Scatter (Gradients) Linear (Nd​)

 

Section 4: The NVIDIA Collective Communications Library (NCCL): The Backbone of Multi-GPU Communication

 

The parallelism and memory management strategies described in the previous sections rely heavily on efficient, high-performance communication between GPUs. The NVIDIA Collective Communications Library (NCCL) is the de facto standard for implementing these communication primitives on NVIDIA GPU platforms.39 It serves as the high-performance communication backend for virtually all major deep learning frameworks, including PyTorch, TensorFlow, and DeepSpeed.7

 

4.1 Core Architecture and Design Philosophy

 

NCCL is not a general-purpose parallel programming framework like MPI, but rather a specialized library focused exclusively on optimizing inter-GPU communication.42 Its design philosophy is centered on several key principles:

  • Performance: NCCL is optimized to achieve the maximum possible bandwidth and lowest latency over a variety of hardware interconnects, including PCIe, NVLink, NVSwitch, and InfiniBand.39
  • Topology Awareness: A critical feature of NCCL is its ability to automatically detect the underlying hardware topology of the system. It builds a graph of the GPUs and their interconnects and uses this information to select the optimal communication path and algorithm for any given operation.39
  • Ease of Integration: NCCL provides a simple C API that closely follows the well-established MPI standard for collective operations, making it easy for framework developers to integrate.39
  • Compatibility: It is designed to work with various parallel programming models, including single-threaded control of multiple GPUs, multi-threaded applications (one thread per GPU), and multi-process applications using MPI.42

 

4.2 Anatomy of Collective Primitives

 

NCCL implements a set of standard collective communication primitives, which are fundamental building blocks for distributed training algorithms.42 The most critical primitives for the strategies discussed in this report are:

  • ncclAllReduce: This operation takes an input buffer from each of the N participating GPUs (ranks), performs a reduction operation (e.g., sum, max, min) across all input buffers element-wise, and writes the final, identical result to the output buffer on all N GPUs. This is the cornerstone of traditional data parallelism for averaging gradients.7
  • ncclAllGather: In this operation, each of the N GPUs contributes a buffer of data. The operation gathers all of these buffers and concatenates them, distributing the final, complete buffer (of size N times the input buffer size) to all participating GPUs. This is essential for ZeRO-3 and FSDP, where partitioned parameters must be reconstructed on each GPU before a layer’s computation can proceed.7
  • ncclReduceScatter: This primitive combines a reduction and a scatter operation. It performs an element-wise reduction across the input buffers from all N GPUs, but instead of distributing the entire result, it scatters the result vector into N chunks and delivers a unique chunk to each GPU. This is the key communication pattern for ZeRO-2 and ZeRO-3, allowing for the efficient reduction and partitioning of gradients without ever requiring a single GPU to hold the full gradient tensor.7
  • ncclBroadcast: A one-to-many operation where a single root GPU sends its buffer to all other GPUs in the communicator.7
  • ncclReduce: A many-to-one operation where data from all GPUs is reduced, and the final result is stored only on the specified root GPU.7

 

4.3 Communication Algorithms: Ring vs. Tree

 

For each collective primitive, NCCL can employ different underlying algorithms to execute the communication. The choice of algorithm is made dynamically by NCCL’s internal cost model and is critical for performance.46 The two primary algorithms are Ring and Tree.

 

Ring Algorithm

 

The Ring algorithm is optimized for maximizing bandwidth, especially for large data transfers.

  • Mechanism: The GPUs are arranged in a logical ring, where each GPU only sends data to and receives data from its immediate neighbors. The data is broken down into chunks. In an All-Reduce operation, each GPU sends a chunk to its neighbor while simultaneously receiving a chunk from its other neighbor. It adds the received chunk to its own corresponding chunk and forwards the result. This process repeats 2(N−1) times, with data circulating the ring in a pipelined fashion until all GPUs have a copy of the final, fully reduced data.48
  • Performance Characteristics: The Ring algorithm is bandwidth-optimal because at steady state, all links in the ring are fully utilized. However, the total time to complete the operation includes a latency component that scales linearly with the number of GPUs (N). This makes it highly efficient for large messages where the transfer time dominates the latency, but inefficient for small messages or on very large-scale systems where the linear latency becomes a bottleneck.49

 

Tree Algorithm

 

The Tree algorithm is optimized for minimizing latency, making it ideal for small messages and large numbers of GPUs.

  • Mechanism: For a Reduce operation, GPUs are arranged in a hierarchical binary tree. Data flows up the tree from the leaves to the root, with intermediate nodes performing partial reductions. For a Broadcast, data flows down from the root to the leaves. An All-Reduce is typically implemented as a Reduce followed by a Broadcast. To maximize bandwidth, NCCL often uses two simultaneous binary trees (a “double binary tree”), where each tree handles half of the data.48
  • Performance Characteristics: The number of steps in a tree-based operation is proportional to the depth of the tree, which is logarithmic with respect to the number of GPUs (log(N)). This logarithmic latency scaling makes the Tree algorithm vastly superior to the Ring algorithm for latency-sensitive operations (i.e., small messages) and at massive scales where the Ring’s linear latency would be prohibitive.49 However, tree-based algorithms can sometimes lead to network congestion and may not saturate link bandwidth as effectively as the Ring algorithm for large messages.49

The existence of these two algorithms with opposing performance characteristics necessitates a sophisticated decision-making process within NCCL. Neither algorithm is universally superior; the optimal choice depends on a complex interplay of message size, GPU count, and the specific hardware topology. NCCL’s internal cost model evaluates these factors for each collective call to dynamically select the algorithm predicted to yield the best performance. This dynamic selection is a cornerstone of NCCL’s ability to achieve high performance across a wide range of platforms and workloads, and understanding this trade-off is crucial for advanced performance tuning.

 

Section 5: High-Performance NCCL Tuning and Optimization

 

While NCCL’s automatic topology detection and algorithm selection provide excellent performance out-of-the-box for most scenarios, achieving peak performance on complex, large-scale systems often requires manual tuning. This optimization process involves understanding the interplay between the hardware interconnects, NCCL’s internal algorithms, and software-level configuration parameters.43

 

5.1 The Impact of Hardware Interconnects

 

The physical connections between GPUs are the most critical factor determining communication performance. NCCL is designed to exploit the hierarchy of available interconnects.52

  • Intra-Node Interconnects:
  • PCIe: The standard interconnect for connecting GPUs to the CPU and to each other in consumer-grade systems. While modern versions like PCIe Gen4 and Gen5 offer significant bandwidth, communication between GPUs often has to traverse the CPU’s memory controller, adding latency and consuming CPU-to-GPU bandwidth.53
  • NVLink and NVSwitch: These are NVIDIA’s proprietary high-speed interconnect technologies designed for direct GPU-to-GPU communication. NVLink provides a point-to-point connection between pairs of GPUs, while NVSwitch acts as a full crossbar switch, allowing all-to-all communication between GPUs within a node at extremely high bandwidth and low latency.53 These technologies are essential for efficient tensor and pipeline parallelism, as they bypass the CPU and PCIe bus entirely.53 NCCL will always prioritize NVLink paths when available.53
  • Inter-Node Interconnects:
  • Ethernet: While standard TCP/IP over Ethernet can be used, it often becomes a bottleneck for high-performance distributed training due to CPU overhead and the lack of direct GPU memory access.54
  • InfiniBand with GPUDirect RDMA: This is the standard for high-performance inter-node communication. InfiniBand is a high-bandwidth, low-latency network fabric. When combined with GPUDirect RDMA (Remote Direct Memory Access), it allows a GPU on one node to directly read from or write to the memory of a GPU on another node without involving the CPUs on either node. This dramatically reduces latency and CPU overhead, and is critical for scaling training across multiple servers.54
  • Advanced Topologies and PXN: In complex multi-node systems (like NVIDIA’s DGX systems), the topology can be intricate. The PXN (PCIe over NVLink) feature in NCCL allows a GPU to communicate with a network interface card (NIC) that is not on its local PCIe root complex by routing the traffic through another GPU via NVLink. This avoids slower inter-CPU links (like QPI) and ensures that network traffic can leverage the high-speed NVLink fabric, which is crucial for maintaining performance in hierarchical topologies.56

 

5.2 Software-Level Tuning with Environment Variables

 

NCCL exposes a set of environment variables that allow users to override its default behavior. These are powerful tools for debugging, benchmarking, and extracting maximum performance in specific scenarios, but should be used with caution as they can force suboptimal configurations if misused.46

  • NCCL_ALGO: This variable forces NCCL to use a specific algorithm for its collectives. For example, setting NCCL_ALGO=RING will force the use of the ring algorithm, while NCCL_ALGO=TREE will force the tree algorithm. This is invaluable for performance analysis. If a workload with many small messages is performing poorly, one can force the tree algorithm to see if latency was the bottleneck. Conversely, if a large-message workload is not saturating the network, forcing the ring algorithm can confirm if the default choice was suboptimal.46
  • NCCL_PROTO: This variable controls the protocol used for communication. The main choices are Simple and LL (Low Latency) / LL128. The Simple protocol is optimized for high bandwidth on large data transfers, often at the cost of slightly higher latency. The LL protocols are designed to minimize latency for small messages, sometimes by using more GPU resources or different communication patterns.46 Tuning this can help align the communication strategy with the message size profile of the workload.
  • Resource and Buffering Variables:
  • NCCL_MIN_CTAS / NCCL_MAX_CTAS: These variables control the number of CUDA Cooperative Thread Arrays (CTAs) that NCCL uses to drive communication. Increasing the number of CTAs can sometimes improve bandwidth utilization, but can also lead to resource contention with the model’s computation kernels, potentially harming overall performance. NCCL’s default is designed to be conservative to avoid this interference.43
  • NCCL_BUFFSIZE: This sets the size of the internal buffers used for communication. Larger buffers can improve throughput for bandwidth-bound operations but may increase memory usage and latency.46

 

5.3 Advanced Optimization with Tuner Plugins

 

While environment variables provide a global override, tuner plugins are the recommended method for fine-grained, platform-specific optimization.46 A tuner plugin is a shared library that NCCL can load at runtime. It allows a user or platform vendor to programmatically override the decisions of NCCL’s internal cost model.

The primary function in a tuner plugin, getCollInfo, is called by NCCL before executing a collective. The plugin can inspect the parameters of the collective (operation type, message size, number of ranks, etc.) and return a custom cost for different algorithm/protocol combinations. This allows for highly targeted optimizations. For example, a cluster administrator who has extensively benchmarked their specific network fabric could create a plugin that forces the TREE algorithm for All-Reduce operations below a certain message size threshold and the RING algorithm above it, ensuring optimal performance for their hardware without requiring any changes to the end-user’s application code.46

Environment Variable Possible Values Description & Impact When to Use
NCCL_ALGO Ring, Tree Forces a specific collective algorithm. Ring is bandwidth-optimal but has linear latency scaling. Tree has logarithmic latency scaling but may have lower peak bandwidth. For benchmarking and debugging. Use Tree to diagnose latency-bound issues (small messages, large scale). Use Ring to diagnose bandwidth-bound issues.
NCCL_PROTO Simple, LL, LL128 Forces a specific communication protocol. Simple is optimized for large-message bandwidth. LL/LL128 are optimized for small-message latency. To align the protocol with the workload’s message size profile. Use LL for latency-sensitive applications.
NCCL_MIN_CTAS Integer (e.g., 16) Sets the minimum number of CUDA Thread Arrays used by NCCL. Increasing this can sometimes improve bandwidth but risks resource contention with compute kernels. When profiling shows underutilized network/interconnect bandwidth and low GPU compute utilization during communication phases. Use with caution.
NCCL_BUFFSIZE Bytes (e.g., 4194304) Sets the size of internal communication buffers. Larger buffers can increase throughput for large transfers but also increase memory footprint. For tuning bandwidth-bound workloads. Experimentation is required to find the optimal size for a given platform and model.
NCCL_DEBUG INFO, WARN, TRACE Sets the logging level for NCCL. INFO provides useful information on the chosen algorithm, protocol, and topology. TRACE is extremely verbose. For debugging. NCCL_DEBUG=INFO is the first step to understanding which algorithm/protocol NCCL is selecting automatically.

 

Section 6: Frameworks in Practice: A Comparative Analysis of FSDP and DeepSpeed

 

The memory and communication optimization techniques discussed are implemented within high-level deep learning frameworks. PyTorch and the DeepSpeed library are two of the most prominent ecosystems for large-scale model training. While both leverage the same underlying principles (like ZeRO-style sharding and NCCL for communication), their implementations, APIs, and feature sets have important differences.

 

6.1 PyTorch Fully Sharded Data Parallel (FSDP)

 

Fully Sharded Data Parallel (FSDP) is PyTorch’s native solution for large-scale training, implementing the core ideas of the ZeRO optimizer.60

  • Core Design and Sharding Strategies: FSDP works by wrapping model modules (e.g., individual Transformer layers). It partitions the parameters, gradients, and optimizer states of the wrapped module across the data-parallel workers. FSDP offers several sharding strategies that map directly to the ZeRO stages 60:
  • NO_SHARD: Equivalent to standard DistributedDataParallel (DDP) or ZeRO Stage 0.
  • SHARD_GRAD_OP: Shards gradients and optimizer states, equivalent to ZeRO Stage 2.
  • FULL_SHARD: Shards parameters, gradients, and optimizer states, equivalent to ZeRO Stage 3.
  • HYBRID_SHARD: A multi-node strategy that performs full sharding within a node and replicates parameters across nodes, optimizing for hierarchical network topologies.
  • Communication and Overlap: During the forward and backward passes, FSDP orchestrates all-gather collectives to reconstruct the full parameters for a given module just before computation, and reduce-scatter collectives to aggregate and shard gradients after computation.61 To improve performance, FSDP includes mechanisms for prefetching the parameters for the next module while the current module’s computation is still in progress (
    forward_prefetch, backward_prefetch), effectively overlapping communication and computation.63
  • Communication Hooks: A powerful feature of PyTorch’s distributed ecosystem is the concept of communication hooks. Both DDP and FSDP allow users to register a custom function (a “hook”) that intercepts the gradient communication step.66 This enables advanced optimizations, such as:
  • Gradient Compression: Hooks can apply compression algorithms (e.g., casting gradients to FP16/BF16, or more advanced methods like PowerSGD) before the reduce-scatter operation, reducing the amount of data sent over the network at the cost of some precision.66
  • Custom Communication Logic: Users can implement entirely novel communication strategies tailored to their specific research needs.

 

6.2 Microsoft DeepSpeed

 

DeepSpeed is a comprehensive deep learning optimization library that pioneered and popularized the ZeRO family of optimizers.4 It is designed to be easy to use and highly effective at enabling extreme-scale training.

  • ZeRO Implementation: DeepSpeed provides robust and battle-tested implementations of ZeRO Stages 1, 2, and 3, along with the advanced ZeRO-Offload and ZeRO-Infinity features for leveraging CPU and NVMe memory.5
  • Ease of Use: A hallmark of DeepSpeed is its configuration-driven approach. Most of its powerful features, including the choice of ZeRO stage and offloading strategies, can be enabled and configured through a single JSON file, often requiring minimal changes to the user’s PyTorch training script.5
  • Advanced Features and Optimizers: Beyond the core ZeRO algorithm, DeepSpeed includes a suite of other optimizations, such as custom high-performance Adam optimizers (e.g., 0/1 Adam) that integrate communication compression directly into the update step, and specialized inference engines for large models.31

 

6.3 Key Architectural Differences and Use Cases

 

While FSDP and DeepSpeed are built on the same foundational principles, their design choices lead to important practical differences for users.60

  • Integration and Dependencies: FSDP is a native part of the PyTorch library, ensuring seamless integration and compatibility with the core framework. DeepSpeed is an external library that requires separate installation and sometimes patches or specific versions of PyTorch to work correctly.60
  • Configuration and API: FSDP is configured primarily through Python code, using wrapper classes and policy functions to define how the model is sharded. This offers fine-grained control but can be more verbose. DeepSpeed primarily uses a JSON configuration file, which is declarative and can be easier for standard use cases, but may be less flexible for complex, dynamic models.33
  • Offloading Capabilities: DeepSpeed’s offloading features (ZeRO-Offload and ZeRO-Infinity) are generally considered more mature and feature-rich than FSDP’s. DeepSpeed offers granular control over what is offloaded (e.g., optimizer states only, or parameters as well) and supports offloading to both CPU and NVMe, whereas FSDP’s CPU offloading is more of an “all-or-nothing” switch.33
  • Precision Handling: A subtle but important difference lies in how they handle data types. FSDP typically keeps optimizer states in the same precision as the computation (e.g., BF16), while DeepSpeed upcasts them to FP32 for the update step. DeepSpeed’s approach can be more stable but may incur slightly more memory overhead on a small number of GPUs.65

The choice between FSDP and DeepSpeed often depends on the user’s specific needs. FSDP is an excellent choice for users who want a native, tightly integrated PyTorch solution and are migrating from DDP. DeepSpeed is often favored by those who need the most advanced features for pushing the limits of model scale, such as NVMe offloading, or who prefer its configuration-based setup.71

Feature PyTorch FSDP DeepSpeed ZeRO
Sharding Granularity NO_SHARD (DP), SHARD_GRAD_OP (ZeRO-2), FULL_SHARD (ZeRO-3), HYBRID_SHARD Stage 1 (Optimizer), Stage 2 (Grads+Optim), Stage 3 (Params+Grads+Optim)
Offload Support CPU Offload (all-or-nothing) Granular CPU and NVMe Offload (ZeRO-Offload & ZeRO-Infinity)
Checkpointing Full or Sharded state dicts, configured via API. Sharded by default; can consolidate to a single rank for saving.
Configuration Method Python API (wrapping policies, FullyShardedDataParallelPlugin) JSON configuration file
Framework Integration Native to PyTorch External library
Advanced Features Communication hooks for custom gradient compression. Custom optimizers (0/1 Adam), inference engines, ZeRO++.

 

Section 7: Profiling and Diagnosing Distributed Training Performance

 

Optimizing a complex distributed training job is an iterative process that relies on empirical performance data. Identifying whether a workload is compute-bound, memory-bound, or communication-bound is essential for applying the correct optimization strategies. This requires a multi-layered profiling workflow, starting with high-level framework tools and drilling down into system-level details with specialized profilers like NVIDIA Nsight Systems.74

 

7.1 A Multi-Layered Profiling Workflow

 

A best-practice workflow for diagnosing performance issues involves a top-down approach:

  1. High-Level Framework Profiling: The first step is to use the profiler integrated with the deep learning framework, such as torch.profiler for PyTorch. This tool provides an operator-level view of execution, breaking down the time spent in different parts of the training loop (e.g., data loading, forward pass, backward pass, optimizer step). It can trace both CPU and CUDA activities, helping to identify which high-level operations are the most time-consuming. If the profiler shows significant time spent in collective communication operations (e.g., nccl:all_reduce) or large periods of GPU inactivity, it signals a potential communication or synchronization bottleneck that warrants a deeper investigation.74
  2. System-Level Profiling: Once a potential bottleneck is identified, the next step is to use a system-wide profiler like NVIDIA Nsight Systems (nsys). Nsight Systems captures a detailed timeline of events across the entire system, including CPU threads, OS runtime libraries, CUDA API calls, GPU kernel executions, and, crucially, NCCL communication operations. This allows for a precise correlation between high-level framework operations and low-level hardware activity.76

 

7.2 Deep-Dive Analysis with NVIDIA Nsight Systems (nsys)

 

Profiling a distributed PyTorch application with Nsight Systems provides an unparalleled level of detail for diagnosing performance issues.

  • Capturing a Trace: The nsys profile command-line tool is used to launch the distributed training script (e.g., via torchrun or mpirun). It is critical to enable tracing for the relevant APIs using flags like -t cuda,nvtx,osrt,nccl. The –pytorch=autograd-nvtx flag is also highly recommended, as it automatically inserts NVTX (NVIDIA Tools Extension) ranges around PyTorch operations, making the resulting trace much easier to interpret.74 To focus the profiling on the steady-state training loop and avoid capturing initialization overhead, it is common practice to use warmup iterations and trigger the profiler programmatically using
    torch.cuda.profiler.start() and stop() within the training script, in conjunction with the –capture-range=cudaProfilerApi nsys flag.76
  • Analyzing the Timeline: The captured .nsys-rep file is visualized in the Nsight Systems GUI. A typical trace will display several horizontal rows, each representing a timeline of events for a specific component 76:
  • CPU Rows: Show thread states, OS runtime events, and Python call stack samples, which can help identify CPU-bound operations or data loading bottlenecks.
  • NVTX Rows: Display the NVTX ranges, providing a high-level semantic view of the training loop (e.g., “forward”, “backward”, “optimizer_step”). These are essential for navigating the trace and understanding the context of low-level events.76
  • GPU Rows: Show the execution of CUDA kernels on the GPU’s streaming multiprocessors and memory copy operations. Gaps in this timeline indicate periods of GPU inactivity, which are prime targets for optimization.
  • NCCL Rows: This row is critical for communication analysis. It visualizes the execution of NCCL collective primitives, showing when they are launched and how long they take to complete.77

 

7.3 Identifying Communication Bottlenecks in Nsight Traces

 

By correlating information across these timelines, an engineer can pinpoint the root cause of performance issues. A common workflow involves observing a high-level symptom and drilling down to the underlying cause. For example, if the GPU utilization is low, the Nsight trace can reveal why. A frequent pattern for a communication bottleneck is a gap in the GPU compute kernel timeline that perfectly aligns with a long-running NCCL operation in the NCCL timeline.

This observation can then be refined:

  1. Identify the Operation: Using the NVTX ranges, determine which part of the training step is triggering the long communication call. Is it an all-gather during the forward pass of an FSDP-wrapped module? Or a reduce-scatter during the backward pass?
  2. Analyze the Collective: Examine the properties of the NCCL call. What is the message size? How many GPUs are participating? This information is available in the profiler.
  3. Formulate a Hypothesis: This empirical data allows for the formation of an evidence-based hypothesis. For instance, if a collective with a very small message size is taking a long time across a large number of nodes, this strongly suggests a latency-bound scenario. The default Ring algorithm, with its linear latency scaling, is a likely culprit.
  4. Apply a Targeted Fix and Re-profile: Based on this hypothesis, the engineer can take a targeted action, such as setting the NCCL_ALGO=TREE environment variable to force the more latency-optimal tree algorithm. A new profile is then captured to verify if the communication stall has been reduced and if overall throughput has improved.

Other common communication-related performance patterns to look for in Nsight traces include:

  • Insufficient Computation-Communication Overlap: Gaps between the end of a compute kernel and the start of a NCCL kernel (or vice-versa) indicate that the framework is not effectively hiding communication latency behind computation. This might be solvable by adjusting prefetching settings in FSDP or buffer sizes in DeepSpeed.
  • Small Message Inefficiency: A “sawtooth” pattern of many small, rapid-fire NCCL calls can be inefficient due to the overhead of launching each kernel. This might indicate that the communication bucketing size is too small or that gradient accumulation should be used to increase the effective batch size and amortize the communication cost over more computations.
  • Network Saturation: If large-message collectives are slow, correlating the Nsight trace with network counter metrics (which can also be collected by nsys) can confirm if the physical network bandwidth is the limiting factor.

This systematic, data-driven approach, moving from high-level symptoms to low-level trace analysis and targeted tuning, is fundamental to achieving high performance in large-scale distributed training.

 

Section 8: Conclusion and Future Directions

 

The rapid evolution of deep learning models towards unprecedented scale has necessitated a parallel evolution in the systems and software used to train them. The constraints of single-GPU memory and the performance bottlenecks of inter-device communication have been the primary drivers of innovation, leading to a sophisticated stack of techniques for distributed training. This report has provided a comprehensive analysis of these techniques, from foundational parallelism paradigms to advanced memory management architectures and the intricacies of communication library tuning.

 

Summary of Key Strategies

 

The analysis reveals a clear hierarchy of solutions tailored to different scales and constraints.

  • Beyond Data Parallelism: While simple data parallelism remains effective for smaller models, the memory redundancy it creates makes it untenable for state-of-the-art research. The move towards partitioning model states is a fundamental requirement for training large models.
  • The Power of ZeRO: The Zero Redundancy Optimizer (ZeRO) and its native PyTorch implementation, FSDP, represent the state of the art in memory-efficient data parallelism. By systematically partitioning optimizer states, gradients, and finally parameters, these techniques allow model size to scale linearly with the number of available devices. Extensions like ZeRO-Offload further push these boundaries by integrating CPU and NVMe memory into the memory hierarchy.
  • NCCL as the Communication Substrate: Efficient execution of these sharding strategies is entirely dependent on a high-performance communication backend. The NVIDIA Collective Communications Library (NCCL) provides this critical layer, offering topology-aware, optimized implementations of the necessary collective primitives (All-Reduce, All-Gather, Reduce-Scatter).
  • The Importance of Profiling: Achieving optimal performance is not automatic. It requires a deep understanding of the interactions between the model architecture, the parallelism strategy, the communication library, and the underlying hardware. A systematic profiling workflow, using tools like the PyTorch Profiler and NVIDIA Nsight Systems, is indispensable for identifying and resolving bottlenecks, whether they lie in computation, memory access, or inter-GPU communication.

 

Recommendations for Practitioners

 

For engineers and researchers embarking on a large-scale training project, the choice and configuration of these strategies can be daunting. The following decision-making framework, based on the findings of this report, can serve as a practical guide:

  1. Assess Model Size vs. GPU Memory:
  • If the model and its associated training states fit comfortably on a single GPU but training is too slow, standard Data Parallelism (DDP) is the simplest and often fastest solution.
  • If the model does not fit on a single GPU, a sharding strategy is required. Start with ZeRO Stage 2 / FSDP SHARD_GRAD_OP. This offers a significant 8x memory reduction with a communication pattern that is often more efficient than full parameter sharding.
  1. Escalate Memory Savings as Needed:
  • If ZeRO-2/FSDP is still insufficient, move to ZeRO Stage 3 / FSDP FULL_SHARD. This provides the maximum on-GPU memory savings but introduces more communication overhead. This stage is highly sensitive to interconnect bandwidth.
  • If the model still does not fit within the aggregate GPU memory of a single node, or if you are on a resource-constrained system, utilize ZeRO-Offload to leverage CPU memory. For the most extreme scales, ZeRO-Infinity with NVMe offloading is the final step.
  1. Profile and Tune for Performance:
  • Always begin by establishing a performance baseline.
  • Use a high-level profiler to get an initial understanding of where time is being spent.
  • If communication appears to be a bottleneck, perform a deep-dive analysis with NVIDIA Nsight Systems.
  • Based on the trace analysis, apply targeted NCCL tuning via environment variables (NCCL_ALGO, NCCL_PROTO) or investigate framework-level settings (e.g., FSDP prefetching, DeepSpeed buffer sizes) to address the specific bottleneck identified. Iterate on this profiling and tuning loop.

 

Future Directions

 

The field of distributed training continues to evolve rapidly. Several emerging trends are poised to shape the future of large-scale AI:

  • Advanced Parallelism Strategies: Techniques like sequence parallelism are becoming more mainstream as models are applied to ever-longer contexts. The co-design of model architectures and parallelism strategies will become increasingly important.
  • Hardware and Interconnect Evolution: The development of next-generation interconnects (e.g., NVLink 5.0) will continue to increase bandwidth and lower latency, potentially shifting the optimal balance between different parallelism strategies. The integration of heterogeneous computing resources and novel memory technologies will also present new optimization challenges and opportunities.
  • Compiler-Level Optimizations: There is a growing trend towards integrating communication and parallelism optimizations directly into deep learning compilers. This could automate many of the complex decisions currently left to the user, such as choosing the optimal parallelism strategy or communication algorithm, leading to more “out-of-the-box” performance and simplifying the user experience.

Ultimately, the ability to effectively manage memory and optimize communication will remain a critical enabler for progress in artificial intelligence. The architectures and techniques detailed in this report form the foundation upon which the next generation of transformative AI models will be built.

Works cited

  1. ZeRO: Memory Optimizations Toward Training Trillion Parameter …, accessed on August 5, 2025, https://arxiv.org/pdf/1910.02054
  2. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism – PARSA, accessed on August 5, 2025, https://parsa.epfl.ch/course-info/cs723/papers/Megatron.pdf
  3. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, accessed on August 5, 2025, https://arxiv.org/abs/1909.08053
  4. [1910.02054] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, accessed on August 5, 2025, https://arxiv.org/abs/1910.02054
  5. Zero Redundancy Optimizer – DeepSpeed, accessed on August 5, 2025, https://www.deepspeed.ai/tutorials/zero/
  6. Parallel and Distributed Deep Learning – Stanford University, accessed on August 5, 2025, https://web.stanford.edu/~rezab/classes/cme323/S16/projects_reports/hedge_usmani.pdf
  7. Understanding Communication Patterns in Distributed ML: A Deep Dive into NCCL, MPI, and Gloo – The ML Architect, accessed on August 5, 2025, https://themlarchitect.com/blog/communication-protocols-for-distributed-ml-nccl-mpi-and-key-patterns/
  8. Parallelism and Distributed Training for Maximizing AI Efficiency | Exxact Blog, accessed on August 5, 2025, https://www.exxactcorp.com/blog/deep-learning/parallelization-and-distributed-training
  9. Parallelism methods – Hugging Face, accessed on August 5, 2025, https://huggingface.co/docs/transformers/main/perf_train_gpu_many
  10. What is distributed training? – Azure Machine Learning | Microsoft Learn, accessed on August 5, 2025, https://learn.microsoft.com/en-us/azure/machine-learning/concept-distributed-training?view=azureml-api-2
  11. Data Parallelism in Machine Learning Training | by Soonmo Seong | Cloud Villains, accessed on August 5, 2025, https://medium.com/cloudvillains/data-parallelism-in-machine-learning-training-686ed9ab05fb
  12. Data Parallelism: From Basics to Advanced Distributed Training …, accessed on August 5, 2025, https://www.digitalocean.com/community/conceptual-articles/data-parallelism-distributed-training
  13. Data-Parallel Distributed Training of Deep Learning Models – siboehm, accessed on August 5, 2025, https://siboehm.com/articles/22/data-parallel-training
  14. Paradigms of Parallelism | Colossal-AI, accessed on August 5, 2025, https://colossalai.org/docs/concepts/paradigms_of_parallelism
  15. A Brief Overview of Parallelism Strategies in Deep Learning | Alex McKinney, accessed on August 5, 2025, https://afmck.in/posts/2023-02-26-parallelism/
  16. Tensor Parallelism | Ayar Labs, accessed on August 5, 2025, https://ayarlabs.com/glossary/tensor-parallelism/
  17. Part 4.1: Tensor Parallelism – the UvA Deep Learning Tutorials!, accessed on August 5, 2025, https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/scaling/JAX/tensor_parallel_simple.html
  18. Communication Optimization for Distributed Training: Architecture, Advances, and Opportunities – arXiv, accessed on August 5, 2025, https://arxiv.org/html/2403.07585v1
  19. Megatron-LM: communication primitives involved in a single transformer layer [20]., accessed on August 5, 2025, https://www.researchgate.net/figure/Megatron-LM-communication-primitives-involved-in-a-single-transformer-layer-20_fig2_383754225
  20. Tensor Parallelism Overview — AWS Neuron Documentation, accessed on August 5, 2025, https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-tp-appnote.html
  21. Model Parallelism — transformers 4.11.3 documentation – Hugging Face, accessed on August 5, 2025, https://huggingface.co/transformers/v4.11.3/parallelism.html
  22. A Unified Sequence Parallelism Approach for Long Context Generative AI – arXiv, accessed on August 5, 2025, https://arxiv.org/pdf/2405.07719
  23. [2205.05198] Reducing Activation Recomputation in Large Transformer Models – arXiv, accessed on August 5, 2025, https://arxiv.org/abs/2205.05198
  24. Tensor Parallelism in Three Levels of Difficulty – Determined AI, accessed on August 5, 2025, https://www.determined.ai/blog/tp
  25. Scaling Large Language Models with DeepSpeed ZeRO, ZeRO++, and ZeRO-Offload — A Complete Guide | by Pratish Dewangan | Jul, 2025 | Medium, accessed on August 5, 2025, https://medium.com/@dpratishraj7991/scaling-large-language-models-with-deepspeed-zero-zero-and-zero-offload-a-complete-guide-70d393e311f4
  26. Gradient Checkpoints — PyTorch Training Performance Guide, accessed on August 5, 2025, https://residentmario.github.io/pytorch-training-performance-guide/gradient-checkpoints.html
  27. cybertronai/gradient-checkpointing: Make huge neural nets … – GitHub, accessed on August 5, 2025, https://github.com/cybertronai/gradient-checkpointing
  28. Transcending Runtime-Memory Tradeoffs in Checkpointing by being Fusion Aware – MLSys Proceedings, accessed on August 5, 2025, https://proceedings.mlsys.org/paper_files/paper/2023/file/8a27bb69950c0b46cdb36d10e5514cc8-Paper-mlsys2023.pdf
  29. Training Deep Nets with Sublinear Memory Cost – deepsense.ai, accessed on August 5, 2025, https://deepsense.ai/wp-content/uploads/2023/04/1604.06174.pdf
  30. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, accessed on August 5, 2025, https://www.cs.toronto.edu/~cmaddis/courses/csc2541_w25/presentations/richards_zero.pdf
  31. DeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective. – GitHub, accessed on August 5, 2025, https://github.com/deepspeedai/DeepSpeed
  32. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models – Microsoft, accessed on August 5, 2025, https://www.microsoft.com/en-us/research/publication/zero-memory-optimizations-toward-training-trillion-parameter-models/
  33. DeepSpeed – Hugging Face, accessed on August 5, 2025, https://huggingface.co/docs/accelerate/usage_guides/deepspeed
  34. Deepspeed communication pattern clarification · Issue #264 – GitHub, accessed on August 5, 2025, https://github.com/microsoft/DeepSpeed/issues/264
  35. ZeRO — DeepSpeed 0.17.4 documentation, accessed on August 5, 2025, https://deepspeed.readthedocs.io/en/stable/zero3.html
  36. ZeRO-Offload – DeepSpeed, accessed on August 5, 2025, https://www.deepspeed.ai/tutorials/zero-offload/
  37. ZeRO-Offload: Democratizing Billion-Scale Model Training – USENIX, accessed on August 5, 2025, https://www.usenix.org/system/files/atc21-ren-jie.pdf
  38. [2101.06840] ZeRO-Offload: Democratizing Billion-Scale Model Training – arXiv, accessed on August 5, 2025, https://arxiv.org/abs/2101.06840
  39. NVIDIA Collective Communications Library (NCCL) | NVIDIA …, accessed on August 5, 2025, https://developer.nvidia.com/nccl
  40. NVIDIA/nccl: Optimized primitives for collective multi-GPU communication – GitHub, accessed on August 5, 2025, https://github.com/NVIDIA/nccl
  41. NCCL: Getting Started – NVIDIA Developer, accessed on August 5, 2025, https://developer.nvidia.com/nccl/getting_started
  42. Overview of NCCL — NCCL 2.27.5 documentation – NVIDIA Docs Hub, accessed on August 5, 2025, https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/overview.html
  43. NVIDIA NCCL — NVIDIA GB200 NVL Multi-Node Tuning Guide, accessed on August 5, 2025, https://docs.nvidia.com/multi-node-nvlink-systems/multi-node-tuning-guide/nccl.html
  44. Collective Operations — NCCL 2.27.5 documentation – NVIDIA Docs Hub, accessed on August 5, 2025, https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html
  45. Enable faster training with Amazon SageMaker data parallel library | Artificial Intelligence, accessed on August 5, 2025, https://aws.amazon.com/blogs/machine-learning/enable-faster-training-with-amazon-sagemaker-data-parallel-library/
  46. Understanding NCCL Tuning to Accelerate GPU-to-GPU …, accessed on August 5, 2025, https://developer.nvidia.com/blog/understanding-nccl-tuning-to-accelerate-gpu-to-gpu-communication/
  47. NCCL ALGOs aren’t quite documented · Issue #790 · NVIDIA/nccl – GitHub, accessed on August 5, 2025, https://github.com/NVIDIA/nccl/issues/790
  48. NCCL Basics for Network Engineers – TechShinobi, accessed on August 5, 2025, https://techshinobi.hashnode.dev/network-engineers-introductory-guide-to-nccl
  49. Communication Algorithm-Architecture Co-Design for Distributed Deep Learning, accessed on August 5, 2025, https://par.nsf.gov/servlets/purl/10374122
  50. Optcast: Open-Source In-Network Aggregation for Distributed Deep Learning | by Nariaki Tateiwa | nttlabs | Medium, accessed on August 5, 2025, https://medium.com/nttlabs/optcast-open-source-in-network-aggregation-for-distributed-deep-learning-c25d76878993
  51. Massively Scale Your Deep Learning Training with NCCL 2.4 | NVIDIA Technical Blog, accessed on August 5, 2025, https://developer.nvidia.com/blog/massively-scale-deep-learning-training-nccl-2-4/
  52. NCCL – ALCF User Guides, accessed on August 5, 2025, https://docs.alcf.anl.gov/polaris/applications-and-libraries/libraries/nccl/
  53. Understanding Multi GPU Communication and Nvidia NCCL for finetuning models, accessed on August 5, 2025, https://tensorfuse.io/docs/blogs/multi_gpu_communication_while_training
  54. Scaling Deep Learning Training with NCCL | NVIDIA Technical Blog, accessed on August 5, 2025, https://developer.nvidia.com/blog/scaling-deep-learning-training-nccl/
  55. Understanding Nvidia’s NvLink and NvSwitch Evolution: Topology and Rates – fibermall.com, accessed on August 5, 2025, https://www.fibermall.com/blog/nvidia-nvlink-and-nvswitch-evolution.htm
  56. Network architecture | Technical White Paper–Generative AI in the Enterprise – Model Training, accessed on August 5, 2025, https://infohub.delltechnologies.com/pt-br/l/technical-white-paper-generative-ai-in-the-enterprise-model-training/network-architecture-155/
  57. Networks for AI, accessed on August 5, 2025, https://s201.q4cdn.com/141608511/files/doc_presentations/2023/06/27/Networks_for_AI_IR_Webcast_June_2023-FINAL.pdf
  58. Doubling all2all Performance with NVIDIA Collective Communication Library 2.12, accessed on August 5, 2025, https://developer.nvidia.com/blog/doubling-all2all-performance-with-nvidia-collective-communication-library-2-12/
  59. What are the optimal settings for NCCL on NVIDIA GPUs? – Massed Compute, accessed on August 5, 2025, https://massedcompute.com/faq-answers/?question=What%20are%20the%20optimal%20settings%20for%20NCCL%20on%20NVIDIA%20GPUs?
  60. DeepSpeed vs FSDP: A Comprehensive Comparison – BytePlus, accessed on August 5, 2025, https://www.byteplus.com/en/topic/499106
  61. PyTorch Fully Sharded Data Parallel (FSDP) on AMD GPUs with …, accessed on August 5, 2025, https://rocm.blogs.amd.com/artificial-intelligence/fsdp-training-pytorch/README.html
  62. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel – VLDB Endowment, accessed on August 5, 2025, https://www.vldb.org/pvldb/vol16/p3848-huang.pdf
  63. Fully Sharded Data Parallel – Hugging Face, accessed on August 5, 2025, https://huggingface.co/docs/accelerate/usage_guides/fsdp
  64. FSDP Notes — PyTorch 2.7 documentation, accessed on August 5, 2025, https://pytorch.org/docs/stable/notes/fsdp.html
  65. FSDP vs DeepSpeed – Hugging Face, accessed on August 5, 2025, https://huggingface.co/docs/accelerate/concept_guides/fsdp_and_deepspeed
  66. DDP Communication Hooks – PyTorch documentation, accessed on August 5, 2025, https://docs.pytorch.org/docs/stable/ddp_comm_hooks.html
  67. FullyShardedDataParallel — PyTorch 2.7 documentation, accessed on August 5, 2025, https://docs.pytorch.org/docs/stable/fsdp.html
  68. DDP Communication Hooks – Hugging Face, accessed on August 5, 2025, https://huggingface.co/docs/accelerate/usage_guides/ddp_comm_hook
  69. Maximizing Communication Efficiency for Large-scale Training via 0/1 Adam – DeepSpeed, accessed on August 5, 2025, https://www.deepspeed.ai/tutorials/zero-one-adam/
  70. Benchmarking Advanced Multi‐GPU Training Strategies | by Savyasachi Thati – Medium, accessed on August 5, 2025, https://medium.com/@savyasachi.thati/benchmarking-advanced-multi-gpu-training-strategies-20c9675003db
  71. Train models with billions of parameters — PyTorch Lightning 2.5.2 documentation, accessed on August 5, 2025, https://lightning.ai/docs/pytorch/stable/advanced/model_parallel.html
  72. Accelerate vs. DeepSpeed vs. FSDP – Ben Gubler, accessed on August 5, 2025, https://www.bengubler.com/posts/2023-08-29-accelerate-deepspeed-fsdp
  73. [D] What do you all use for large scale training? Normal pytorch or do you use libraries like HF Accelerate. – Reddit, accessed on August 5, 2025, https://www.reddit.com/r/MachineLearning/comments/1dxtaez/d_what_do_you_all_use_for_large_scale_training/
  74. 17.3. GPU Profiling — Kempner Institute Computing Handbook, accessed on August 5, 2025, https://handbook.eng.kempnerinstitute.harvard.edu/s5_ai_scaling_and_engineering/scalability/gpu_profiling.html
  75. Profiling & Monitoring Deep Learning Training Tasks – GitHub Pages, accessed on August 5, 2025, https://itu-dasyalab.github.io/RAD/publication/papers/euromlsys2023.pdf
  76. Using Nsight Systems to profile GPU workload – NVIDIA CUDA – PyTorch Dev Discussion, accessed on August 5, 2025, https://dev-discuss.pytorch.org/t/using-nsight-systems-to-profile-gpu-workload/59
  77. What are the key features of NVIDIA Nsight Systems for GPU-to-GPU communication?, accessed on August 5, 2025, https://massedcompute.com/faq-answers/?question=What+are+the+key+features+of+NVIDIA+Nsight+Systems+for+GPU-to-GPU+communication%3F
  78. Profile Pytorch code using nsys and nsight step by step | by …, accessed on August 5, 2025, https://medium.com/@yuanzhedong/profile-pytorch-code-using-nsys-and-nsight-step-by-step-9c3f01995fd3
  79. Nsight Systems – DL Profiling Argonne National Labs 2022-06-30, accessed on August 5, 2025, https://www.alcf.anl.gov/sites/default/files/2024-07/Nsight-Systems-DL-Profiling-2022-06-30.pdf
  80. Speed Up PyTorch Training by 3x with NVIDIA Nsight and PyTorch 2.0 Tricks | Practical ML, accessed on August 5, 2025, https://arikpoz.github.io/posts/2025-05-25-speed-up-pytorch-training-by-3x-with-nvidia-nsight-and-pytorch-2-tricks/
  81. Analyzing NCCL Usage with NVIDIA Nsight Systems – YouTube, accessed on August 5, 2025, https://www.youtube.com/watch?v=bda5Er27jqo
  82. Analyzing NCCL Usage with NVIDIA Nsight Systems | Supercomputing 2020, accessed on August 5, 2025, https://www.nvidia.com/en-us/on-demand/session/supercomputing2020-sc2024/