Data Parallelism (DP)

Data Parallelism (DP) stands as the most straightforward and commonly employed strategy for distributing the computational load of training deep learning models, including large language models. As highlighted in the chapter introduction, training large models often encounters bottlenecks related to computation time and memory constraints on single devices. Data parallelism directly addresses the computation time aspect by processing different parts of the data simultaneously across multiple devices.

The core principle is simple: replicate the entire model onto each available processing unit (like a GPU), and then divide each global data batch into smaller pieces, called micro-batches. Each device processes its assigned micro-batch independently.

How Data Parallelism Works

Let's break down the typical steps involved in one training iteration using data parallelism with $K$ devices:

1. Model Replication: Before training begins, an identical copy of the entire model (parameters, buffers) is loaded onto each of the $K$ devices.
2. Data Sharding: The global batch of training data is split into $K$ micro-batches. Each device receives one micro-batch. For instance, if the global batch size is $B$ and we have $K$ devices, each device typically processes a micro-batch of size $B/K$.
3. Forward Pass: Each device performs a forward pass using its local model copy and its assigned micro-batch. This computes the loss for that portion of the data.
4. Backward Pass: Each device computes the gradients of the loss with respect to its local model parameters via backpropagation. At this point, each device $i$ holds a set of gradients $g_i$, calculated based only on its micro-batch.
5. Gradient Synchronization: This is the defining step of data parallelism. The gradients computed independently on each device must be combined to represent the gradient for the entire global batch. This is typically achieved using a collective communication operation called AllReduce. The AllReduce operation sums (or averages) the gradients $g_i$ from all $K$ devices and distributes the resulting synchronized gradient $g_{sync}$ back to every device. Often, averaging is used: $$g_{sync} = \frac{1}{K} \sum_{i=1}^{K} g_i$$
6. Optimizer Step: Each device uses the synchronized gradient $g_{sync}$ to update its local copy of the model parameters using the chosen optimizer (e.g., AdamW). Because all devices start with the same parameters and apply the same update based on the synchronized gradient, the model replicas remain identical after the optimizer step.

This cycle repeats for each batch in the training dataset.

Data Parallelism workflow: The global batch is split, processed in parallel on devices holding model replicas, gradients are synchronized via AllReduce, and synchronized updates are applied to each replica.

Advantages of Data Parallelism

• Simplicity: Frameworks like PyTorch (DistributedDataParallel) provide high-level abstractions that make implementing data parallelism relatively straightforward, often requiring only minor modifications to standard single-device training code.
• Increased Throughput: By processing data in parallel, DP significantly reduces the time per training step, allowing for faster iteration through large datasets or the use of larger effective batch sizes, which can sometimes improve model convergence and generalization.
• Wide Applicability: Data parallelism doesn't require fundamental changes to the model architecture itself. It works "out of the box" for most standard network designs.

Limitations of Data Parallelism

Despite its advantages, data parallelism has a significant limitation, especially for the large language models this course focuses on:

• Memory Constraint: The primary drawback is that each device must hold a complete copy of the model's parameters, gradients, and optimizer states, as well as the activations computed during the forward pass for its micro-batch. For models with billions or trillions of parameters, the memory required often exceeds the capacity of even the largest available accelerators (GPUs/TPUs). This fundamentally limits the maximum model size that can be trained using data parallelism alone.
• Communication Overhead: The AllReduce operation requires communicating gradients across all devices involved. The amount of data transferred is proportional to the size of the model parameters. As the number of devices ($K$) increases, or if the interconnect bandwidth between devices is limited, this synchronization step can become a significant bottleneck, diminishing the speedup gained from parallel computation. The time taken for AllReduce can sometimes dominate the computation time, especially for smaller models or faster computations per device.

Implementation Sketch with PyTorch DistributedDataParallel (DDP)

PyTorch's torch.nn.parallel.DistributedDataParallel module is the standard way to implement data parallelism in a multi-process setting, which is generally preferred over the older torch.nn.DataParallel for performance and flexibility. Here's an outline:

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
import os

# Assume a simple model definition exists:
# class MyLargeModel(nn.Module): ...

def setup(rank, world_size):
    """Initialize the process group."""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    # Initialize the process group
    # Use 'nccl' backend for NVIDIA GPUs
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    """Destroy the process group."""
    dist.destroy_process_group()

def train_worker(rank, world_size, model_args, data_args, train_args):
    """Main training function for a single worker process."""
    print(f"Running DDP training on rank {rank}.")
    setup(rank, world_size)

    # Instantiate the model and move it to the assigned GPU
    model = MyLargeModel(**model_args).to(rank)
    # Wrap the model with DDP
    # This handles gradient synchronization automatically
    ddp_model = DDP(model, device_ids=[rank])

    # Prepare dataset and DistributedSampler
    # DistributedSampler ensures each process gets a different
    # slice of data
    dataset = YourDataset(**data_args)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    # Use pin_memory=True and num_workers > 0 for optimized data loading
    dataloader = DataLoader(
        dataset,
        batch_size=train_args['micro_batch_size'],
        sampler=sampler,
        pin_memory=True,
        num_workers=4
    )

    optimizer = torch.optim.AdamW(
        ddp_model.parameters(),
        lr=train_args['learning_rate']
    )

    # --- Simplified Training Loop ---
    ddp_model.train()
    for epoch in range(train_args['num_epochs']):
        sampler.set_epoch(epoch) # Important for shuffling with DDP
        for batch in dataloader:
            inputs = batch['input_ids'].to(rank)
            labels = batch['labels'].to(rank)

            optimizer.zero_grad()
            outputs = ddp_model(inputs, labels=labels) # Forward pass
            loss = outputs.loss
            # Backward pass - DDP automatically averages gradients
            loss.backward()
            # Optimizer step - applies update based on averaged grads
            optimizer.step()

            if rank == 0: # Log only on the main process
                print(f"Epoch: {epoch}, Loss: {loss.item()}")
        # --- End Simplified Loop ---

    cleanup()

if __name__ == '__main__':
    # Example configuration (replace with actual args)
    world_size = torch.cuda.device_count() # e.g., 4 GPUs
    model_args = {'vocab_size': 50257, 'hidden_size': 768, 'num_layers': 12}
    data_args = {'data_path': '/path/to/data'}
    train_args = {
        'micro_batch_size': 8,
        'learning_rate': 1e-4,
        'num_epochs': 3
    }
    # Note: Global batch size = micro_batch_size * world_size

    # Spawn worker processes
    mp.spawn(
        train_worker,
        args=(world_size, model_args, data_args, train_args),
        nprocs=world_size,
        join=True
    )

In this sketch:

• setup initializes the distributed environment. Each process gets a unique rank from 0 to world_size - 1.
• DistributedSampler is used with the DataLoader to ensure each process gets non-overlapping data partitions.
• The core change is wrapping the model: ddp_model = DDP(model, device_ids=[rank]). DDP intercepts the backward pass, performs the AllReduce operation on the gradients automatically, and ensures all processes have the averaged gradients before the optimizer.step() is called.

Data parallelism is a fundamental technique for scaling training throughput. However, its memory limitations necessitate exploring other strategies like tensor and pipeline parallelism, especially when dealing with the enormous scale of modern large language models. Often, the most effective approaches combine data parallelism with these other techniques, which we will examine next.