Multi-GPU training (DDP)

[annalena-k]

Multi-GPU training with PyTorch DDP

Summary
This PR adds data-parallel multi-GPU training to Dingo using https://pytorch.org/docs/stable/notes/ddp.html. It also adds gradient accumulation and automatic mixed precision (AMP) as complementary features for memory-efficient training. Together these allow significantly faster training on multi-GPU nodes and greater flexibility in effective batch size.

What's new

Multi-GPU training (DDP)

• Setting num_gpus > 1 in local: spawns one process per GPU via torch.multiprocessing. Each GPU processes a non-overlapping shard of the mini-batch; gradients are averaged across GPUs via NCCL before each parameter update.
• BatchNorm layers are automatically converted to SyncBatchNorm so statistics are computed correctly across all GPUs.
• The DistributedSampler ensures each GPU sees a different, non-overlapping slice of data each epoch.
• Checkpoints are written only by rank 0 and are always saved without the DDP module prefix, so they can be loaded on any number of GPUs (or a single GPU / CPU) without modification.
• Loss and timing statistics are aggregated across ranks so the training log is identical to the single-GPU case.

Gradient accumulation

• New per-stage setting gradient_updates_per_optimizer_step (default: 1). Setting this to $k$ accumulates gradients over $k$ forward-backward passes before calling the optimizer, simulating an effective batch size of $k \times$ batch_size without extra GPU memory.

Automatic mixed precision (AMP)

• New per-stage setting automatic_mixed_precision (default: False). When enabled, the forward pass runs in FP16 while optimizer state and parameter updates stay in FP32. This roughly halves GPU memory usage and increases throughput on GPUs with Tensor Core hardware (A100, V100, H100).

dingo_train_condor

• request_gpus in the HTCondor submission file is now set automatically from local.num_gpus. Users no longer need to specify it in the condor block.

Settings file changes

New optional keys in local:

local:
   device: cuda
   num_gpus: 4        # NEW — default 1 (single GPU). Set >1 for DDP.
   num_workers: 8

New optional keys per training stage

training:
  stage_0:
    batch_size: 512           # Always the TOTAL effective batch size across all GPUs.
    gradient_updates_per_optimizer_step: 2   # NEW — default 1
    automatic_mixed_precision: True          # NEW — default False

condor block: num_gpus removed
num_gpus is no longer needed (or expected) inside condor. It is injected automatically from local.num_gpus when building the submission file:

# Before (old):
  local:
    condor:
      num_gpus: 4   # ← remove this

 # After (new):
  local:
    num_gpus: 4     # ← specify here instead
    condor:
      ...           # num_gpus no longer needed here

Backward compatibility

• num_gpus under local.condor: Existing settings files that specify num_gpus inside the condor block continue to work. A DeprecationWarning is raised at runtime prompting users to move the key to local.num_gpus. This fallback will be removed in a future release.
• Single-GPU / CPU: No changes to existing behaviour. The num_gpus key defaults to 1 if absent.
• Checkpoints: All existing checkpoints load without modification. The new iteration counter (total optimizer steps) is read with a default of 0 for old checkpoints that lack it.
• batch_size semantics: batch_size has always been documented as the total effective batch size. This PR enforces that interpretation — when num_gpus > 1, each GPU receives batch_size / num_gpus samples. No change for single-GPU users.

Requirements

• PyTorch ≥ 2.0 (required for torch.amp.GradScaler and torch.amp.autocast; the deprecated torch.cuda.amp namespace is no longer used)
• NCCL-capable PyTorch build for multi-GPU (standard for any CUDA-enabled install)

Tests

• tests/core/test_multi_gpu.py — covers get_num_gpus (including deprecation warning), build_train_and_test_loaders with and without DDP samplers, LossInfo and RuntimeLimits in a real two-process gloo group, DDP state-dict stripping, and utility functions (set_seed_based_on_rank, replace_BatchNorm_with_SyncBatchNorm).
• tests/core/test_amp.py — covers AMP imports, train_epoch without AMP (including gradient accumulation), and train_epoch with AMP on CUDA (skipped if no GPU).

Integration tests on an htcondor cluster with the NPE example:

• dingo_train on 1 GPU: Training works on A100 8 GPU node with num_gpus: 1
• dingo_train on 8 GPUs: Training works
• dingo_train_condor on 1 GPU: Training works
• dingo_train_condor on 8 GPUs: Training works

To test the performance, I investigated how the training time per epoch scales with the batch size on an NVIDIA A100 8-GPU node with 80 GB GPU memory:

• With a single GPU, we cannot increase the batch size beyond 16384 due to limited GPU memory.
• With 8 GPUs, we can increase the total effective batch size beyond 16384 since each GPU only processes batch_size / num_gpus. However, the overhead from gradient accumulation leads to a scaling which is worse than the ideal projection given by
  $\mathrm{time}(1\mathrm{GPU}) \cdot \frac{\mathrm{batchsize}(1\mathrm{GPU})}{\mathrm{batchsize}(8\mathrm{GPUs}})$
