Parallelism in generative AI models

Compare data, model, and hybrid parallelism by following what each GPU stores and computes. Examine how synchronization, network traffic, memory, and failure recovery affect the useful speedup of a distributed training job.

A training job may exceed one device’s memory, take too long on one device, or both. These are different constraints: copying the model to more GPUs can increase data throughput, but it does not make each copy smaller.

A historical estimate: The 2021 Megatron-LM study estimated approximately 288 years for GPT-3-scale training on one V100 under its assumptions. Such an extrapolation illustrates compute demand; it is not a feasible single-GPU deployment plan because memory capacity is also a constraint.

Distributed machine learning (DML) spreads training across devices or machines. It can divide data, model computation, or both. Useful speedup depends on how much work can run concurrently after accounting for communication and waiting.

Distribution introduces three responsibilities:

• Communication: Transfer gradients, parameters, or activations over links with finite bandwidth and nonzero latency.

• Synchronization: Define which updates belong to a training step and when each worker may advance.

• Recovery: Detect failures and restore a consistent training state, including optimizer and progress information.

Data parallelism divides examples; model parallelism divides a model’s storage or computation. The following sections trace each choice and then combine them.

Data parallelism

In ordinary replicated data parallelism, each worker holds the same model and processes a different local batch. Workers compute gradients, combine them according to the training rule, and update their replicas consistently. PyTorch DDP illustrates this gradient-synchronization approach. The application is responsible for partitioning the input data.

The Megatron-LM paper estimated 34 days for a 175-billion-parameter model trained on 300 billion tokens using 1,024 A100 GPUs and its measured throughput. That configuration combined tensor, pipeline, and data parallelism. It is not evidence that data parallelism alone fits the model or achieves the same time on an arbitrary cluster.

The local batch assignment must agree with both device capacity and the global update:

• In a homogeneous cluster, equal local batches are a useful starting point. Sequence length and preprocessing can still make their execution times differ.

• In a heterogeneous cluster, faster workers may process larger local batches. The gradient aggregation must then preserve the intended contribution of each example; an unweighted average of unequal local-batch means changes those contributions.

The synchronization design determines how local updates become a consistent model step.

Parameter server

A parameter-server design assigns shared parameter management to one or more servers. Workers send updates such as gradients; the server applies the chosen update procedure and supplies parameters to workers. A single unreplicated server can become a capacity bottleneck or a point of failure. Sharding and replication change those limits but add their own consistency and recovery work.

Peer-to-peer synchronization

Collective synchronization combines worker updates without requiring one central parameter server. The important contract is the collective’s result, separate from the communication algorithm that produces it:

• AllReduce: Combine corresponding values from all participating workers and make the reduced result available to each. Sum, minimum, and maximum are possible reductions; averaging gradients also requires the appropriate normalization. NCCL’s collective definitions distinguish the operation from its implementation. AllReduce is not synonymous with every worker directly sending a full buffer to every peer.

• Ring AllReduce: Arrange communication in a logical ring and exchange chunks through neighboring workers. A typical implementation uses a reduce-scatter phase followed by all-gather, so communication can overlap across chunks. The following small diagrams illustrate aggregation and distribution, not a complete schedule for an optimized ring implementation. Ring performance depends on message size, topology, and slow participants; it is not universally faster than other AllReduce algorithms.

• Hierarchical AllReduce: Organize communication around groups, often matching faster links within a machine and slower links between machines. Reduce within a group, combine group results, and distribute the result back. The implementation may use selected workers or specialized network support for the cross-group stage; a separate coordinator service is not mandatory.

Follow the same gradient contribution through three boundaries in the next figures: local workers, group aggregation, and the final result returned to every worker. Each contribution must be included with the correct weight.

Replicated reduced values are not a complete fault-tolerance mechanism. A failed participant can still interrupt a collective, and remaining copies do not automatically repair membership or restore optimizer state. Use an explicit restart or reconfiguration procedure and checkpoints; do not infer recovery from the diagram’s redundant arrows.

