A concise multi-GPU training loop leverages high-level framework abstractions to orchestrate data-parallel execution. The training function begins by wrapping the neural network in a distributed execution module (such as nn.DataParallel in PyTorch) and moving it to the primary processing device. During each epoch, it iterates over the training minibatches, calling the corresponding minibatch training function to perform the synchronized forward and backward passes. It accumulates the training loss and accuracy, and periodically evaluates the model on the test dataset using multiple GPUs, tracking performance metrics efficiently without manually managing individual device state aggregations.

Concise Multi-GPU Training Loop Implementation

When training a minibatch across multiple GPUs using high-level deep learning APIs, the implementation becomes significantly simpler than a from-scratch approach. The primary simplification is the delegation of gradient synchronization and parameter updates to the framework's optimization algorithms (e.g., calling trainer.step()). Depending on the framework, the data distribution may also be automated. For example, PyTorch's DataParallel allows developers to move the entire minibatch to the primary device, letting the framework automatically scatter the data and parallelize the forward and backward passes. In MXNet, the batch is partitioned across devices manually using a splitting function, but high-level tools handle the parallel gradient aggregation and parameter updates seamlessly.

Claude

The `train_batch` function implements data-parallel training for a single minibatch across multiple GPUs. The procedure follows four sequential stages:

1. **Data splitting:** The minibatch of features and labels is divided across the available devices using `split_batch`.
2. **Per-GPU forward pass and loss:** Each GPU independently computes the model's output and the loss on its local data shard. The losses are summed per device.
3. **Per-GPU backpropagation:** Backpropagation is performed separately on each GPU to compute local gradients.
4. **Gradient synchronization and update:** An `allreduce` operation sums and broadcasts all gradients across GPUs within a `torch.no_grad()` context. Finally, each GPU independently updates its own copy of the model parameters using SGD, scaling the update by the full (unsplit) batch size.

```python
def train_batch(X, y, device_params, devices, lr):
    X_shards, y_shards = split_batch(X, y, devices)
    ls = [loss(lenet(X_shard, device_W), y_shard).sum()
          for X_shard, y_shard, device_W in zip(
              X_shards, y_shards, device_params)]
    for l in ls:
        l.backward()
    with torch.no_grad():
        for i in range(len(device_params[0])):
            allreduce([device_params[c][i].grad
                       for c in range(len(devices))])
    for param in device_params:
        d2l.sgd(param, lr, X.shape[0])
```

Because there are no cross-device dependencies within the computational graph for a single minibatch, the per-GPU computations execute in parallel automatically.

Multi-GPU Minibatch Training Implementation

Dive into Deep Learning

The `train` function orchestrates the full multi-GPU data-parallel training loop. It accepts the number of GPUs, batch size, and learning rate as arguments. The procedure is:

1. **Setup:** Load the Fashion-MNIST dataset, allocate the specified number of GPU devices using `d2l.try_gpu(i)`, and replicate the model parameters to each GPU via `get_params`.
2. **Epoch loop:** For each of $$10$$ epochs, iterate over every minibatch in the training set and call `train_batch` to perform the data-parallel forward pass, gradient synchronization, and parameter update. A synchronization barrier (`torch.cuda.synchronize()`) is called after every minibatch to ensure all GPU operations complete before the next minibatch begins, which is also necessary for accurate epoch timing.
3. **Evaluation:** After each epoch, test accuracy is computed on a single GPU (GPU 0) using `d2l.evaluate_accuracy_gpu`, passing a lambda that wraps the model with the first device's parameters. Although evaluating on only one GPU leaves the others idle, it simplifies the code.

```python
def train(num_gpus, batch_size, lr):
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    device_params = [get_params(params, d) for d in devices]
    num_epochs = 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    timer = d2l.Timer()
    for epoch in range(num_epochs):
        timer.start()
        for X, y in train_iter:
            train_batch(X, y, device_params, devices, lr)
            torch.cuda.synchronize()
        timer.stop()
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(
            lambda x: lenet(x, device_params[0]),
            test_iter, devices[0]),))
    print(f'test acc: {animator.Y[0][-1]:.2f}, '
          f'{timer.avg():.1f} sec/epoch on {str(devices)}')
```

Learn Before

Related

Learn After