• Automatic mixed precision decreases the run time slightly on multiple GPUs.

[nihargupte-ph]

I'm reviewing this PR at the moment and will leave a more detailed response on the code soon. But for now, I was able to reproduce your plots but with 4 GPUs. Strangely, the time per epoch on my machine also increased at the last test point!

Interestingly, if I replace the Syncbatchnorms with just Batchnorms then the time for epoch does go down a bit more. I think this is because the Syncbatchnorm is not asynchronous unlike the gradient accumulation. This is what I'm showing in the light blue curve above and in red is the current PR. I think the grad accumulation will begin summing the gradients at the beginning of the graph while the gradients for the end of the graph are still being computed, but Syncbatchnorm has to wait.

One possibility we discussed which could be causing the suboptimal scaling was if increasing the batch size wasn't helping because we were compute bound. However, in my experiments I tracked the FLOPs/possible FLOPs and I found this hovered between 60-80%:

so I don't think we're compute bound yet. I am now investigating why we're not getting 100% utilization.

[stephengreen]

Thanks. What if you replace BatchNorm with LayerNorm everywhere? No syncing.

[annalena-k]

Thanks Nihar!
I only included the SynchBatchNorm to ensure that nothing goes wrong for users, but I never used it myself.
If we remove this, I would just add an error message to let users know that BatchNorm doesn't work with multiple GPUs.
Happy to change this if people think this is the way forward :)

[nihargupte-ph]

Right, I think batch norm or layer norm could both work. When computing the syncbatchnorm you sync the batch statistics across all the GPUs. This results in two slowdowns:

1). The communication between each GPU to compute the batch statistics
2). The fact that the GPUs can't run fully async because they need to wait for each other to reach the syncbatch layer.

Of course this is no different than an all reduce. However, the all reduce latency is hidden because the backward pass is occurring asynchronously with the all reduce. However, in the case of the syncbatchnorm, you can't finish the forward pass until you have batch statistics and therefore there is no async.

In our case, the batches are large enough that I think we can get away with just BatchNorm. Layernorm could also work. The only thing is we haven't tested our enets using LayerNorm. What do people think, should we use LayerNorm going forward? This seems like a slightly larger change than just switching to BatchNorm.

[nihargupte-ph]

I was able to find a way to speed up multi and single-GPU training further while profiling the multi-GPU. In particular, I noticed the GPU utilization was not reaching 100% even when the batch size is maxxing out the high bandwidth memory out. This implies we are not compute bound. This can be traced to the following reasoning

1). The neural spline flow fires 25k CUDA kernels per forward pass.
2). Each kernel launch costs 5 micro-s
3). Each kernel runs for 10 micro-s
4). Therefore, T_kernel_launch (5 micro-s * 25k) < T_kernel_run (A100's capacity)

Note this is not related to the dataloader. This is after the data is on the HBM, how much are we utilizing the GPUs FLOPs. So the GPU is finishing computations faster than they can be launched! The solution is kernel fusing using torch.compile(pm.network). This reduces the number of kernels used to 10k. The GPU may still be finishing faster (which is maybe why I'm not getting to 100% utilization on my experiments), but it is faster. One could also up the intensity of the computation (make the kernels take longer).

The problem is torch.compile requires the entire computation graph to have static array shapes. This requires a change to glasflows. We could also consider switching to zuko.

The speedup with torch.compile is not negligible. It can be 2x. The speedup after torch.compile is the green. In red is after using batchnorm instead of syncbatchnorm

  barrier()                    # all 4 ranks aligned
  cuda.synchronize(); t0 = perf_counter()
    forward, backward         # backward includes the DDP gradient all-reduce 
    optimizer.step()
  cuda.synchronize(); dt = perf_counter() - t0

Then:

  step_time  = median(dt over steps 20–120)          # 20 warmup steps discarded
  epoch_time = step_time * num_batches # num batches is 10 million / effective batch size

A general point is that I guess we haven't really gone into the weeds of profiling how much we're getting out of the GPUs if we're maxing out the FLOPS.

[nihargupte-ph]

One final point is that it would be interesting to do a study to determine what is the optimal batch size for our problem. Of course we could use as many GPUs as possible but at a certain batch size, your convergence speed starts to drop off. This was explored in https://arxiv.org/abs/1812.06162. When we start launching large training runs, this would help us determine how many GPUs we should really be using.

[nihargupte-ph]

Regarding the above three points, I am fine to either

A). submit my review for the current PR and then we merge into main. Then we can consider these optimizations later.

B). Take into account either the point about removing the syncbatchnorm or using torch.compile and directly merging into this PR. I think the third point about the batch sizes will require more time and is outside of the scope of this PR.

Just let me know what your preference/thoughts on the above optimizations are @annalena-k @stephengreen . I think either could work, but I'm leaning more towards B to keep things organized.