title multigpu distributed training with jax fchollethttps twitter comfchollet date created 20230711 last modified 20230711 description guide to multigputpu training for keras models with jax accelerator gpu introduction there are generally two ways to distribute computation across multiple devices data parallelism where a single model gets replicated on multiple devices or multiple machines each of them processes different batches of data then they merge their results there exist many variants of this setup that differ in how the different model replicas merge results in whether they stay in sync at every batch or whether they are more loosely coupled etc model parallelism where different parts of a single model run on different devices processing a single batch of data together this works best with models that have a naturallyparallel architecture such as models that feature multiple branches this guide focuses on data parallelism in particular synchronous data parallelism where the different replicas of the model stay in sync after each batch they process synchronicity keeps the model convergence behavior identical to what you would see for singledevice training specifically this guide teaches you how to use jax sharding apis to train keras models with minimal changes to your code on multiple gpus or tpus typically 2 to 16 installed on a single machine single host multidevice training this is the most common setup for researchers and smallscale industry workflows setup let s start by defining the function that creates the model that we will train and the function that creates the dataset we will train on mnist in this case make a simple convnet with batch normalization and dropout load the data and split it between train and test sets scale images to the 0 1 range make sure images have shape 28 28 1 create tf datasets singlehost multidevice synchronous training in this setup you have one machine with several gpus or tpus on it typically 2 to 16 each device will run a copy of your model called a replica for simplicity in what follows we ll assume we re dealing with 8 gpus at no loss of generality how it works at each step of training the current batch of data called global batch is split into 8 different subbatches called local batches for instance if the global batch has 512 samples each of the 8 local batches will have 64 samples each of the 8 replicas independently processes a local batch they run a forward pass then a backward pass outputting the gradient of the weights with respect to the loss of the model on the local batch the weight updates originating from local gradients are efficiently merged across the 8 replicas because this is done at the end of every step the replicas always stay in sync in practice the process of synchronously updating the weights of the model replicas is handled at the level of each individual weight variable this is done through a using a jax sharding namedsharding that is configured to replicate the variables how to use it to do singlehost multidevice synchronous training with a keras model you would use the jax sharding features here s how it works we first create a device mesh using meshutils createdevicemesh we use jax sharding mesh jax sharding namedsharding and jax sharding partitionspec to define how to partition jax arrays we specify that we want to replicate the model and optimizer variables across all devices by using a spec with no axis we specify that we want to shard the data across devices by using a spec that splits along the batch dimension we use jax deviceput to replicate the model and optimizer variables across devices this happens once at the beginning in the training loop for each batch that we process we use jax deviceput to split the batch across devices before invoking the train step here s the flow where each step is split into its own utility function config initialize all state with build this is the loss function that will be differentiated keras provides a pure functional forward pass model statelesscall function to compute gradients training step keras provides a pure functional optimizer statelessapply replicate the model and optimizer variable on all devices all variables will be replicated on all devices in namedsharding axes not mentioned are replicated all axes here apply the distribution settings to the model variables combine all state in a tuple data will be split along the batch axis display data sharding custom training loop postprocessing model state update to write them back into the model that s it introduction there are generally two ways to distribute computation across multiple devices data parallelism where a single model gets replicated on multiple devices or multiple machines each of them processes different batches of data then they merge their results there exist many variants of this setup that differ in how the different model replicas merge results in whether they stay in sync at every batch or whether they are more loosely coupled etc model parallelism where different parts of a single model run on different devices processing a single batch of data together this works best with models that have a naturally parallel architecture such as models that feature multiple branches this guide focuses on data parallelism in particular synchronous data parallelism where the different replicas of the model stay in sync after each batch they process synchronicity keeps the model convergence behavior identical to what you would see for single device training specifically this guide teaches you how to use jax sharding apis to train keras models with minimal changes to your code on multiple gpus or tpus typically 2 to 16 installed on a single machine single host multi device training this is the most common setup for researchers and small scale industry workflows setup let s start by defining the function that creates the model that we will train and the function that creates the dataset we will train on mnist in this case make a simple convnet with batch normalization and dropout load the data and split it between train and test sets scale images to the 0 1 range make sure images have shape 28 28 1 create tf datasets single host multi device synchronous training in this setup you have one machine with several gpus or tpus on it typically 2 to 16 each device will run a copy of your model called a replica for simplicity in what follows we ll assume we re dealing with 8 gpus at no loss of generality how it works at each step of training the current batch of data called global batch is split into 8 different sub batches called local batches for instance if the global batch has 512 samples each of the 8 local batches will have 64 samples each of the 8 replicas independently processes a local batch they run a forward pass then a backward pass outputting the gradient of the weights with respect to the loss of the model on the local batch the weight updates originating from local gradients are efficiently merged across the 8 replicas because this is done at the end of every step the replicas always stay in sync in practice the process of synchronously updating the weights of the model replicas is handled at the level of each individual weight variable this is done through a using a jax sharding namedsharding that is configured to replicate the variables how to use it to do single host multi device synchronous training with a keras model you would use the jax sharding features here s how it works we first create a device mesh using mesh_utils create_device_mesh we use jax sharding mesh jax sharding namedsharding and jax sharding partitionspec to define how to partition jax arrays we specify that we want to replicate the model and optimizer variables across all devices by using a spec with no axis we specify that we want to shard the data across devices by using a spec that splits along the batch dimension we use jax device_put to replicate the model and optimizer variables across devices this happens once at the beginning in the training loop for each batch that we process we use jax device_put to split the batch across devices before invoking the train step here s the flow where each step is split into its own utility function config initialize all state with build this is the loss function that will be differentiated keras provides a pure functional forward pass model stateless_call function to compute gradients training step keras provides a pure functional optimizer stateless_apply replicate the model and optimizer variable on all devices all variables will be replicated on all devices in namedsharding axes not mentioned are replicated all axes here apply the distribution settings to the model variables combine all state in a tuple data will be split along the batch axis naming axes of the mesh naming axes of the sharded partition display data sharding custom training loop post processing model state update to write them back into the model that s it
import os

