Data Parallelism
Distributed Data Parallelism
Model Conversion
PyTorch
Deep Learning

How to convert model.module.fc when using DP, if using DDP

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

When working in the field of deep learning, particularly with large-scale models and data, distributing the workload efficiently across multiple GPUs can significantly enhance computational efficiency and reduce training time. The two common strategies for this purpose are Data Parallelism (DP) and Distributed Data Parallelism (DDP). Both approaches handle the division of computational tasks differently and have implications on how model components, such as fully connected layers (e.g., model.module.fc), are accessed and modified.

Understanding Data Parallelism (DP) and Distributed Data Parallelism (DDP)

Data Parallelism (DP)

DP involves splitting the input data across different GPUs. Each GPU receives a copy of the model and processes a different part of the input data. After processing, the results (typically gradients during backpropagation) are gathered and averaged across all GPUs. In PyTorch, this can be implemented very easily using torch.nn.DataParallel. When using DataParallel, your model is automatically replicated across the GPUs, and each attribute of your model can be accessed typically through an additional .module (e.g., model.module.fc for the fully connected layer).

Distributed Data Parallelism (DDP)

DDP, on the other hand, is designed for distributed training across multiple machines/processors. Unlike DP, each process in DDP manages one GPU and only a part of the model if model parallel is used, greatly reducing the bandwidth pressure on the network. DDP in PyTorch is implemented using torch.nn.parallel.DistributedDataParallel. It synchronizes gradients by applying all-reduce across all processes in its group and does not replicate the model across all GPUs. Instead, each GPU processes a unique batch and the gradients are synchronized.

Converting model.module.fc When Using DP to DDP

Switching from DP to DDP can lead to confusion regarding how to access model attributes such as the fully connected layer (i.e., model.module.fc). In DP, since the model is replicated, you access the layer with the additional .module. However, in DDP, you interact directly with the model’s attributes since the model isn't wrapped with an additional module layer.

Here are the steps and considerations when converting:

  1. Initialize the model as usual. When defining the model, you might typically instantiate your fully connected layer with something like self.fc = nn.Linear(...) in your model's class definition.
  2. Apply DistributedDataParallel: When you wrap your model with DistributedDataParallel, you do so typically after the model instantiation:
python
   model = MyModel()
   model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])
  1. Accessing the fully connected layer:
    • With DP: output = model.module.fc(input)
    • With DDP: output = model.fc(input)

Key Points of Transition from DP to DDP

AspectData Parallel (DP)Distributed Data Parallel (DDP)
Model ReplicationAcross all GPUsOnly one GPU per process
Attribute AccessVia model.moduleDirect attribute access
Gradient SynchronizationAt the end of backward passesContinuously during backward passes
ScalabilityLimited to a single machineScales across multiple nodes

Additional Tips for Switching from DP to DDP

  • Ensure proper GPU setup: Each process should ideally be tied to a single GPU. Ensure that you set the CUDA device correctly with torch.cuda.set_device(rank).
  • Balance the load: Ensure that the data is evenly divided across all processes to prevent some GPUs from becoming bottlenecks.
  • NCCL Backend: For best performance in DDP, consider using the NCCL backend which is optimized for high-performance training on NVIDIA GPUs.

Switching from DP to DDP can lead to a significant improvement in performance, especially for training jobs that can scale across multiple nodes. Correctly transforming the way you access model components such as the fully connected layer is crucial for a seamless transition.


Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.