How to run distributed training in Keras

Data parallel training keeps identical model weights on several devices while each device processes a different slice of the same global batch. Keras can manage that split through the familiar fit() loop when a DataParallel distribution is active before the model is created.

The keras.distribution implementation currently uses the JAX backend. TensorFlow and PyTorch Keras projects use tf.distribute or DistributedDataParallel instead, so select the JAX backend before the first Keras import for this path.

A CPU rehearsal can expose two logical JAX devices through XLA_FLAGS without changing the host configuration. Replace the CPU device selection with gpu or tpu on an accelerator host, remove the CPU-only flag, and keep the global batch size divisible by the number of devices.

1. Create the distribution section in train_distributed.py.

   train_distributed.py

   import numpy as np
    
   import keras
   from keras import layers
    
    
   keras.utils.set_random_seed(7)
    
   devices = keras.distribution.list_devices("cpu")
   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)

   DataParallel replicates model variables across the listed devices and splits each batch along its first dimension.

2. Append the reproducible training arrays after set_distribution(distribution).

   rng = np.random.default_rng(7)
   x = rng.normal(size=(128, 4)).astype("float32")
   y = (x[:, 0] > 0).astype("float32")

3. Append the model definition after the target array assignment.

   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"],
   )

4. Append the training and evaluation block at the end of the compilation section.

   model.fit(x, y, epochs=4, batch_size=16, verbose=2)
   _, accuracy = model.evaluate(x, y, verbose=0)
    
   if accuracy < 0.85:
       raise RuntimeError(f"Evaluation accuracy below 0.85: {accuracy:.4f}")
    
   print(f"Devices: {devices}")
   print(f"Distribution: {type(keras.distribution.distribution()).__name__}")
   print(f"Evaluation accuracy: {accuracy:.4f}")

   The evaluation uses the trained model on the generated samples and exits with an error when accuracy stays below the stated threshold.

5. Review the assembled program as one complete file.

   train_distributed.py

   import numpy as np
    
   import keras
   from keras import layers
    
    
   keras.utils.set_random_seed(7)
    
   devices = keras.distribution.list_devices("cpu")
   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)
    
   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"],
   )
    
   model.fit(x, y, epochs=4, batch_size=16, verbose=2)
   _, accuracy = model.evaluate(x, y, verbose=0)
    
   if accuracy < 0.85:
       raise RuntimeError(f"Evaluation accuracy below 0.85: {accuracy:.4f}")
    
   print(f"Devices: {devices}")
   print(f"Distribution: {type(keras.distribution.distribution()).__name__}")
   print(f"Evaluation accuracy: {accuracy:.4f}")

6. Run the completed script to verify DataParallel training across two logical CPU devices.

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

   Multiple accelerators use gpu or tpu device selection without the CPU-only XLA_FLAGS value.