Data parallelism speeds up training by giving each GPU a complete copy of the model and a different share of each batch. For example, with two GPUs, a batch of 128 becomes two sub-batches of 64, processed at the same time. Training then merges the updates from the replicas. The model must fit on each device. section 18.2.1 section 18.4
keras.distribution.set_distribution(keras.distribution.DataParallel())
With Keras and JAX, set DataParallel before building the model to use data parallelism. section 18.2.2
A few practical checks matter:
- To choose particular GPUs, list the available devices and pass their names to DataParallel. section 18.2.2
- Keep the global batch large enough that each GPU receives enough work. Merging updates adds overhead: reported speedups are about 2× on two GPUs, 3.8× on four and 7.3× on eight. section 18.2.2
- If the model cannot fit on one device, use model parallelism to split its variables and computation; the two approaches can also be combined. section 18.2.1 section 18.4







