Scaling distributed plug-and-play inference

Distributing an inverse-problem solver across many GPUs makes it possible to reconstruct images that would not fit on a single device. But using more GPUs does not automatically make the reconstruction proportionally faster. The computation must be divided between workers, overlapping image regions may be processed several times, and intermediate results must be exchanged and synchronized.

This study investigates how our distributed plug-and-play inference pipeline behaves as both the problem size and the number of GPUs increase.

Rather than only asking whether the implementation can run on many GPUs, we want to answer three practical questions:

1. Can distributed inference preserve reconstruction quality while enabling larger reconstructions and reducing computation time?
2. How efficiently does the reconstruction accelerate when more GPUs are added?
3. When scaling eventually slows down, is the limitation caused by the computation itself or by communication and synchronization?

We first verify that distributed execution preserves the reconstruction trajectory. We then study how iteration time decreases as additional GPUs are added, before analyzing parallel efficiency and separating out the additional work introduced by overlapping tiles. Computation and communication are separated on the communication page.

Experimental setting

Each plug-and-play iteration alternates between two main operations:

• a gradient step using the acquisition model;
• a DRUNet denoising step acting as the learned image prior.

Both operations expose parallelism, but in different ways. The acquisition operators are distributed across workers, while the image passed to the denoiser is split into overlapping tiles. The overlap provides the spatial context needed by the neural network near tile boundaries, but it also introduces additional computation and data movement.

The strong-scaling experiment uses noisy multi-frame super-resolution problems with eight acquisition operators. We reconstruct RGB images with side lengths of 2048, 4096, and 8192 pixels on Tesla V100-SXM2 GPUs with 32 GB of memory, using four GPUs per node.

Distributed denoising uses an inner tile size of 512 × 512 pixels, surrounded by a 32-pixel halo on each side. Up to four tiles are processed in one batch. The full experiment configuration is available in benchmark_inference/configs/experiments/strong_scaling_inference.yml.

Timing values are averaged over iterations 3 to 10, after excluding the first two iterations to avoid initialization and warm-up effects. The quality comparison includes initialization (iteration 0) and every PnP iteration from 1 to 10. The 2048² and 4096² experiments include a non-distributed single-GPU reference. The 8192² image does not fit in the corresponding single-GPU configuration, so its scaling study starts from four GPUs.

The measured iteration time also covers the benchmark’s own per-iteration PSNR, SSIM, and MSE evaluation on the full reconstruction. This is measurement rather than reconstruction, so it is reported as a separate band and excluded from the efficiency analysis.

Can distributed inference preserve quality and provide useful speedups?

Before evaluating performance, we first verify that distributing the reconstruction does not alter the optimization itself. Plug-and-play reconstruction alternates physics updates and learned denoising operations, so numerical differences introduced by distributed execution could, in principle, accumulate over successive iterations and lead to a different reconstruction.

To assess this, we compare the PSNR obtained at every iteration with the corresponding single-GPU reference reconstruction.

The reconstruction trajectories remain virtually identical across all tested GPU counts. The largest observed deviation is below 0.002 dB, several orders of magnitude smaller than the approximately 18.6 dB improvement achieved during the reconstruction itself. More importantly, these differences remain bounded throughout the optimization and show no tendency to accumulate over successive iterations.

These results indicate that the distributed implementation preserves the behavior of the original algorithm to numerical precision. Although small floating-point differences naturally arise from tiled execution and distributed reductions, they do not translate into any measurable degradation of the reconstruction quality.

Having established that distributed execution preserves the reconstruction trajectory, we can now evaluate its computational performance.

The figure below shows the average time spent in each stage of one steady-state plug-and-play iteration, separating the acquisition physics, learned denoising, and the per-iteration metric evaluation.

All three problem sizes benefit from additional GPUs, although the amount of acceleration naturally depends on the size of the reconstruction. The largest 8192² problem, which cannot be reconstructed in the corresponding single-GPU configuration, reaches sub-second iteration times on 64 GPUs while continuing to improve throughout the explored scaling range.

Smaller problems exhibit the expected saturation behavior. The 2048² reconstruction accelerates significantly as GPUs are added, and its scaling study stops at four GPUs, where each worker already holds only about one megapixel of useful image. The 4096² problem shows the same effect from eight GPUs onwards, because each worker receives too little computation to fully amortize distributed overhead. Larger images, by contrast, maintain enough local work to keep additional GPUs productive.

Taken together, these two figures validate the distributed inference strategy. The implementation preserves the numerical behavior of the original plug-and-play algorithm while extending it to much larger images and substantially reducing reconstruction time as computational resources increase.

How efficiently are additional GPUs used?

The previous section showed that distributed execution reduces reconstruction time while preserving the reconstruction trajectory. A natural follow-up question is whether these speedups remain proportional to the amount of hardware used.

