Multi-GPU Training

Dingo supports data-parallel training across multiple GPUs using PyTorch DDP (DistributedDataParallel). This tutorial explains how DDP works in Dingo, why simply setting num_gpus: 8 is not enough to achieve a speedup, how to read the extended training log, and which hyperparameters require tuning.

How DDP works in Dingo

The goal of DDP is to increase the data throughput during training by processing subsets of the mini-batch in parallel on separate GPUs. Each GPU holds a full copy of the DINGO model, but each model sees a different, random subset of the waveform data set. After every backward pass, DDP performs a gradient all-reduce: it aggregates gradients across all GPUs resulting in a gradient update based on the full effective batch size, not the per-GPU batch size. The resulting averaged gradient is used to update all models equivalently such that every model replica stays the same. The optimizer step is identical to single-GPU training with the effective batch size.

Through DDP, it is possible to increase the effective batch size beyond the single-GPU memory limit. Now, the limiting factor for training throughput is the memory limit of a single GPU which processes the effective batch size divided by the number of GPUs. Out-of-memory errors can be addressed by either increasing the number of GPUs or by reducing the effective batch size.

There are certain caveats to DDP:

  • Since the gradients have to be synced across GPUs, the interconnect between GPUs becomes highly relevant. Slow interconnects can significantly slow down training. Therefore, it is recommended to only run DDP on GPUs that are located on the same node where fast interconnect is independent of network traffic.

  • In addition to the full model replica (with its gradients and optimizer state) and the batch subset, additional GPU memory is required for syncing and aggregating gradients. Therefore, the per-GPU batch size can be lower than in single-GPU training.

Compute requirements

  • A single node with N GPUs connected by fast interconnect (tested on A100 and H100 nodes).

  • An NCCL-capable PyTorch build (the default for any CUDA-enabled installation).

  • When submitting via HTCondor, request_gpus is derived automatically from num_gpus — no changes to the condor: block are needed. Cluster-specific submit-file directives (e.g., full-node templates) can be passed through verbatim via condor: extra_submit_lines:.

Settings changes

Multi-GPU training is enabled within Dingo by changing num_gpus: 1 to the number of available GPUs. Additionally, it is recommended to scale batch_size and lr in each training stage (see explanations below):

# single-GPU baseline
local:
  device: cuda
  num_workers: 16
  num_gpus: 1
training:
  stage_0:
    batch_size: 4096
    optimizer:
      lr: 5.0e-5
  stage_1:
    batch_size: 4096
    optimizer:
      lr: 1.0e-5

# 8-GPU DDP — scale batch_size and lr by num_gpus
local:
  device: cuda
  num_workers: 32   # total across GPUs (= 4 per GPU); scale with num_gpus
  num_gpus: 8
  # ddp_port: 12355  # change when running several DDP jobs on one node
training:
  stage_0:
    batch_size: 32768   # = 4096 × 8
    optimizer:
      lr: 4.0e-4        # = 5e-5 × 8
  stage_1:
    batch_size: 32768
    optimizer:
      lr: 8.0e-5        # = 1e-5 × 8

Increasing the batch_size

Dingo interprets batch_size as the total effective batch size across all GPUs and divides it equally:

per-GPU batch size = batch_size / num_gpus

Pytorch’s DistributedSampler also splits the waveform dataset equally, so each GPU sees 1/N of the data per epoch. Putting both together, the number of optimizer steps per epoch is:

steps / epoch  =  dataset_size × train_fraction / batch_size

This formula is independent of num_gpus. The table below shows the consequences:

num_gpus

batch_size

per-GPU batch size

steps/epoch

outcome

1

4 096

4 096

N

baseline

8

4 096

512

N

same steps + all-reduce overhead → slower than the baseline

8

32 768

4 096

N/8

fewer steps, good GPU utilisation → faster than the baseline

When the batch_size is scaled by num_gpus, the actual wall-clock speedup is somewhat less than 8× because the gradient all-reduce adds overhead per step (visible as inflated Time Network in the log); a factor of 4–6× is typical.

In general, it is recommended to increase the batch size as much as possible to fully utilize the GPU memory.

Learning rate scaling

When the effective batch size increases by a factor of N, the learning rate should be increased by the same factor (linear scaling rule, [8]):

lr_multi  =  lr_single × num_gpus

