title multigpu distributed training with pytorch fchollethttps twitter comfchollet date created 20230629 last modified 20230629 description guide to multigpu training for keras models with pytorch 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 pytorch s distributeddataparallel module wrapper to train keras with minimal changes to your code on multiple gpus 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 a tensordataset next let s define a simple pytorch training loop that targets a gpu note the calls to cuda forward pass backward and optimize print loss statistics singlehost multidevice synchronous training in this setup you have one machine with several gpus 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 mirrored variable object how to use it to do singlehost multidevice synchronous training with a keras model you would use the torch nn parallel distributeddataparallel module wrapper here s how it works we use torch multiprocessing startprocesses to start multiple python processes one per device each process will run the perdevicelaunchfn function the perdevicelaunchfn function does the following it uses torch distributed initprocessgroup and torch cuda setdevice to configure the device to be used for that process it uses torch utils data distributed distributedsampler and torch utils data dataloader to turn our data into a distributed data loader it also uses torch nn parallel distributeddataparallel to turn our model into a distributed pytorch module it then calls the trainmodel function the trainmodel function will then run in each process with the model using a separate device in each process here s the flow where each step is split into its own utility function config device setup setup the process groups prepare the dataloader instantiate the torch optimizer instantiate the torch loss function put model on device time to start multiple processes we use the fork method rather than spawn to support notebooks 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 pytorch s distributeddataparallel module wrapper to train keras with minimal changes to your code on multiple gpus 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 a tensordataset next let s define a simple pytorch training loop that targets a gpu note the calls to cuda forward pass backward and optimize print loss statistics single host multi device synchronous training in this setup you have one machine with several gpus 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 mirrored variable object how to use it to do single host multi device synchronous training with a keras model you would use the torch nn parallel distributeddataparallel module wrapper here s how it works we use torch multiprocessing start_processes to start multiple python processes one per device each process will run the per_device_launch_fn function the per_device_launch_fn function does the following it uses torch distributed init_process_group and torch cuda set_device to configure the device to be used for that process it uses torch utils data distributed distributedsampler and torch utils data dataloader to turn our data into a distributed data loader it also uses torch nn parallel distributeddataparallel to turn our model into a distributed pytorch module it then calls the train_model function the train_model function will then run in each process with the model using a separate device in each process here s the flow where each step is split into its own utility function config device setup setup the process groups prepare the dataloader instantiate the torch optimizer instantiate the torch loss function put model on device time to start multiple processes we use the fork method rather than spawn to support notebooks that s it
import os

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

import torch
import numpy as np
import keras


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_dataset():
    
    (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)

    
    dataset = torch.utils.data.TensorDataset(
        torch.from_numpy(x_train), torch.from_numpy(y_train)
    )
    return dataset





def train_model(model, dataloader, num_epochs, optimizer, loss_fn):
    for epoch in range(num_epochs):
        running_loss = 0.0
        running_loss_count = 0
        for batch_idx, (inputs, targets) in enumerate(dataloader):
            inputs = inputs.cuda(non_blocking=True)
            targets = targets.cuda(non_blocking=True)

            
            outputs = model(inputs)
            loss = loss_fn(outputs, targets)

            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            running_loss_count += 1

        
        print(
            f"Epoch {epoch + 1}/{num_epochs}, "
            f"Loss: {running_loss / running_loss_count}"
        )





num_gpu = torch.cuda.device_count()
num_epochs = 2
batch_size = 64
print(f"Running on {num_gpu} GPUs")


def setup_device(current_gpu_index, num_gpus):
    
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "56492"
    device = torch.device("cuda:{}".format(current_gpu_index))
    torch.distributed.init_process_group(
        backend="nccl",
        init_method="env://",
        world_size=num_gpus,
        rank=current_gpu_index,
    )
    torch.cuda.set_device(device)


def cleanup():
    torch.distributed.destroy_process_group()


def prepare_dataloader(dataset, current_gpu_index, num_gpus, batch_size):
    sampler = torch.utils.data.distributed.DistributedSampler(
        dataset,
        num_replicas=num_gpus,
        rank=current_gpu_index,
        shuffle=False,
    )
    dataloader = torch.utils.data.DataLoader(
        dataset,
        sampler=sampler,
        batch_size=batch_size,
        shuffle=False,
    )
    return dataloader


def per_device_launch_fn(current_gpu_index, num_gpu):
    
    setup_device(current_gpu_index, num_gpu)

    dataset = get_dataset()
    model = get_model()

    
    dataloader = prepare_dataloader(
        dataset, current_gpu_index, num_gpu, batch_size
    )

    
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    
    loss_fn = torch.nn.CrossEntropyLoss()

    
    model = model.to(current_gpu_index)
    ddp_model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[current_gpu_index], output_device=current_gpu_index
    )

    train_model(ddp_model, dataloader, num_epochs, optimizer, loss_fn)

    cleanup()




if __name__ == "__main__":
    
    torch.multiprocessing.start_processes(
        per_device_launch_fn,
        args=(num_gpu,),
        nprocs=num_gpu,
        join=True,
        start_method="fork",
    )