기술 요약: 텐서 병렬성 및 완전 샤딩 데이터 병렬성을 이용한 신경망 검증 스케일링

1. 문제 정의

형식적 신경망 검증(Formal neural network verification)은 네트워크가 지정된 입력 도메인 내의 모든 입력에 대해 안전 속성을 만족함을 수학적으로 증명하는 것을 목표로 한다. IBP, CROWN, $\alpha$-CROWN와 같은 경계 전파(bound-propagation) 알고리즘이 이 작업의 최첨단 기술이지만, 이들의 실제 적용은 GPU 메모리에 의해 심하게 제약된다. 표준 구현에서는 가중치 행렬과 선형 완화(linear-relaxation) $A$-행렬이 전체가 단일 가속기에 상주해야 한다. 대규모 네트워크(예: 너비 $h=262,144$인 단일 은닉층)의 경우, 가중치 행렬만으로도 float32 기준 4 GB를 초과할 수 있어, 검증에 필요한 $A$-행렬, 배칭 또는 중간 텐서를 위한 메모리가 부족해진다. 본 논문은 대규모 모델 학습을 위한 파라미터 샤딩 전략을 검증 프레임워크에 적응시켜 이러한 단일 GPU 메모리 병목 현상을 제거하는 문제를 다룬다.

2. 방법론

저자들은 두 가지 병렬화 기술인 텐서 병렬성(Tensor Parallelism, TP)과 완전 샤딩 데이터 병렬성(Fully Sharded Data Parallelism, FSDP)을 auto LiRPA / $\alpha, \beta$-CROWN 검증 프레임워크에 적응시켰다.

2.1 텐서 병렬성 (TP)

TP는 단일 레이어 내의 행렬 곱셈을 $P$개의 GPU로 분할한다.

• 구현: 저자들은 커스텀 BoundLinearTP Col 및 BoundLinearTP Row 연산자를 구현하였다. 가중치는 출력(열) 또는 입력(행) 차원을 따라 샤딩된다.
• 메커니즘: Column–Row 쌍에서 계산은 분산되어 수행되며, 결과를 복구하기 위해 한 번의 AllReduce 연산이 필요하다.
• 제한 사항: 샤딩된 구역(Column과 Row 레이어 사이) 내에 ReLU 활성화 함수가 발생하는 경우, CROWN 역전파는 완화 기울기($\lambda_i$)를 결정하기 위해 중간 경계값(intermediate bounds)을 필요로 한다. 이를 정확하게 계산하려면 매 단계마다 GPU 간 통신이 필요하다. 국부성(locality)을 유지하기 위해, 저자들은 이러한 중간 경계값에 대해 더 단순한 구간 경계 전파(Interval Bound Propagation, IBP)를 대체하여 사용한다. 이는 "래핑 효과(wrapping effect)"를 유발하여, 샤딩된 구역의 깊이가 깊어질수록 경계의 타이트함(tightness)이 저하된다.

2.2 완전 샤딩 데이터 병렬성 (FSDP)

FSDP는 파라미터 텐서를 $P$개의 GPU에 균등하게 샤딩하지만, 계산을 위해 이를 재구성한다.

• 구현: 약 100줄의 파이썬 코드로 구현된 FSDP는 BoundParams 텐서를 샤드(shard)로 대체한다. 레이어 계산 전, AllGather를 통해 전체 가중치 행렬을 조립하며, 직후에 전체 텐서를 해제하고 샤드만을 유지한다.
• 메커니즘: 이 접근 방식은 계산 시점에 최대 하나의 전체 가중치 행렬만이 GPU 메모리에 존재하도록 보장한다. 결정적으로, FSDP는 산술 연산의 순서나 유형을 변경하지 않으며, 단지 가중치의 저장 위치만을 변경한다.
• 정확성: 산술 연산이 단일 GPU 베이스라인과 동일하기 때문에, FSDP는 비트 단위로 동일한 경계($\Delta = 0$)를 생성한다.
• 통합: FSDP는 완전한 검증($\beta$-CROWN + Branch-and-Bound) 및 합성곱 레이어(BoundConv)와 원활하게 통합된다.

2.3 실험 설정

실험은 NVIDIA A40 GPU에서 PyTorch 2.x 및 auto LiRPA 라이브러리를 사용하여 수행되었다. 저자들은 단일 GPU 경로에서의 자동 배치 크기 조정 및 프루닝을 비활성화하여 단일 GPU와 분산 모드 간의 공정한 비교를 보장하기 위해 bab.force_synchronous 플래그를 도입함으로써 동일한 워크로드를 보장하였다.

3. 주요 기여