Intuition: with N× more samples per gradient update the gradient estimate has lower variance, so a proportionally larger step can be taken without destabilising training.

Every training stage has to be adapted independently.

Normalization layers

BatchNorm computes its statistics from the batch on each GPU, so under DDP every replica normalizes with slightly different statistics, and synchronizing them (SyncBatchNorm) stalls the forward pass at every normalization layer. Use LayerNorm instead, which normalizes each sample independently and needs no communication between GPUs. The norm option of the flow and embedding network settings selects the normalization layer (LayerNorm, BatchNorm or null for none). The single-GPU example settings use BatchNorm; for DDP, switch both to LayerNorm:

model:
  posterior_kwargs:
    base_transform_kwargs:
      norm: LayerNorm
  embedding_kwargs:
    norm: LayerNorm

Checkpoints trained with the former batch_norm: True setting keep their BatchNorm layers when loaded. If a network with BatchNorm layers is trained on multiple GPUs, they are converted to SyncBatchNorm (with a warning), so that the batch statistics are still computed over the full effective batch.

LayerNorm is only available for base_transform_type: rq-coupling. For rq-autoregressive transforms it would break the causal structure of the MADE layers, so BatchNorm or null must be used there.

Compiling the network (torch_compile)

The neural spline flow launches tens of thousands of small CUDA kernels per step, so on a modern GPU the step is bound by kernel-launch overhead rather than by arithmetic (GPU utilization plateaus well below 100% even at large batch sizes). Setting torch_compile: true in the local section wraps the network with torch.compile, which fuses these kernels. The gain depends on how launch-bound the network is: for the npe_model example network (30 flow steps, hidden_dim 1024) trained on a 10M-waveform dataset at a per-GPU batch size of 4096 (A100), the training step is 1.22× faster on one GPU and 1.14× faster on four GPUs; with hidden_dim 512 it is 1.4× / 1.35×. Peak GPU memory drops by about 20%. The speedup is independent of num_gpus and larger at smaller per-GPU batch sizes.

local:
  torch_compile: true                    # default: false
  torch_compile_cache_dir: /scratch/tmp  # optional; see note below

Notes:

  • Compilation is slow: 4–12 minutes for a production-size network, paid on the first training step of a run and again at every stage boundary that changes which parameters are trainable. The compiled steps that follow are the fast ones, so torch_compile pays off for trainings of tens of epochs or more, not for short runs. To avoid further compilations the training loader drops the last, smaller batch of each epoch (the compiled graph is specialized to the batch shape; the test loader keeps it) and the test epoch runs the network eagerly (an eval-mode graph would cost another compilation that a short test epoch never amortizes).

  • A stage boundary that changes freeze_rb_layer discards the compiled graphs and the next step recompiles (a compiled graph does not track which parameters are trainable).

  • torch_compile_cache_dir sets the base directory of the on-disk Inductor/Triton cache (default: the system temp directory); under DDP each rank gets its own subdirectory. That cache must live on node-local disk. If the system temp directory is a shared network filesystem, set it to a node-local path (e.g. the HTCondor scratch directory); otherwise a just-compiled kernel can be unloadable on another rank and the run hangs on an NCCL timeout.

  • Once the network step is faster, the dataloader can become the bottleneck: watch Time Dataloader in the log and raise num_workers if needed.

TensorFloat-32 matrix multiplications (float32_matmul_precision)

PyTorch runs float32 matrix multiplications at full precision by default, which leaves the tensor cores of Ampere and newer GPUs unused. Setting

local:
  float32_matmul_precision: high   # default: highest

lets them use TensorFloat-32: matmul inputs are rounded to 10 mantissa bits (the range of float32 is kept, and accumulation, weights, gradients and optimizer state stay float32). For the npe_model network (hidden_dim 1024) this roughly doubles the speed of the network step and combines with torch_compile: with both (and the fused optimizer below) the network step drops from 0.85 s to 0.28 s per 4096 samples on an A100, a 3× gain in GPU time. All non-matmul operations (splines, normalization, the optimizer) are unaffected.

TF32 and automatic_mixed_precision are alternatives, not a stack: under AMP the matmuls already run in float16, so float32_matmul_precision: high changes nothing there. On the npe_model network TF32 alone and AMP alone give the same speedup (0.90 to 0.47 s per step); TF32 is the option for trainings that stay in float32. The value in effect is stored in the checkpoint metadata (float32_matmul_precision).

