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 finally separating computation from communication to understand what limits scaling.

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 448 × 448 pixels, surrounded by a 32-pixel halo on each side. Up to 32 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 two GPUs.

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 remaining distributed overhead.

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, but improvements become marginal beyond approximately 32 GPUs 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 decreases rapidly as the image is divided among more workers. Although runtime continues to improve up to around 32 GPUs, each additional GPU contributes progressively less 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, while the 8192² reconstruction continues to scale well throughout the experiment. 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 63% 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 62.7% to nearly 98% 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 slightly exceeds 100%. This 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.

Where does the remaining overhead come from?

The efficiency analysis shows that scaling eventually becomes less effective as the number of GPUs increases. This naturally raises the next question: what is actually limiting performance?

One possibility is that the computations themselves stop scaling. Another is that communication and synchronization progressively dominate the execution as less work remains on each GPU.

To distinguish between these effects, we profiled a second set of experiments using the PyTorch profiler, separating CUDA kernels performing useful computation from those associated with distributed communication.

The 2D experiment uses the same 4096² multi-frame super-resolution problem as the previous sections, while the 3D experiment reconstructs a simulated 512³ volume. Both experiments were executed on Tesla V100 GPUs using up to 64 GPUs.

The reported quantities are diagnostics extracted from the profiler rather than a decomposition of the total iteration time. Compute and communication can overlap, and synchronization outside the measured sections is reported independently.

The first observation is reassuring: the compute kernels themselves continue to scale almost ideally throughout the experiments.

For the 2D reconstruction, the compute portion accelerates by a factor of 54.4 between one and 64 GPUs, close to the ideal speedup of 64. The 3D experiment is even closer to ideal, reaching a 31× acceleration between two and 64 GPUs, compared with the theoretical value of 32.

This confirms that the computational building blocks of the distributed solver remain highly scalable. The decreasing end-to-end efficiency observed in the previous section is therefore not caused by the computation itself becoming inefficient.

The second observation is that communication gradually occupies a larger fraction of the execution.

This does not necessarily mean that communication becomes dramatically more expensive. Instead, useful computation is reduced much faster than communication. As the local workload assigned to each GPU decreases, the fixed cost of exchanging halo regions and synchronizing workers represents an increasing fraction of the remaining iteration time.

The effect is clearly visible for the 2D problem. Communication accounts for only a small portion of the CUDA activity at low GPU counts, but approaches half of the profiled execution at 64 GPUs as the compute kernels become extremely short.

The 3D experiment behaves differently. Although the absolute communication time is larger because considerably more data must be exchanged, each GPU also performs substantially more local computation. Communication therefore remains a relatively small fraction of the execution, allowing the reconstruction to maintain excellent scaling even at the largest GPU counts.

Together with the previous efficiency measurements, these results indicate that the current implementation is fundamentally compute-efficient. The main scaling limitations appear once communication, synchronization, and other distributed overheads become comparable to the shrinking amount of useful work performed by each GPU.

Future analyses will investigate these communication patterns in greater detail, including which collective operations dominate execution and how their cost depends on the reconstruction algorithm and hardware configuration.

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.