os.environ["KERAS_BACKEND"] = "jax"

import jax
import numpy as np
import tensorflow as tf
import keras

from jax.experimental import mesh_utils
from jax.sharding import Mesh
from jax.sharding import NamedSharding
from jax.sharding import PartitionSpec as P


def get_model():
    
    inputs = keras.Input(shape=(28, 28, 1))
    x = keras.layers.Rescaling(1.0 / 255.0)(inputs)
    x = keras.layers.Conv2D(
        filters=12, kernel_size=3, padding="same", use_bias=False
    )(x)
    x = keras.layers.BatchNormalization(scale=False, center=True)(x)
    x = keras.layers.ReLU()(x)
    x = keras.layers.Conv2D(
        filters=24,
        kernel_size=6,
        use_bias=False,
        strides=2,
    )(x)
    x = keras.layers.BatchNormalization(scale=False, center=True)(x)
    x = keras.layers.ReLU()(x)
    x = keras.layers.Conv2D(
        filters=32,
        kernel_size=6,
        padding="same",
        strides=2,
        name="large_k",
    )(x)
    x = keras.layers.BatchNormalization(scale=False, center=True)(x)
    x = keras.layers.ReLU()(x)
    x = keras.layers.GlobalAveragePooling2D()(x)
    x = keras.layers.Dense(256, activation="relu")(x)
    x = keras.layers.Dropout(0.5)(x)
    outputs = keras.layers.Dense(10)(x)
    model = keras.Model(inputs, outputs)
    return model


def get_datasets():
    
    (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

    
    x_train = x_train.astype("float32")
    x_test = x_test.astype("float32")
    
    x_train = np.expand_dims(x_train, -1)
    x_test = np.expand_dims(x_test, -1)
    print("x_train shape:", x_train.shape)
    print(x_train.shape[0], "train samples")
    print(x_test.shape[0], "test samples")

    
    train_data = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    eval_data = tf.data.Dataset.from_tensor_slices((x_test, y_test))
    return train_data, eval_data





num_epochs = 2
batch_size = 64

train_data, eval_data = get_datasets()
train_data = train_data.batch(batch_size, drop_remainder=True)

model = get_model()
optimizer = keras.optimizers.Adam(1e-3)
loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)


(one_batch, one_batch_labels) = next(iter(train_data))
model.build(one_batch)
optimizer.build(model.trainable_variables)




def compute_loss(trainable_variables, non_trainable_variables, x, y):
    y_pred, updated_non_trainable_variables = model.stateless_call(
        trainable_variables, non_trainable_variables, x
    )
    loss_value = loss(y, y_pred)
    return loss_value, updated_non_trainable_variables



compute_gradients = jax.value_and_grad(compute_loss, has_aux=True)



@jax.jit
def train_step(train_state, x, y):
    (
        trainable_variables,
        non_trainable_variables,
        optimizer_variables,
    ) = train_state
    (loss_value, non_trainable_variables), grads = compute_gradients(
        trainable_variables, non_trainable_variables, x, y
    )

    trainable_variables, optimizer_variables = optimizer.stateless_apply(
        optimizer_variables, grads, trainable_variables
    )

    return loss_value, (
        trainable_variables,
        non_trainable_variables,
        optimizer_variables,
    )



def get_replicated_train_state(devices):
    
    var_mesh = Mesh(devices, axis_names=("_"))
    
    var_replication = NamedSharding(var_mesh, P())

    
    trainable_variables = jax.device_put(
        model.trainable_variables, var_replication
    )
    non_trainable_variables = jax.device_put(
        model.non_trainable_variables, var_replication
    )
    optimizer_variables = jax.device_put(optimizer.variables, var_replication)

    
    return (trainable_variables, non_trainable_variables, optimizer_variables)


num_devices = len(jax.local_devices())
print(f"Running on {num_devices} devices: {jax.local_devices()}")
devices = mesh_utils.create_device_mesh((num_devices,))


data_mesh = Mesh(devices, axis_names=("batch",))  
data_sharding = NamedSharding(
    data_mesh,
    P(
        "batch",
    ),
)  


x, y = next(iter(train_data))
sharded_x = jax.device_put(x.numpy(), data_sharding)
print("Data sharding")
jax.debug.visualize_array_sharding(jax.numpy.reshape(sharded_x, [-1, 28 * 28]))

train_state = get_replicated_train_state(devices)


for epoch in range(num_epochs):
    data_iter = iter(train_data)
    for data in data_iter:
        x, y = data
        sharded_x = jax.device_put(x.numpy(), data_sharding)
        loss_value, train_state = train_step(train_state, sharded_x, y.numpy())
    print("Epoch", epoch, "loss:", loss_value)


trainable_variables, non_trainable_variables, optimizer_variables = train_state
for variable, value in zip(model.trainable_variables, trainable_variables):
    variable.assign(value)
for variable, value in zip(
    model.non_trainable_variables, non_trainable_variables
):
    variable.assign(value)