AEGIS: Scaling Long-Sequence Homomorphic Encrypted Transformer Inference via Hybrid Parallelism on Multi-GPU Systems

1. 문제 정의 (Problem Statement)

완전 동형 암호화 (FHE, Fully Homomorphic Encryption) 는 클라우드 환경에서 데이터 프라이버시를 보호하며 머신러닝 추론을 가능하게 하지만, 긴 시퀀스 (Long-Sequence) Transformer 모델을 다중 GPU 에서 실행할 때 심각한 확장성 문제를 야기합니다.

• 메모리 폭발: 암호화된 가중치와 시퀀스 길이에 비례하여 급격히 증가하는 암호화된 활성화 (Activations) 데이터는 단일 GPU 의 메모리 용량을 쉽게 초과합니다. (예: 2048 토큰 입력 시 BERT-Base 모델만으로도 700GB 이상의 인코딩 가중치와 110GB 이상의 암호화 활성화가 필요함)
• 기존 방식의 한계:
  • 플레인텍스트 (Plaintext) 병렬화 적용 불가: 기존 데이터 병렬 (DP) 또는 텐서 병렬 (TP) 방식은 암호화 환경에서 적용하기 어렵습니다. 암호화된 데이터는 고정된 슬롯 (Slot) 레이아웃을 가지며, 임의의 분할이 불가능하고, 통신 시 전체 암호문 (Ciphertext) 을 전송해야 하므로 통신 오버헤드가 기하급수적으로 증가합니다.
  • RNS 기반 병렬화 (Limb Parallelism) 의 비효율: 동형 암호 (CKKS) 의 잔여수 체계 (RNS) 를 이용한 '사지 (Limb) 병렬화'는 키 스위칭 (Key Switching) 및 회전 (Rotation) 연산 시 장치 간 통신이 빈번하게 발생하여, 오히려 단일 GPU 실행보다 속도가 느려지는 역효과를 낳습니다.
  • 통신 증폭: 암호화 환경에서는 소량의 데이터 전송도 수 GB 단위의 암호문 전송으로 이어져 (예: 1024 토큰 전송 시 평문 2MB vs 암호문 1.3GB) 통신 지연이 병목이 됩니다.

2. 제안 방법론: AEGIS

저자들은 AEGIS (Application-Encryption Guided Inference System) 를 제안합니다. 이는 애플리케이션 수준의 데이터 흐름과 암호화 수준의 RNS 결합 (Coupling) 을 동시에 고려하여 다중 GPU 환경에서 긴 시퀀스 암호화 Transformer 추론을 확장 가능하게 만드는 컴파일러 - 런타임 협업 프레임워크입니다.

핵심 기술 요소

1. 의존성 인식 작업 분할 (Dependency-Aware Workload Partitioning):

   • 모듈러스 일관성 배치 (Modulus-coherent placement): CKKS 의 키 스위칭 및 회전 연산에서 발생하는 RNS 사지 (Limb) 간의 의존성을 분석하여, 동일한 의존성 그룹에 속하는 다항식들을 동일한 GPU 에 배치합니다. 이를 통해 장치 간 동기화를 최소화하고 키 스위칭 연산을 로컬에서 처리합니다.
   • 토큰 일관성 배치 (Token-coherent placement): 애플리케이션 수준에서 동일한 토큰에 해당하는 암호문 조각들이 서로 다른 GPU 에 흩어지지 않도록 배치합니다. 이는 토큰 단위 집계 (Aggregation) 시 불필요한 통신을 방지합니다.
   • 배치 계층 구조: RNS 모듈러스 의존성을 최우선으로 하고, 그 다음으로 토큰 일관성을 고려하는 엄격한 배치 계층을 따릅니다.

2. FFN 및 Self-Attention 작업 분할 전략:

   • FFN (Feed-Forward Networks): 동일한 출력 그룹에 속하는 가중치 조각들을 동일한 GPU 에 배치하여, 암호화된 활성화의 AllGather 통신을 최소화하고 로컬 연산을 최대화합니다.
   • Self-Attention: QKV 프로젝션 후, 어텐션 점수 계산 시 부분적인 AllReduce-like 연산을 수행하고, Softmax 전후로 ReduceScatter 및 AllGather 를 효율적으로 삽입하여 통신을 줄입니다.

