TopolearnLarge models need weights split across accelerators; Keras exposes this declaratively without rewriting the model.
`keras.distribution` describes how to place variables and data on many devices. It is currently implemented for the JAX backend (other backends use their native tools, like tf.distribute or PyTorch DDP/FSDP).
The building blocks:
keras.distribution.list_devices()lists accelerators.- A `DeviceMesh(shape, axis_names, devices)` arranges devices in a named grid, for example 2 x 4 with axes
("batch", "model"). - `DataParallel()`: replicate the model, split each batch across all devices. One line for multi-GPU training.
- `ModelParallel(layout_map=..., batch_dim_name="batch")`: shard variables according to a `LayoutMap`, which maps regular expressions over variable paths to layouts. A layout like
(None, "model")splits a kernel's second axis across the mesh'smodelaxis and replicates the first.
Activate a distribution with keras.distribution.set_distribution(...) before creating the model, so variables are created already sharded; then fit() works as usual and the batch is split along the batch axis.
Sharded dimensions must be divisible by the mesh axis size. Under the hood Keras uses JAX's sharding and the XLA compiler inserts the required communication, so the model code itself stays unchanged.
Example
import os
os.environ["KERAS_BACKEND"] = "jax"
import keras
devices = keras.distribution.list_devices() # e.g. 8 GPUs or TPU cores
mesh = keras.distribution.DeviceMesh(
shape=(2, len(devices) // 2), axis_names=("batch", "model"), devices=devices)
layout_map = keras.distribution.LayoutMap(mesh)
layout_map["dense_1/kernel"] = (None, "model") # split output columns across "model"
layout_map["dense_1/bias"] = ("model",)
keras.distribution.set_distribution(
keras.distribution.ModelParallel(layout_map=layout_map, batch_dim_name="batch"))
model = keras.Sequential([keras.Input((1024,)),
keras.layers.Dense(4096, name="dense_1"),
keras.layers.Dense(10, name="dense_2")])
# Variables are created sharded; fit() splits each batch along "batch".
# Pure data parallelism instead:
# keras.distribution.set_distribution(keras.distribution.DataParallel())Remember
- DeviceMesh names device axes; LayoutMap maps variable paths to layouts
- DataParallel for replication, ModelParallel for sharded weights
- Set the distribution before building the model (JAX backend)
Watch out for
- Setting the distribution after the model exists, so its variables are not sharded.