How to run distributed training in Keras

Keras can keep the familiar fit() training loop while a distribution object decides how batches and variables move across devices. Use DataParallel when a JAX-backed Keras project needs a small multi-device training smoke test before the same model code moves to GPUs or TPUs.

Keras 3's keras.distribution path uses the JAX backend for DataParallel and ModelParallel distribution. TensorFlow and PyTorch Keras jobs use their backend-native distribution APIs, so set the backend before the first Keras import and choose the distribution path that matches the runtime.

The CPU rehearsal exposes two JAX CPU devices with XLA_FLAGS so the device list, distribution object, and fit() output can be checked on a workstation. On an accelerator host, keep the same Keras sequence, change device_type to gpu or tpu, and remove the CPU-only flag.

1. Create a training script that sets the distribution before building the model.

   train_distributed.py

   import numpy as np
    
   import keras
   from keras import layers
    
    
   keras.utils.set_random_seed(7)
    
   device_type = "cpu"
   devices = keras.distribution.list_devices(device_type)
   if len(devices) < 2:
       raise SystemExit(f"Expected at least 2 CPU devices, found: {devices}")
    
   distribution = keras.distribution.DataParallel(devices=devices)
   keras.distribution.set_distribution(distribution)
    
   print(f"Devices: {devices}")
   print(f"Distribution: {type(keras.distribution.distribution()).__name__}")
    
   rng = np.random.default_rng(7)
   x = rng.normal(size=(128, 4)).astype("float32")
   y = (x[:, 0] > 0).astype("float32")
    
   model = keras.Sequential(
       [
           layers.Input(shape=(4,)),
           layers.Dense(1, activation="sigmoid"),
       ]
   )
   model.compile(
       optimizer=keras.optimizers.SGD(learning_rate=0.2),
       loss="binary_crossentropy",
       metrics=["accuracy"],
   )
    
   history = model.fit(x, y, epochs=4, batch_size=16, verbose=2)
   print(f"Final accuracy: {history.history['accuracy'][-1]:.4f}")
    
   keras.distribution.set_distribution(None)

   DataParallel replicates model variables and splits each batch across the listed devices. Keep model construction and compile() after set_distribution() so newly created variables belong to the distribution.

2. Run the script with the JAX backend and two CPU devices visible to JAX.

   $ KERAS_BACKEND=jax XLA_FLAGS=--xla_force_host_platform_device_count=2 python train_distributed.py
   Devices: ['cpu:0', 'cpu:1']
   Distribution: DataParallel
   Epoch 1/4
   8/8 - 0s - 51ms/step - accuracy: 0.7734 - loss: 0.4660
   Epoch 2/4
   8/8 - 0s - 800us/step - accuracy: 0.8672 - loss: 0.3990
   Epoch 3/4
   8/8 - 0s - 729us/step - accuracy: 0.8906 - loss: 0.3547
   Epoch 4/4
   8/8 - 0s - 697us/step - accuracy: 0.9453 - loss: 0.3243
   Final accuracy: 0.9453

   The XLA_FLAGS value is only for CPU-only rehearsal. On a host with multiple accelerators, remove it and set device_type to gpu or tpu in the script.

3. Clear the CPU-only device override before launching other JAX jobs from the same shell.

   $ unset XLA_FLAGS