Technical Summary: Scaling Neural Network Verification with Tensor Parallelism and Fully Sharded Data Parallelism

1. Problem Statement

Formal neural network verification aims to mathematically prove that a network satisfies safety properties for all inputs within a specified domain. While bound-propagation algorithms like IBP, CROWN, and $\alpha$-CROWN are the state-of-the-art for this task, their practical application is severely constrained by GPU memory. Standard implementations require weight matrices and linear-relaxation $A$-matrices to reside entirely on a single accelerator. For large networks (e.g., a single hidden layer with width $h=262,144$), the weight matrices alone can exceed 4 GB in float32, leaving insufficient memory for the $A$-matrices, batching, or intermediate tensors required for verification. This paper addresses the need to remove this single-GPU memory bottleneck by adapting parameter-sharding strategies from large-scale model training to the verification framework.

2. Methodology

The authors adapt two parallelism techniques—Tensor Parallelism (TP) and Fully Sharded Data Parallelism (FSDP)—to the auto LiRPA / $\alpha, \beta$-CROWN verification framework.

2.1 Tensor Parallelism (TP)

TP splits matrix multiplications within a single layer across $P$ GPUs.

• Implementation: The authors implemented custom BoundLinearTP Col and BoundLinearTP Row operators. Weights are sharded along the output (column) or input (row) dimensions.
• Mechanism: In a Column–Row pair, the computation is distributed, requiring one AllReduce operation to recover the result.
• Limitation: When ReLU activations occur inside a sharded zone (between Column and Row layers), the CROWN backward pass requires intermediate bounds to determine relaxation slopes ($\lambda_i$). Computing these exactly would require inter-GPU communication at every step. To preserve locality, the authors substitute the simpler Interval Bound Propagation (IBP) for these intermediate bounds. This introduces a "wrapping effect," causing bound tightness to degrade as the depth of the sharded zone increases.

2.2 Fully Sharded Data Parallelism (FSDP)

FSDP shards parameter tensors equally across $P$ GPUs but reconstructs them for computation.

• Implementation: Implemented with approximately 100 lines of Python, FSDP replaces BoundParams tensors with shards. Before a layer's computation, an AllGather assembles the full weight matrix; immediately after, the full tensor is freed, retaining only the shard.
• Mechanism: This approach ensures that at most one full weight matrix resides in GPU memory at any instant. Crucially, FSDP does not alter the order or type of arithmetic operations; it merely changes the storage location of weights.
• Correctness: Because the arithmetic remains identical to the single-GPU baseline, FSDP produces bitwise-identical bounds ($\Delta = 0$).
• Integration: FSDP integrates seamlessly with complete verification ($\beta$-CROWN + Branch-and-Bound) and convolutional layers (BoundConv).

2.3 Experimental Setup

Experiments were conducted on NVIDIA A40 GPUs using PyTorch 2.x and the auto LiRPA library. The authors introduced a bab.force_synchronous flag to ensure fair comparisons between single-GPU and distributed modes by disabling automatic batch resizing and pruning in the single-GPU path, ensuring identical workloads.

3. Key Contributions

1. TP Implementation for Verification: The authors implemented TP for auto LiRPA using custom operators, achieving approximately $2\times$ peak-memory reduction at $P=2$. They formally analyzed the trade-off, confirming that while soundness is preserved, bound tightness degrades with the number of sharded zones due to IBP substitution.
2. FSDP Implementation: They implemented FSDP for bound propagation, achieving 80–90% baseline memory reduction and 34–39% peak memory reduction on wide MLPs. Crucially, this method produces bitwise-identical bounds to the single-GPU baseline.
3. Integration with Complete Verification: FSDP was successfully integrated with $\beta$-CROWN and Branch-and-Bound (BaB), enabling the first complete verification (UNSAT result) of a CIFAR-100 ResNet-large model under FSDP.
4. Bottleneck Identification: Through extensive experimentation, the authors identified that in $\alpha$-CROWN + BaB mode, the primary memory bottleneck is not the weight matrices, but the per-neuron alpha tensors (ReLU slopes).

4. Results

4.1 Memory Reduction

• TP: Achieved a peak memory reduction factor of $\approx 1.985$ (close to theoretical $P=2$) for a wide MLP ($h=262,144$).
• FSDP:
  • Baseline Memory: Reduced by 80–90% across various model sizes, as only shards are allocated before bound computation.
  • Peak Memory: Reduced by 34–39% on wide MLPs ($h \ge 4096$). The savings are lower than baseline savings because the unsharded $A$-matrices (which scale with width squared) remain in memory.

4.2 Bound Accuracy and Soundness

• TP:
  • Models with one sharded zone (no intermediate ReLU) achieved machine-precision agreement ($\sim 10^{-8}$) with the single-GPU reference.
  • Models with multiple sharded zones showed significant deviation (e.g., $\sim 750$ units for a 3-zone model) due to IBP substitution, though soundness (one-sided over-approximation) was preserved.
• FSDP:
  • Achieved exactly zero deviation ($\Delta lb = \Delta ub = 0.00$) across all tested models (MLPs and ONNX benchmarks) for both IBP and CROWN.
  • Successfully verified a Transformer model (ViT) and convolutional ResNets with bitwise-identical results to the baseline.

4.3 Complete Verification (BaB)

In $\beta$-CROWN + BaB mode, the memory consumption is dominated by alpha tensors (e.g., 1.5 GB for alpha tensors vs. 1.5 MB for weights in a specific MNIST network).

• Sharding weights via FSDP provided negligible peak memory savings in this mode because the alpha tensors were not sharded.
• However, FSDP successfully enabled the verification of a ResNet-large model (VNN-COMP 2024) that would otherwise be infeasible or difficult to manage, confirming an UNSAT property without BaB branching.

5. Significance and Limitations

The paper demonstrates that parameter-sharding strategies from training can be transferred to formal verification with modest engineering effort.

• Significance: FSDP offers a path to verify larger networks with exact bounds, while TP offers the highest memory reduction at the cost of bound tightness for deep networks. The work establishes that FSDP is compatible with the complex requirements of complete verification (BaB) and diverse architectures (CNNs, Transformers).
• Limitations:
  • Scalability: Experiments were limited to $P=2$. The authors note that as $P$ increases, the unsharded $A$-matrices will likely become the dominant cost, leading to diminishing returns.
  • TP Tightness: The degradation of bound tightness in TP for deep networks limits its utility for precise verification of deep architectures.
  • Bottleneck Shift: The experiments reveal that future multi-GPU verification research must focus on sharding per-neuron alpha tensors, not just weights, to achieve significant scaling in complete verification modes.

The authors conclude that while FSDP solves the weight-storage bottleneck, the alpha tensor bottleneck remains the critical challenge for scaling complete verification to larger models.