Scaling Neural Network Verification with Tensor Parallelism and Fully Sharded Data Parallelism
2026-06-08 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors address the challenge of verifying neural networks using limited GPU memory. They adapt two methods from large model training, called Tensor Parallelism (TP) and Fully Sharded Data Parallelism (FSDP), to split the work across multiple GPUs. TP reduces memory usage but slightly lowers verification accuracy, while FSDP cuts memory use significantly without any loss in accuracy. Their techniques work well on various benchmarks and large models, but the main memory bottleneck remains the storage of certain key intermediate data (alpha tensors), suggesting where future improvements could focus.
formal neural network verificationbound propagationIBPCROWNalpha-CROWNTensor ParallelismFully Sharded Data ParallelismGPU memoryBranch-and-Boundverification benchmarks
Authors
Sergei Vorobyov, Eugene Ilyushin
Abstract
Formal neural network verification -- proving that a network satisfies safety properties for \emph{all} inputs in a specified domain -- is bounded in practice by GPU memory: standard implementations of bound-propagation algorithms (IBP, CROWN, $α$-CROWN) require weight and relaxation-coefficient matrices to reside entirely on one accelerator. We adapt two parallelism techniques originally developed for large-scale model training to the \texttt{auto\_LiRPA}\,/\,$α,β$-CROWN verification framework. \textbf{Tensor Parallelism (TP)} shards both weight and $A$-matrices across GPUs, achieving ${\approx}2\times$ peak-memory reduction at $P{=}2$; soundness is confirmed on VNN-COMP 2022 MNIST-FC benchmarks, though bound tightness degrades with the number of sharded zones due to forced IBP substitution for intermediate bounds inside sharded zones. \textbf{Fully Sharded Data Parallelism (FSDP)} shards only weight matrices with a per-layer \texttt{AllGather}, producing bounds that are \emph{bitwise identical} to the single-GPU baseline: baseline memory drops by 80--90\%, peak memory by 34--39\% on wide MLPs. FSDP integrates cleanly with complete verification ($β$-CROWN + Branch-and-Bound) and with convolutional layers (\texttt{BoundConv}); a complete \emph{unsat} result is obtained for CIFAR-100 ResNet-large (VNN-COMP 2024) under FSDP. Across all experiments the memory bottleneck in $α$-CROWN+BaB mode proves to be per-neuron alpha tensors, not weight matrices, pointing to the key direction for future work.