Efficient Training on Multiple GPUs

Efficient Training on Multiple GPUs

• single GPU is too slow or the model weights don’t fit in a single GPUs memory → mutli-GPU setup
• requires some form of parallelism as the work needs to be distributed
• parallism such as data, tensor, or pipeline parallism

have a look at it before diving into the following sections: single GPU section

Concepts

• DataParallel (DP)
  • the same setup is replicated multiple times
  • each being fed a slice of the data
  • The processing is done in parallel and all setups are synchronized at the end of each training step.
• TensorParallel (TP)
  • each tensor is split up into multiple chunks
  • so instead of having the whole tensor reside on a single gpu, each shard of the tensor resides on its designated gpu
  • During processing each shard gets processed separately and in parallel on different GPUs and the results are synced at the end of the step
  • This is what one may call horizontal parallelism, as the splitting happens on horizontal level
• PipelineParallel (PP)
  • the model is split up vertically (layer-level) across multiple GPUs,
  • so that only one or several layers of the model are places on a single gpu
  • Each gpu processes in parallel different stages of the pipeline and working on a small chunk of the batch.
• Zero Redundancy Optimizer (ZeRO)
  • Also performs sharding of the tensors somewhat similar to TP, except the whole tensor gets reconstructed in time for a forward or backward computation
  • therefore the model doesn’t need to be modified
  • It also supports various offloading techniques to compensate for limited GPU memory.
• Sharded DDP
  • another name for the foundational ZeRO concept as used by various other implementations of ZeRO.

Scalability Strategy

• Single Node / Multi-GPU

  • Model fits onto a single GPU

    1. DDP - Distributed DP
    2. ZeRO - may or may not be faster depending on the situation and configuration used

  • Model doesn’t fit onto a single GPU

    1. PP

    2. ZeRO

    3. TP

       With very fast intra-node connectivity of NVLINK or NVSwitch

       all three should be mostly on par,

       Without these

       PP will be faster than TP or ZeRO.

       The degree of TP may also make a difference. Best to experiment to find the winner on your particular setup.

       TP is almost always used within a single node. That is TP size <= gpus per node.

  • Largest Layer not fitting into a single GPU

    1. If not using ZeRO - must use TP, as PP alone won’t be able to fit.
    2. With ZeRO see the same entry for “Single GPU” above
• Multi-Node / Multi-GPU
  • When you have fast inter-node connectivity:
    1. ZeRO - as it requires close to no modifications to the model
    2. PP+TP+DP - less communications, but requires massive changes to the model
  • when you have slow inter-node connectivity and still low on GPU memory:
    1. DP+PP+TP+ZeRO-1