Distributed training uses multiple GPUs to train a model when one GPU limits training speed or cannot hold the model. The two main approaches are data parallelism, which splits the data, and model parallelism, which splits the model. section 18.2.1
keras$distribution$set_distribution( keras$distribution$DataParallel() )
For data parallelism with Keras and JAX, set the distribution strategy before building the model. Keras then places a model replica on each GPU. section 18.2.2
A batch across two GPUs
With a batch of 128, two replicas each process 64 samples. During training, their gradients are averaged and the model copies are updated. This is the usual choice when the model fits on one GPU. section 18.2.1
Model parallelism instead places different parts of one model on different devices to process a batch together. It is chiefly useful when the model is too large for one device, not as the usual way to speed up a model that already fits. Both approaches can be combined: four devices can hold one split model while another four hold its replica. section 18.2.1
device_mesh <- keras$distribution$DeviceMesh(
shape = shape(2, 4), #1
axis_names = c("data", "model") #2
)
A two-by-four device mesh names separate data and model axes for that combined arrangement. Model parallelism also requires a layout map specifying how model variables are distributed. section 18.2.2
More GPUs do not guarantee proportional speedups: combining updates takes time, and each GPU needs enough samples to stay busy. A small global batch can leave the resulting per-GPU batches too small. section 18.2.2
