Data Parallelism & Model Parallelism

Let’s break down the two concept and how to implement it.

The Difference: Data Parallelism vs. Model Parallelism

Think of it like a factory:

• Data Parallelism (nn.DataParallel / nn.DistributedDataParallel): You have several identical, small factories (your GPUs). To speed up production, you send different sets of raw materials (your data batches) to each factory. They all perform the same tasks on their own set of materials. This works only if the entire factory blueprint (your model) can fit in each building.
• Model Parallelism: You have one enormous, complex assembly line (your model) that is too long or too big to fit in a single building (a single GPU). So, you build the first half of the assembly line in Building 1 (cuda:0) and the second half in Building 2 (cuda:1). Raw materials (your data) go into Building 1, get partially assembled, and then the semi-finished product is moved to Building 2 for the final steps.

PyTorch Underutilizing GPUs? Unraveling the Mystery of nn.Parallel and Single-GPU Usage

You’ve equipped your machine with four powerful GPUs, ready to accelerate your deep learning model training with PyTorch’s nn.Parallel. However, to your frustration, only a single GPU is shouldering the entire workload, leaving the other three idle. This is a common hurdle for many developers, and the solution often lies in a few key areas, from environment settings to the nuances of how Pytoch handles parallel processing.

This guide will walk you through the most frequent culprits behind this issue and provide actionable solutions to ensure all your available GPUs are being leveraged effectively.

The Likely Suspects and How to Address Them

The primary reasons for nn.Parallel (specifically nn.DataParallel) failing to utilize all available GPUs can be categorized as follows:

1. The CUDA_VISIBLE_DEVICES Environment Variable:

This is one of the most common and often overlooked reasons. The CUDA_VISIBLE_DEVICES environment variable dictates which GPUs are visible to your PyTorch script.¹ If this variable is not set or is set to a single GPU index, PyTorch will only “see” and be able to use that specific GPU.²

How to Fix:

• In your terminal: Before running your Python script, you can set the variable to include all your GPUs.³ The indexing for GPUs starts from 0. For four GPUs, you would do: Bashexport CUDA_VISIBLE_DEVICES=0,1,2,3 python your_script.py
• Within your Python script: You can set the environment variable at the beginning of your script, before importing torch. Pythonimport os os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" import torch

2. Verifying GPU Availability in PyTorch:

It’s crucial to programmatically check how many GPUs PyTorch can actually detect.

How to Check:

Use the following code snippet to see what PyTorch reports:

Python

import torch

if torch.cuda.is_available():
    print(f"CUDA is available. Number of GPUs: {torch.cuda.device_count()}")
    for i in range(torch.cuda.device_count()):
        print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
else:
    print("CUDA is not available.")

If the output of torch.cuda.device_count() is 1, it confirms that PyTorch only sees a single GPU, pointing back to a potential issue with CUDA_VISIBLE_DEVICES or your CUDA installation.

3. The Limitations of nn.DataParallel and the Rise of nn.DistributedDataParallel:

While nn.DataParallel is straightforward to implement, it comes with a significant drawback: it creates an imbalance in GPU workload. The primary GPU (usually GPU 0) not only processes its share of the data but also manages the other GPUs, gathers the results, and computes the final loss. This often leads to the master GPU being heavily loaded while the others are underutilized.

For more efficient and balanced multi-GPU training, the PyTorch team strongly recommends using nn.DistributedDataParallel (DDP).⁴ DDP uses a different parallelization strategy that distributes the workload more evenly across all GPUs.

Recommendation:

Migrating your training script from nn.DataParallel to nn.DistributedDataParallel is the most robust solution for multi-GPU training. While it requires a bit more setup, the performance gains and balanced utilization are well worth the effort.

4. Incorrect Usage of nn.DataParallel:

Even if you stick with nn.DataParallel for its simplicity, incorrect implementation can lead to single-GPU usage.

Common Mistakes:

• Batch Size: nn.DataParallel works by splitting your batch of data along the first dimension and sending a chunk to each GPU.⁵ If your batch size is 1, there’s nothing to split, and the entire batch will be processed on the master GPU. Ensure your batch size is larger than the number of GPUs you intend to use.
• Data Placement: You need to ensure that your input tensors are moved to the correct device in each training iteration. A common pattern is: Pythondevice = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = MyModel() if torch.cuda.device_count() > 1: print(f"Using {torch.cuda.device_count()} GPUs!") model = nn.DataParallel(model) model.to(device) # In your training loop for inputs, labels in data_loader: inputs, labels = inputs.to(device), labels.to(device) # This is crucial! # ... rest of your training logic If you only move the model to the device and not the input data for each batch, the computation will likely default to the CPU or a single GPU.
• Model Wrapping: The nn.DataParallel wrapper should be applied to your model before you start your training loop.

