Scaling Neural Network Verification with Tensor Parallelism and Fully Sharded Data Parallelism

arXiv:2606.09377v1 Announce Type: cross Abstract: Formal neural network verification -- proving that a network satisfies safety properties for all inputs in a specified domain -- is bounded in practice by GPU memory: standard implementations of bound-propagation algorithms (IBP, CROWN, alpha-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 exttt{auto_LiRPA},/,alpha,eta-CROWN verification framework. extbf{Tensor Parallelism (TP)} shards both weight and A-matrices across GPUs, achieving {approx}2imes 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. extbf{Fully Sharded Data Parallelism (FSDP)} shards only weight matrices with a per-layer exttt{AllGather}, producing bounds that are 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 (eta-CROWN + Branch-and-Bound) and with convolutional layers (exttt{BoundConv}); a complete unsat result is obtained for CIFAR-100 ResNet-large (VNN-COMP 2024) under FSDP. Across all experiments the memory bottleneck in alpha-CROWN+BaB mode proves to be per-neuron alpha tensors, not weight matrices, pointing to the key direction for future work.

Source: arXiv cs.AI | 2026-06-09