A single GPU limits how quickly a growing model can train—and sometimes whether it fits at all. Distributed training uses multiple devices in two main ways: split the data across copies of a model, or split the model itself across devices. section 18.2.1
Two GPUs, one batch
With data parallelism, each GPU holds a complete model replica. Split a batch of 128 into two batches of 64; each replica processes one half. During training, the results from the devices must then be merged. This is the usual choice when the model fits on one GPU. section 18.2.1 section 18.2.2
Choose the split that matches the constraint:
- Data parallelism splits samples across model replicas. Use a sufficiently large global batch so each GPU has enough work. section 18.2.1 section 18.2.2
- Model parallelism splits one model across devices. Reserve it chiefly for models too large for one device, rather than as a speed trick for ordinary models. section 18.2.1
- Combine them when needed: split a model across four devices, then replicate that arrangement twice to use eight devices. section 18.2.1
keras.distribution.set_distribution(keras.distribution.DataParallel())
For Keras with the JAX backend, set data parallelism before building the model. Device names can also be supplied explicitly when only selected GPUs should participate. section 18.2.2
device_mesh = keras.distribution.DeviceMesh( shape=(2, 4), #1 axis_names=["data", "model"], #2 )
For a combined setup, a device mesh names separate data and model axes; this example assigns two positions to data and four to the model. section 18.2.2
Model-parallel configuration also needs a layout map specifying how model variables use the mesh, followed by a ModelParallel distribution setting. Scaling is not free: merging updates adds overhead. Approximate speedups given for two, four, and eight GPUs are 2×, 3.8×, and 7.3×, assuming batches large enough to keep them busy. A rented multi-GPU VM can avoid much of the hardware setup. section 18.2.2 section 18.2.2
