A bigger model

Concise Implementation for Multiple GPUs

Concise multi-GPU training

The previous section did data-parallel training the hard way — manual all_reduce, manual replica management. In practice, every framework wraps it in a one-liner:

  • PyTorch: nn.DataParallel(net) (multi-GPU on one host) or nn.parallel.DistributedDataParallel (multi-host).
  • MXNet: gluon.Trainer(..., kvstore='device').
  • TensorFlow: tf.distribute.MirroredStrategy().

Same numerical result; orders of magnitude less boilerplate; NCCL all-reduce under the hood.

from d2l import tensorflow as d2l
import tensorflow as tf
import keras

We use a small ResNet for these experiments — the speedup from data parallelism only matters once the per-GPU compute is non-trivial:

def resnet18(num_classes, in_channels=1):
    """A slightly modified ResNet-18 model built with Keras."""
    def resnet_block(num_channels, num_residuals, first_block=False):
        blk = tf.keras.Sequential()
        for i in range(num_residuals):
            if i == 0 and not first_block:
                blk.add(d2l.Residual(num_channels, use_1x1conv=True,
                                     strides=2))
            else:
                blk.add(d2l.Residual(num_channels))
        return blk

    # Smaller conv, no max-pool (same as the PT version)
    net = tf.keras.Sequential([
        tf.keras.layers.Conv2D(64, kernel_size=3, strides=1, padding='same'),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Activation('relu'),
        resnet_block(64, 2, first_block=True),
        resnet_block(128, 2),
        resnet_block(256, 2),
        resnet_block(512, 2),
        tf.keras.layers.GlobalAveragePooling2D(),
        tf.keras.layers.Dense(num_classes),
    ])
    return net

Multi-GPU initialization

Wrap the model in the framework’s data-parallel container. Parameters are replicated to each GPU automatically:

# MirroredStrategy distributes training across all available GPUs
strategy = tf.distribute.MirroredStrategy()
print(f'Number of devices: {strategy.num_replicas_in_sync}')
# The model will be created inside strategy.scope() in the training function
Number of devices: 4

Parallel evaluation

The wrapper also handles inference — splits the input minibatch across replicas, gathers outputs:

Training loop

The loop looks like ordinary single-GPU training because the wrapper owns the distributed work:

  • scatter each minibatch across devices;
  • run the same model replica on each shard;
  • average gradients across replicas;
  • step one synchronized set of parameters.

The important lesson is the interface: after wrapping the model, most training code should not need to know how many GPUs are present.

Single-GPU baseline

train(num_gpus=1, batch_size=256, lr=0.1)
test acc: 0.93, 105.1 sec total on ['/device:GPU:0']

Use this as the throughput baseline before the data-parallel wrapper adds replication and gradient averaging.

Two GPUs

train(num_gpus=2, batch_size=512, lr=0.2)
test acc: 0.91, 92.0 sec total on ['/device:GPU:0', '/device:GPU:1']

The training loop is unchanged; the wrapper splits the minibatch and synchronizes gradients under the hood.

Recap

  • Framework wrappers (DataParallel, MirroredStrategy) reduce data-parallel SGD to one line of setup.
  • Same numerical recipe as the from-scratch version: replicate, split, all-reduce, identical step.
  • For multi-host distributed training, use DistributedDataParallel / MultiWorkerMirroredStrategy — same idea, NCCL/Gloo across the network.