A Quick Checklist for Troubleshooting

1. Check CUDA_VISIBLE_DEVICES: Is it set correctly to include all your GPU indices?
2. Verify with PyTorch: Run torch.cuda.device_count() to see how many GPUs PyTorch detects.⁶
3. Consider nn.DistributedDataParallel: For serious multi-GPU training, this is the recommended approach for better performance and balanced load.
4. Review your nn.DataParallel implementation:
   • Is your batch size greater than the number of GPUs?
   • Are you moving your input data to the correct device in every iteration of your training loop?
   • Is your model correctly wrapped with nn.DataParallel?

By systematically going through these points, you can identify the bottleneck and unlock the full potential of your multi-GPU setup, significantly accelerating your deep learning research and development.

How to Implement Model Parallelism in PyTorch (The Manual Way)

The most straightforward way to achieve this is to manually assign different parts (layers) of your model to different GPUs. You then have to explicitly move the data (intermediate tensors) between these GPUs in the forward pass.

Here’s a conceptual example. Let’s say you have a large model that you can split into two sequential parts.

Python

import torch
import torch.nn as nn

class LargeModel(nn.Module):
    def __init__(self):
        super(LargeModel, self).__init__()

        # --- Part 1: To be loaded on GPU 0 ---
        # Define the first set of layers
        self.part1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
            # ... many more layers that fit on one GPU
            nn.Linear(1024, 512)
        ).to('cuda:0') # <-- Explicitly move this part to the first GPU

        # --- Part 2: To be loaded on GPU 1 ---
        # Define the second set of layers
        self.part2 = nn.Sequential(
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            # ... the rest of your model's layers
            nn.Linear(256, 10) # Output layer
        ).to('cuda:1') # <-- Explicitly move this part to the second GPU

    def forward(self, x):
        # 1. Start the computation on the device of the input data
        #    Let's assume the input `x` is initially on cuda:0
        x = x.to('cuda:0')

        # 2. Process the data through the first part of the model on GPU 0
        x = self.part1(x)

        # 3. !! CRITICAL STEP !!
        #    Move the intermediate output from GPU 0 to GPU 1
        x = x.to('cuda:1')

        # 4. Process the data through the second part of the model on GPU 1
        x = self.part2(x)

        return x

# --- How to use it ---

# First, check if you have at least 2 GPUs
if torch.cuda.device_count() < 2:
    print("This example requires at least 2 GPUs.")
else:
    # Instantiate the model. The __init__ method handles device placement.
    model = LargeModel()

    # Create a dummy input tensor on the first GPU
    # The batch size can be whatever your GPUs can handle
    dummy_input = torch.randn(64, 3, 224, 224).to('cuda:0')

    # Run the forward pass
    output = model(dummy_input)

    # The final output will be on cuda:1
    print(f"Output tensor is on device: {output.device}")

    # Now you can calculate the loss. The target tensor also needs to be on the correct device.
    target = torch.randint(0, 10, (64,)).to('cuda:1')
    loss = nn.CrossEntropyLoss()(output, target)
    print(f"Loss tensor is on device: {loss.device}")

    # Backpropagation works automatically across GPUs!
    loss.backward()

    print("Gradients calculated successfully.")
    # You can now call your optimizer.step()

Key Considerations and Modern Alternatives

1. Bottlenecks: The manual approach is simple to understand but can be inefficient. While part2 on cuda:1 is computing, part1 on cuda:0 is sitting idle. This is known as a “pipeline bubble.” Moving data between GPUs (x.to('cuda:1')) also introduces overhead.
2. Pipeline Parallelism: To mitigate the “bubble” effect, you can use a more advanced technique called pipeline parallelism, where you micro-batch the input and start feeding new data to GPU 0 while GPU 1 is still working on a previous chunk. The torch.distributed.pipeline module can help automate this, but it’s more complex to set up.
3. Fully Sharded Data Parallelism (FSDP): This is the modern, state-of-the-art solution in PyTorch for training massive models. FSDP is a type of data parallelism, but with a crucial difference: instead of replicating the entire model on each GPU, it shards (splits) the model’s parameters, gradients, and optimizer states across all the GPUs. It’s incredibly memory-efficient and is now the recommended approach for large-scale model training. It elegantly solves the problem of both model size and training speed.

Recommendation

• For understanding: Start by implementing the manual model parallel approach shown above. It will give you a clear understanding of what’s happening under the hood.
• For serious training: If you are working with truly large models like Transformers, you should invest the time to learn and implement Fully Sharded Data Parallelism (FSDP). It is far more efficient and scalable.