This page looks best with JavaScript enabled

Model Parallel Training Techniques

 ·  ☕ 4 min read

1. Data Parallelism

Training steps:

  1. The master device loads the model and copies the model parameters to every worker device
  2. The master device splits the training data along the batch dimension and passes each batch to every worker device
  3. Each worker device performs training
  4. The master device aggregates the gradients from each worker device and updates the model parameters
  5. The master device broadcasts the model parameters to every worker device, preparing for the next batch of training

Core idea:

Split the training data along the batch dimension and distribute it to multiple worker devices for parallel computation, thereby speeding up training. Each worker keeps a full copy of the model parameters and independently performs the forward and backward passes. Finally, a single master device aggregates the gradients from all workers, updates the model parameters uniformly, and broadcasts them back to the workers.

Applicable scenarios:

  • Training scenarios with large-scale datasets

Characteristics:

  • The model parameter count is relatively small and can be loaded on a single GPU
  • Hardware requirements are relatively low; a multi-GPU server is enough
  • Communication overhead is small, mainly gradient synchronization
  • Suited to small-model training for most CV/NLP tasks

2. Tensor Parallelism

Training steps:

  1. Split certain tensors of the model (such as weight matrices, activation values, and so on) by row or column across different devices
  2. During the forward pass, each device computes the operations for the portion of the tensor assigned to it
  3. Different devices need to shard and gather activation values in order to pass partial results to the next processing device
  4. After all devices finish computing, aggregate the gradients from each device and update the corresponding model tensor parameters.
  5. Distribute the updated model tensor parameters back to each device, preparing for the next batch of training.

Core idea:

Split a single layer or weight matrix by row/column across different worker devices for parallel computation. This is an extremely fine-grained form of model parallelism that can maximize the use of every device’s compute power, but it also adds a large amount of communication overhead. Different workers need to frequently shard and gather intermediate activation values and gradients.

Applicable scenarios:

  • Extremely large models where a single layer has a huge parameter count, such as billion-parameter-scale large language models like GPT-3

Characteristics:

  • Requires the cluster to provide a large amount of GPU memory
  • The largest communication overhead; intermediate results must be passed frequently
  • The tensor-splitting strategy must be designed carefully to reduce overhead

3. Pipeline Parallelism

Characteristics:

  1. Split the entire model by layer into multiple consecutive stages, with each stage’s computation handled by one device.
  2. At the start of each training iteration, the first device takes a batch of input data and performs the forward computation.
  3. The first device passes the computed intermediate activations to the device of the second stage.
  4. After receiving the activations, the second device continues the forward computation based on this input and passes the result to the next stage, and so on.
  5. This continues until the last stage finishes the forward computation and produces the final output.
  6. Based on the output, compute the loss function and perform backpropagation.
  7. During backpropagation, each stage computes the gradients it needs and passes the upstream gradients to the previous stage.
  8. After all stages finish computing, aggregate their respective gradients and update the corresponding model parameters.
  9. Distribute the updated model parameters to the corresponding devices, preparing for the next batch’s training iteration.

Core idea:

Split the model into multiple consecutive stages, each handled by one worker device. Multiple workers can participate in the forward/backward passes of different batches at the same time, forming pipeline-style parallelism and thereby improving overall throughput. Different stages must shard and pass activations and gradients across devices, achieving a fusion of model parallelism and pipeline parallelism.

Applicable scenarios:

  • Suited to training long models over sequence data

Characteristics:

  • The model can be relatively large, but each stage uses less GPU memory
  • Depends on the cluster providing enough device resources
  • Communication overhead is moderate, mainly passing data at stage transitions
  • Splitting stages sensibly maximizes pipeline parallelism efficiency

WeChat Official Account
WRITTEN BY
WeChat Official Account