In an ideal strong-scaling experiment, doubling the number of GPUs would halve the runtime, and quadrupling the resources would make the computation four times faster. In practice, distributed execution always introduces some overhead, so the observed speedup gradually deviates from this ideal behavior.

To quantify this effect, we measure the parallel efficiency

where $T_N$ is the measured iteration time on $N$ GPUs and the baseline configuration uses $N_0$ GPUs. An efficiency of 100% corresponds to ideal strong scaling, while lower values indicate that additional GPUs provide progressively smaller reductions in runtime.

The expected trend immediately appears: larger images remain efficient over a much wider range of GPU counts than smaller ones.

For the 2048² reconstruction, efficiency drops to 78.2% as soon as the image is divided between two workers and stays there at four, because the local amount of work becomes too small to fully amortize communication and synchronization costs.

The 4096² reconstruction remains efficient over a larger range, holding 76.6% at four GPUs before falling to 64.4% at eight and 60.1% at sixteen, while the 8192² reconstruction stays at 95.0% on eight GPUs and 79.8% on sixteen before decaying to 45.7% on sixty-four. These results reinforce the main observation from the timing measurements: the usefulness of additional GPUs is primarily determined by the amount of work assigned to each worker.

At first sight, however, the initial efficiency values may appear surprisingly low. For example, the 4096² reconstruction reaches only about 79% efficiency when moving from one GPU to two GPUs. This apparent loss is largely explained by an important difference between the reference and distributed executions.

Accounting for the work introduced by overlapping tiles

The distributed denoiser operates on overlapping image tiles. Each tile is extended with a halo before being processed so that the neural network has sufficient spatial context near tile boundaries.

As a result, the distributed implementation evaluates more pixels than are present in the original image. Neighboring tiles process overlapping regions multiple times before their central regions are cropped and stitched back together.

This additional computation is an intrinsic consequence of the tiling strategy, not of inefficient parallel execution. The first transition from a single-GPU reconstruction to distributed inference therefore combines two effects:

1. the efficiency of distributing the computation across multiple GPUs;
2. the additional work introduced by overlapping tile halos.

To separate these effects, we also report a work-normalized efficiency

where $W_N$ denotes the amount of halo-padded image data processed by the denoiser relative to the useful image area.

This normalization does not assume that inference time is perfectly proportional to the number of pixels. Rather, it removes the known increase in work introduced purely by the tiling geometry.

The normalized curves reveal that most of the apparent efficiency loss at low GPU counts is simply the cost of processing overlapping halos. For the 4096² problem, the two-GPU efficiency increases from 79.3% to 100.4% after normalization, showing that the distributed implementation itself remains close to ideal.

The same effect is observed for the 2048² reconstruction, where the normalized efficiency reaches 99.0%. The 4096² value slightly exceeding 100% should not be interpreted as superlinear scaling of the overall algorithm. Rather, it indicates that a simple pixel-count correction cannot perfectly model every implementation detail, such as differences in neural-network batching, GPU occupancy, or kernel selection.

As the GPU count continues to increase, both the absolute and normalized curves decrease together. At this stage, halo computation is no longer the dominant factor. Instead, each worker receives progressively less useful computation, making communication, synchronization, and load imbalance increasingly important.

These experiments deliberately keep the tiling parameters fixed so that GPU scaling can be studied in isolation. In the following benchmark studies, we investigate how tile size and workload distribution can themselves be optimized to extend the range over which distributed inference scales efficiently.

The remaining cost — what is real data transfer and what is ranks waiting on each other — is profiled on the communication page, which uses the same solver.

Conclusions

This study demonstrates that distributed plug-and-play inference can scale to much larger reconstruction problems without sacrificing either numerical correctness or computational efficiency.

First, distributing the computation preserves the reconstruction trajectory. Where a direct comparison is possible, the distributed implementation follows the same optimization path as the single-GPU reference, with only negligible floating-point differences introduced by tiled execution and distributed reductions.

Second, the implementation delivers substantial speedups while extending the range of solvable problems. The largest 8192² reconstruction, which cannot be executed in the reference single-GPU configuration, reaches sub-second iterations on 64 GPUs, demonstrating that distributed execution is both a scalability and an enabling technology.

Finally, the experiments reveal what ultimately governs distributed performance. Scaling is not determined solely by the number of GPUs, but by the amount of useful computation assigned to each worker. When the local workload remains sufficiently large, the distributed implementation achieves near-ideal scaling. As the workload per GPU decreases, communication, synchronization, and other fixed overheads naturally become an increasing fraction of the execution time.

Taken together, these results validate the overall distributed inference strategy implemented in the benchmark. They also motivate the next stage of the analysis: understanding how tiling strategies, workload balancing, and communication patterns can be optimized to extend this efficient scaling regime to even larger systems.