Data parallelism speeds up training by putting a copy of the same model on each device and giving each copy different data to process. It is a practical way past the speed limit of a single GPU—as long as the whole model fits on one device. section 18.2.1 section 18.4
A batch shared by two GPUs
Suppose the batch size is 128. Two GPUs each hold a model replica and process a different sub-batch of 64 samples. During training, the gradients from the sub-batches are averaged to update the replicas. During inference, their predictions are concatenated to recover predictions for the full batch. section 18.2.1
keras$distribution$set_distribution( keras$distribution$DataParallel() )
With Keras and JAX, set a DataParallel distribution before building the model to use data parallelism. section 18.2.2
keras$distribution$set_distribution(
keras$distribution$DataParallel(devices = c("gpu:0", "gpu:1"))
)
Pass device names to DataParallel when you want to choose the GPUs explicitly. Use keras$distribution$list_devices() to find their names. section 18.2.2
Two practical limits matter:
- Keep the global batch large enough. Splitting a small batch can leave each GPU with too little work to stay busy. section 18.2.2
- Expect less than a perfect N-fold speedup on N GPUs because merging results across devices takes time. Model parallelism serves a different need: splitting a model that is too large for one device. The two approaches can also be combined. section 18.2.1 section 18.2.2