1. 검증을 위한 TP 구현: 저자들은 커스텀 연산자를 사용하여 auto LiRPA를 위한 TP를 구현하였으며, $P=2$일 때 약 $2\times$의 피크 메모리 감소를 달성하였다. 저자들은 트레이드오프를 공식적으로 분석하여, 건전성(soundness)은 유지되지만 IBP 대체로 인해 샤딩된 구역의 수가 늘어남에 따라 경계의 타이트함이 저하됨을 확인하였다.
2. FSDP 구현: 저자들은 경계 전파를 위한 FSDP를 구현하여, 다양한 모델 크기에서 80–90%의 베이스라인 메모리 감소와 넓은 MLP에서 34–39%의 피크 메모리 감소를 달성하였다. 결정적으로, 이 방법은 단일 GPU 베이스라인과 비트 단위로 동일한 경계를 생성한다.
3. 완전 검증과의 통합: FSDP는 $\beta$-CROWN 및 Branch-and-Bound(BaB)와 성공적으로 통합되어, FSDP 환경에서 CIFAR-100 ResNet-large 모델의 첫 번째 완전한 검증(UNSAT 결과)을 가능하게 했다.
4. 병목 현상 식별: 광범위한 실험을 통해 저자들은 $\alpha$-CROWN + BaB 모드에서 주요 메모리 병목 현상이 가중치 행렬이 아니라 뉴런당 알파 텐서(ReLU 기울기)임을 식별하였다.

4. 결과

4.1 메모리 감소

• TP: 넓은 MLP ($h=262,144$)에 대해 이론적 수치인 $P=2$에 근접한 약 $1.985$의 피크 메모리 감소율을 달了.
• FSDP:
  • 베이스라인 메모리: 경계 계산 전에 샤드만 할당되므로 다양한 모델 크기에 걸쳐 80–90% 감소하였다.
  • 피크 메모리: 넓은 MLP ($h \ge 4096$)에서 34–39% 감소하였다. 샤딩되지 않은 $A$-행렬(너비의 제곱에 비례)이 메모리에 남아 있기 때문에 절감 폭은 베이스라인 절감 폭보다 낮다.

4.2 경계 정확도 및 건전성

• TP:
  • 하나의 샤딩 구역(중간 ReLU 없음)을 가진 모델은 단일 GPU 참조 모델과 기계 정밀도 수준의 일치($\sim 10^{-8}$)를 보였다.
  • 여러 샤딩 구역을 가진 모델은 IBP 대체로 인해 상당한 편차(예: 3개 구역 모델에서 $\sim 750$ 유닛)를 보였으나, 건전성(한 방향 과대 근사)은 유지되었다.
• FSDP:
  • 테스트된 모든 모델(MLP 및 ONNX 벤치마크)에 대해 IBP 및 CROWN 모두에서 정확히 제로 편차($\Delta lb = \Delta ub = 0.00$)를 달성하였다.
  • Transformer 모델(ViT) 및 합성곱 ResNet을 베이스라인과 비트 단위로 동일한 결과로 성공적으로 검증하였다.

4.3 완전 검정 (BaB)

$\beta$-CROWN + BaB 모드에서 메모리 소비는 알파 텐서에 의해 지배된다 (예: 특정 MNIST 네트워크에서 알파 텐서 1.5 GB vs 가중치 1.5 MB).

• FSDP를 통한 가중치 샤딩은 알파 텐서가 샤딩되지 않았기 때문에 이 모드에서 피크 메모리 절감 효과가 미미했다.
• 그러나 FSDP는 BaB 브랜칭 없이도 UNSAT 속성을 확인하는 것을 통해, 원래는 불가능하거나 관리하기 어려웠던 ResNet-large 모델(VNN-COMP 2024)의 검증을 성공적으로 수행하였다.

5. 의의 및 한계

본 논문은 학습을 위한 파라미터 샤딩 전략이 적은 엔지니어링 노력으로 형식적 검증에 전이될 수 있음을 보여준다.

• 의의: FSDP는 정확한 경계를 가진 더 큰 네트워크를 검증할 수 있는 경로를 제공하며, TP는 정밀한 검증을 위한 비용으로 경계의 타이트함 저하를 감수하면서 가장 높은 메모리 감소를 제공한다. 이 연구는 FSDP가 복잡한 요구 사항(BaB) 및 다양한 아키텍처(CNN, Transformer)와 호환됨을 입증하였다.
• 한계:
  • 확장성: 실험은 $P=2$로 제한되었다. 저자들은 $P$가 증가함에 따라 샤딩되지 않은 $A$-행렬이 지배적인 비용이 되어 수익 체감(diminishing returns)이 발생할 것이라고 언급하였다.
  • TP 타이트함: 깊은 네트워크에서 TP의 경계 타이트함 저하는 정밀한 검증을 위한 활용도를 제한한다.
  • 병목 현상의 이동: 실험 결과, 향에 대한 완전한 검증을 위해 스케일링하려면 가중치뿐만 아니라 뉴런당 알파 텐서를 샤딩하는 것에 미래의 다중 GPU 검증 연구가 집중되어야 함을 보여준다.

저자들은 FSDP가 가중치 저장 병목 현상은 해결하지만, 알파 텐서 병목 현상이 더 큰 모델에 대한 완전한 검증을 위한 핵심적인 과제로 남아 있다고 결론지었다.