Compare the update and recovery responsibilities in the two designs:

Model parallelism

Model parallelism partitions a model’s storage or computation across devices. It can make a model fit when one device cannot hold the required training or inference state. It also introduces transfers between partitions, so adding devices does not automatically reduce latency. Larger GPU memory changes the capacity calculation; it does not eliminate model parallelism for models whose full state still exceeds that capacity.

There are different ways to partition a model, each with its trade-offs:

1. Layer-wise partitioning: Assign groups of layers to different devices. Activations flow forward across boundaries and gradients flow backward during training. A pipeline schedule can overlap different microbatches, while stage imbalance and pipeline fill or drain time can leave devices idle.

2. Operator-wise or tensor partitioning: Split an operation, such as a large matrix multiplication, across devices. This distributes computation within a layer but requires communication to combine or exchange partial results.

The next diagrams show layer and operation boundaries on a small network.

Hybrid parallelism

Hybrid parallelism combines these approaches. A group of GPUs can hold one partitioned model, while several groups process different data batches as data-parallel replicas. The groups need not correspond exactly to physical machines: placement should reflect memory requirements and the cost of the links each communication pattern uses.

Note: Later case studies often begin with replicated data parallelism as a simplifying assumption. Verify that the full training state fits before accepting it. Parameters, gradients, optimizer state, activations, and temporary buffers all use memory; model weights fitting alone is insufficient.

Challenges in parallelizing GenAI models

The design is useful only if it trains the intended model correctly and makes progress under the cluster’s constraints. Review the following issues before treating GPU count as a speedup estimate.

Fault tolerance

In large distributed systems, the risk of node failure or communication errors increases, potentially leading to training interruptions.

A recovery plan needs more than spare compute:

• Checkpointing: Save enough training state to resume consistently, and test restoration. Choose the interval by comparing write overhead with work that could be lost.

• Recovery capacity: Keep or provision replacement workers according to the acceptable restart delay. A spare worker does not know the training state until it is restored.

• Monitoring: Detect failed workers, stalled collectives, and storage errors. Connect the alert to a bounded recovery procedure.

Hardware heterogeneity

Not all GPUs or servers in a distributed setup may have the same compute power, memory, or architecture, leading to inefficiencies and bottlenecks.

Two choices help manage a heterogeneous cluster:

• Device-specific work: Measure the actual step time on each device and adjust partitioning or local batch assignment while preserving the intended update.

• Compatible software and hardware: Use a homogeneous cluster where practical, or verify the kernels, precision support, memory limits, and communication paths across device types.

Load imbalance

If certain GPUs handle more work than others, this results in idle time for some devices and reduces overall efficiency. This is the phenomenon of load imbalance.

Balance measured execution time, not just layer or example counts:

• Data assignment: Consider sequence lengths, padding, and preprocessing as well as examples per worker. Unequal local batches need correct gradient weighting.

• Partition placement: Profile layer compute and activation transfers before assigning stages. A balanced compute estimate can still hide an expensive boundary transfer.

Measure step time, memory, and communication on the actual workload before distributing work across GPU types. Peak specifications must use matching precision and sparsity assumptions. NVIDIA’s H100 specifications mark the listed Tensor Core rates as sparse, so they cannot be divided directly by a dense rate from another device to infer a training-speed ratio.

Conclusion

Choose parallelism by tracing three things: what each device stores, what it computes, and what crosses each link. Data parallelism replicates a model across batches; model parallelism divides its computation; hybrid designs combine both. Each choice needs a consistent update rule and a recovery plan.

For a hypothetical eight-GPU design, compare eight complete replicas with two groups of four GPUs holding partitioned models. Which arrangement fits the full state, how many batches advance concurrently, and where are gradients or activations exchanged? Answer those questions before predicting training time.