Fused optimizer (fused)

Unlike the settings above, this one is not in the local section: it belongs to the optimizer block of an individual training stage, and is passed straight through to the PyTorch optimizer. Setting fused: true selects the fused CUDA kernel, which performs the whole optimizer update in one kernel instead of one per step of the update math, and is about 10% faster per step for large networks:

training:
  stage_0:
    optimizer:
      type: adam
      lr: 0.0001
      fused: true   # default: false

Each stage builds its own optimizer, so this has to be repeated in every stage that should use it. It is supported by adam, adamw, adagrad and sgd (not lbfgs); the fused kernel exists for CUDA and (since PyTorch 2.4) CPU tensors, the speedup matters on the GPU.

Freezing layers

It is currently not possible to set freeze_rb_layer: True in DDP. The reason is that when starting the separate DDP processes, it is fixed which network parameters have to be synced across GPUs and GPU memory is allocated accordingly. In the current implementation, the stages are initialized (and therefore freeze_rb_layer: True is set) afterward. If the layer is frozen within a DDP process, the all-reduce fails since the reduced tensors have a different shape than expected. Allowing freeze_rb_layer: True with DDP would require significant restructuring of the code.

Reading the training log

Within Dingo, the separate processes for DDP are spawned after loading the waveform dataset, building the Dingo model, and computing the SVD which initializes the first layer of the embedding network. At this point, a statement like Process group initialised: backend=nccl, rank=0, world_size=8 is printed for every GPU into info.out. Afterward, every print statement within the Dingo code is executed by every GPU process (if this is not specifically prevented). As a result, duplicate information can appear in info.out. Since each GPU process prints outputs to info.out in parallel, the order of these statements can be mixed up compared to single-GPU training. If information is not duplicated, it is only printed by rank = 0.

During training, each printed line looks like the following in single-GPU mode:

Train Epoch: 1 [4096/4750000 (0%)]   Loss: -2.771 (-2.799)   Time Dataloader: 0.043 (0.043)   Time Network: 0.415 (0.415)

In multi-GPU mode, rank 0 additionally prints a third timing column:

Train Epoch: 1 [4096/4750000 (0%)]   Loss: -2.771 (-2.799)   Time Dataloader: 0.018 (0.018)   Time Network: 0.700 (0.700)   Time Loss Aggregation: 0.090 (0.090)

Each value is the current step; the value in parentheses is the running average over the epoch. In DDP mode all three timings are reported as the maximum across all GPUs (dist.ReduceOp.MAX), i.e. they reflect the slowest rank.

Time Dataloader

Wall-clock time spent loading and pre-processing one batch (waveform decompression, projection to detectors, noise generation). Expect this to be smaller than in single-GPU mode because each GPU loads a smaller per-GPU batch.

Time Network

Forward pass + backward pass + DDP gradient all-reduce + optimizer step. Unlike single-GPU training this includes gradient synchronisation, so Time Network per step will typically be larger than in single-GPU mode even though each GPU processes fewer samples. The speedup comes from fewer steps per epoch, not from faster individual steps.

Time Loss Aggregation (DDP only)

Time required to reduce per-GPU losses to rank 0 and compute the global average. Includes a dist.barrier() that waits for the slowest GPU to finish its network pass before proceeding. Occasional spikes can indicate increased wait time at the barrier. Sustained high values could indicate a load imbalance between GPUs.

Guidelines for tuning num_workers

  • Time Dataloader should be clearly smaller than Time Network. If dataloader time dominates, increase num_workers. Like batch_size, num_workers is the total across all GPUs and is divided equally between them (aim for 4–8 per GPU).

  • Time Network being larger per step than single-GPU is expected; this is the all-reduce cost.

  • Time Loss Aggregation should be small relative to Time Network.

Monitoring with WandB

Only rank 0 writes to WandB, so the run appears in the dashboard exactly as a single-GPU run. The loss logged is the globally reduced value (weighted average over all GPU contributions), making curves directly comparable across single- and multi-GPU runs.

Enable WandB in the local section:

local:
  wandb:
    project: dingo
    group: O4   # optional: group related runs for comparison