[논문 리뷰] A Fast, Performant, Secure Distributed Training Framework For Large Language Model

Wei Huang, Yinggui Wang외 4명|arXiv (Cornell University)|2024. 01. 18.

Topic Modeling인용 수 4

한 줄 요약

이 논문은 모델 분할, 신뢰 실행 환경(TEE), 경량 암호화를 활용하여 서버 측 및 클라이언트 측 모델/데이터 泄露를 방지하는 안전하고 고성능의 분산 학습 프레임워크를 제안한다. 클라이언트와 서버 양쪽에 TEE를 배포하고, 특히 후행 레이어를 서버 측 TEE에 배치함으로써 강력한 보안을 확보하면서도 정확도 손실을 최소화하며, 기준 방법 대비 정확도와 효율성에서 뛰어난 성능을 발휘한다.

ABSTRACT

The distributed (federated) LLM is an important method for co-training the domain-specific LLM using siloed data. However, maliciously stealing model parameters and data from the server or client side has become an urgent problem to be solved. In this paper, we propose a secure distributed LLM based on model slicing. In this case, we deploy the Trusted Execution Environment (TEE) on both the client and server side, and put the fine-tuned structure (LoRA or embedding of P-tuning v2) into the TEE. Then, secure communication is executed in the TEE and general environments through lightweight encryption. In order to further reduce the equipment cost as well as increase the model performance and accuracy, we propose a split fine-tuning scheme. In particular, we split the LLM by layers and place the latter layers in a server-side TEE (the client does not need a TEE). We then combine the proposed Sparsification Parameter Fine-tuning (SPF) with the LoRA part to improve the accuracy of the downstream task. Numerous experiments have shown that our method guarantees accuracy while maintaining security.

연구 동기 및 목표

• 악성 서버 및 클라이언트로부터 발생할 수 있는 분산(연합) LLM 학습에서의 모델 매개변수 및 데이터 泄露라는 핵심 보안 과제를 해결하기 위해.
• 높은 모델 정확도를 유지하면서도 민감한 데이터와 모델 가중치를 보호하는 안전하고 효율적이며 확장 가능한 학습 프레임워크를 설계하기 위해.
• 소비자용(소메모리) 및 산업용(대메모리) TEE 솔루션을 모두 지원함으로써 다양한 하드웨어 환경에서의 실용적 구현을 가능하게 하기 위해.
• TEE 내에서 업데이트되는 매개변수 수를 최소화하는 새로운 분할 미세조정 전략을 통해 학습 오버헤드와 장비 비용을 절감하기 위해.

제안 방법

• 프레임워크는 LLM을 분할하여 후행 레이어를 서버 측 TEE에, 전행 레이어를 클라이언트의 일반 환경에 배치함으로써 구현된다.
• LoRA 및 P-tuning v2 매개변수는 클라이언트와 서버 양측의 TEE 내부에 저장되고 계산되어 미세조정된 가중치와 임bedding의 기밀성을 확보한다.
• 기울기와 중간 출력 간의 안전한 통신은 기울기 및 중간 출력에 대해 한 번 사용한 패드(OTP) 암호화를 통해 달성된다.
• 분할 미세조정 전략을 도입하여, 서버 측 TEE에서 레이어의 일부(예: 마지막 4개 레이어)만을 미세조정함으로써 매개변수 업데이트 범위와 비용을 감소시킨다.
• 이 방법은 두 가지 TEE 구성 방식을 지원한다: Intel SGX(소메모리, GPU 없음)와 Intel TDX/SGX(대메모리, GPU 없음)이며, 각각에 맞는 맞춤형 기법을 제공한다.
• 희소화 매개변수 미세조정(SPF)을 LoRA와 결합하여 모델 크기를 증가시키지 않고도 최종 작업 정확도를 추가로 향상시킨다.

실험 결과

연구 질문

• RQ1 분산 LLM 학습에서 서버 측 및 클라이언트 측의 모델 매개변수 및 데이터 泄露를 어떻게 방지할 수 있는가?
• RQ2 TEE 기반 분산 LLM 프레임워크에서 보안, 모델 정확도, 학습 효율성 간 최적의 트레이드오프는 무엇인가?
• RQ3 후행 레이어를 서버 측 TEE에 격리하는 분할 미세조정 전략이 매개변수 처리 수를 줄이면서도 높은 정확도를 유지할 수 있는가?
• RQ4 경량 암호화(OTP)와 TEE의 조합이 차별적 개인정보 보호 기법(예: 차등적 개인정보 보장 또는 MPC)과 성능 및 보안 측면에서 어떻게 비교되는가?
• RQ5 다른 TEE 하드웨어 구성(예: SGX 대비 TDX/SGX)이 학습 지연 시간과 시스템 비용에 미치는 영향은 무엇인가?

주요 결과

• Method2는 QKV 및 밀도 레이어에 대해 각각 25% 및 50%의 매개변수 비율로 4개 레이어를 서버 측 TEE에 배치하여, 5개의 의료 데이터셋에서 평균 정확도가 가장 높게 나타났으며, FL-LLM 및 SWMT를 모두 초월했다.
• Method2의 평균 학습 시간은 1개 샘플당 4.33초(LoRA 기준)로, SWMT(17.92초) 및 Method1(17.06초)보다 현저히 낮아 높은 효율성을 입증했다.
• Method1은 암호화로 인한 FP16 절삭 오류로 인해 FL-LLM(평판)보다 略적으로 낮은 정확도를 기록했지만, 여전히 강력한 보안 보장을 제공했다.
• Method2의 학습 가능한 매개변수 수는 Method1보다 약 5배 많았지만, 훨씬 뛰어난 성능을 발휘하여 분할 미세조정의 효과를 입증했다.
• SWMT는 TEE 전용 CPU 처리로 인해 가장 높은 학습 및 추론 지연 시간을 보였고, 반면 Method2는 대부분의 계산을 GPU로 오프로드하여 TEE 사용을 최소화했다.
• 제거 실험 결과, TEE 내 레이어 수를 늘릴수록 정확도는 향상되지만 지연 시간과 매개변수 수 증가의 비용이 수반됨을 확인하여, 4개 레이어, (25%,50%) 비율 조합이 균형 잡힌 선택임을 검증했다.