3. 연산자 재순서화 (Polynomial-Operator Reordering) 및 통신 숨김:

   • 계산 그래프의 저수준 연산 순서를 재구성하여, 통신 (Collective Communication) 과 계산을 겹치게 (Overlap) 합니다.
   • 예를 들어, 회전 (Rotation) 오프셋별 계산을 교차하여 수행함으로써, 일부 결과가 준비되는 대로 통신을 시작하게 하여 GPU 의 대기 시간 (Stall) 을 제거합니다. 이는 런타임 휴리스틱 없이도 결정론적으로 통신 지연을 숨깁니다.

4. 컴파일러 파이프라인:

   • 고수준 애플리케이션 IR 에서 저수준 다항식 IR 로 변환하는 과정에서 토큰/시퀀스 의존성과 RNS 의존성을 분석하고, 최적의 통신 삽입 시점과 연산자 스케줄링을 결정합니다.

3. 주요 기여 (Key Contributions)

• AEGIS 프레임워크 개발: 애플리케이션 데이터 흐름과 암호화 원시 연산을 조율하여 다중 GPU 에서 확장 가능한 암호화 Transformer 추론을 가능하게 하는 최초의 시스템.
• 하이브리드 병렬화 전략: 애플리케이션 의미론 (Token/Embedding) 과 암호화 제약 (RNS/Modulus) 을 결합한 새로운 배치 전략 제시.
• 통신 숨김 메커니즘: 런타임 오버헤드 없이 통신과 계산을 겹치게 하는 연산자 재순서화 기술 개발.
• 성능 및 메모리 최적화: 기존 방식 대비 통신량 감소, 메모리 효율성 향상, 확장성 확보.

4. 실험 결과 (Results)

실험은 2048 토큰 입력을 가진 BERT-Base 모델을 대상으로 2 개 및 4 개의 GPU (NVIDIA RTX A6000 및 A100) 환경에서 수행되었습니다.

• 통신 감소:
  • FFN 레이어에서 57.9% 감소.
  • Self-Attention 레이어에서 81.3% 감소 (기존 최첨단 설계 대비).
• 확장성 (Scaling Efficiency):
  • 4 개 GPU 환경에서 96.62% 의 확장 효율 달성 (단일 GPU 대비).
  • 2 개 GPU 환경에서는 92.98% 의 효율 달성.
• 속도 향상 (Speedup):
  • 4 개 GPU 환경에서 3.86 배의 엔드투엔드 속도 향상.
  • 2 개 GPU 환경에서 1.86 배 속도 향상.
• 메모리 효율성:
  • 중간 암호문을 복제하는 방식 (Tensor Parallel 등) 대비 69.1% 의 단일 장치당 메모리 감소.
  • 긴 시퀀스 (2048 토큰) 에서도 단일 GPU 메모리 부족 (OOM) 없이 실행 가능.
• 비교 대상 (Baselines) 대비 우위:
  • HEBooster (Limb Parallel): 통신 오버헤드로 인해 확장성이 낮음 (2 GPU 에서 오히려 느려짐).
  • Cinnamon/Hydra: 하드웨어 의존적이거나 메모리 복제로 인해 긴 시퀀스에서 OOM 발생 또는 효율 저하.
  • Plaintext Tensor Parallel: 암호화 환경에 직접 적용 시 메모리 낭비와 불필요한 계산으로 인해 효율이 극히 낮음.

5. 의의 및 결론 (Significance)

이 논문은 암호화된 Transformer 추론이 단일 장치를 넘어 대규모 다중 GPU 시스템으로 확장되기 위해서는 애플리케이션 레벨의 데이터 구조와 암호화 레벨의 연산 제약을 분리하여 고려해서는 안 되며, 이를 통합적으로 조율 (Coordinated) 해야 함을 증명했습니다.

• 실용적 기여: 긴 컨텍스트 (Long-Context) 를 가진 암호화 LLM 서비스의 실현 가능성을 높였으며, 메모리 제약과 통신 병목 문제를 동시에 해결하는 구체적인 아키텍처를 제시했습니다.
• 미래 방향: 향후 동형 암호 시스템 설계 시 패킹 (Packing), 배치 (Placement), 통신 스케줄링을 교차 계층 (Cross-layer) 최적화 문제로 접근해야 함을 시사합니다.

요약하자면, AEGIS는 암호화 연산의 고유한 특성 (RNS, 키 스위칭) 과 Transformer 의 구조적 특성을 결합하여, 긴 시퀀스 데이터 처리 시 발생하는 메모리 및 통신 병목을 획기적으로 개선한 혁신적인 시스템입니다.