Multi-GPU Training

Two ways to spread a model

When one GPU isn't enough, you split something across many. There are two things you can split, and they answer different shortages:

Data parallelism — put a full copy of the model on every GPU and feed each a different slice of the batch. The gradients are averaged across devices every step. This is the answer when the model fits but training is slow.
Model parallelism — the model itself is too big for one GPU, so you split its weights across devices. This is the answer when you hit out-of-memory before you ever hit slow.

The Keras 3 distribution API

Keras 3 unifies both under keras.distribution. For data parallel you build a DataParallel object from your devices and call set_distribution() before you build the model — then fit() shards the batches and reduces the gradients for you, with no change to the training code. Model parallel is more deliberate: you declare a DeviceMesh, then a LayoutMap that says how each weight tensor is sharded across the mesh axes.

Reach for the simple one first

Data parallel scales close to linearly out to roughly 8 GPUs and is almost zero extra code. Model parallel is fiddly and only justified once a single replica genuinely won't fit — so measure before you reach for it.

Code

Data parallel vs model parallel with keras.distribution·python

# Data parallelism: same model on each GPU, split data
devices = keras.distribution.list_devices("gpu")
data_parallel = keras.distribution.DataParallel(devices=devices)

# Set distribution before building the model
keras.distribution.set_distribution(data_parallel)

model = build_model()
model.compile(optimizer="adam", loss="mse")
model.fit(x_train, y_train)  # Automatically distributed!

# Model parallelism: split model across GPUs
device_mesh = keras.distribution.DeviceMesh(
    shape=(2,), axis_names=["model"], devices=devices
)
layout_map = keras.distribution.LayoutMap(device_mesh)
layout_map["dense/kernel"] = keras.distribution.TensorLayout(["model", None])

Two ways to spread a model

The Keras 3 distribution API

Reach for the simple one first

Code

External links

Exercise

Progress

Comments 0