[논문 리뷰] The Future of Large Language Model Pre-training is Federated

Lorenzo Sani, Alex Iacob외 9명|arXiv (Cornell University)|2024. 05. 17.

Topic Modeling인용 수 4

한 줄 요약

이 논문은 데이터가 풍부한 기관들이 중앙 집중식 데이터 센터 없이도 분산된 방식으로 빌리언 파라미터 규모의 대규모 언어 모델(Large Language Models, LLMs)을 사전 훈련할 수 있도록 하는 분산형, 협업 기반의 연합 학습(Federated Learning, FL) 프레임워크를 제안한다. 이는 이질적인 FL 환경에서 빌리언 스케일 LLM을 성공적으로 훈련한 최초의 사례로, 중앙 집중식 훈련과 경쟁 가능한 성능을 달성하면서도 계산 자원과 데이터 자원에 대한 접근을 민주화한다.

ABSTRACT

Generative pre-trained large language models (LLMs) have demonstrated impressive performance over a wide range of tasks, thanks to the unprecedented amount of data they have been trained on. As established scaling laws indicate, LLMs' future performance improvement depends on the amount of computing and data sources they can leverage for pre-training. Federated learning (FL) has the potential to unleash the majority of the planet's data and computational resources, which are underutilized by the data-center-focused training methodology of current LLM practice. Our work presents a robust, flexible, reproducible FL approach that enables large-scale collaboration across institutions to train LLMs. We propose a scalable deployment system called Photon to enable the investigation and development of this new training paradigm for LLM pre-training. We show that Photon can be used by organizations interested in collaborating with their private data sources and computational resources for pre-training LLMs with billions of parameters. This paradigm would mobilize more computational and data resources while matching or potentially exceeding centralized performance. We further show the effectiveness of the federated training scales with model size and present our approach for training billion-scale federated LLMs using limited resources. Thus far, we have used Photon to train LLM models to the size of 7B parameters and anticipate larger models being completed in the near future. Finally, we show that LLM training is highly resilient to the classical challenges of federated statistical and hardware heterogeneity. Furthermore, we show that convergence is robust to partial participation, opening the avenue for compute-efficient collaborative training. Photon will help data-rich actors to become the protagonists of LLMs pre-training instead of leaving the stage to compute-rich actors alone.

연구 동기 및 목표

• 중앙 집중식 LLM 사전 훈련에서의 증가하는 데이터 및 계산 자원의 병목 현상을 해결하기 위해, 활용도가 낮은 전 세계의 데이터와 계산 자원을 활용하고자 한다.
• 계산 능력이나 네트워크 연결성에 관계없이 데이터가 풍부한 기관들이 원시 데이터를 공유하지 않고도 LLM 사전 훈련에 기여할 수 있도록 하고자 한다.
• 연합 학습이 빌리언 파라미터 규모의 모델로도 확장 가능하며, 중앙 집중식 훈련과 경쟁 가능한 성능을 달성할 수 있음을 입증하고자 한다.
• 이질적인 클라이언트와 개인 정보 보호를 지원하는, 재현 가능한 오픈소스 시스템을 개발하여 연합 LLM 사전 훈련을 가능하게 하고자 한다.
• 더 큰 연합 모델이 더 작은 모델보다 수렴이 더 쉽게 이루어지는지 조사하고, 기존의 통념을 도전하고자 한다.

제안 방법

• 다양한 하드웨어와 데이터를 가진 여러 클라이언트 간의 분산 훈련을 조율하기 위해 개방형 FL 프레임워크인 Flower를 사용한다.
• 최적화의 안정성과 이질적인 노드 간의 클라이언트 기울기 정렬을 위해 모멘텀 기반 평균화를 사용하는 FedAvg를 적용한다.
• 통신 오버헤드를 줄이고 수렴 성능을 향상시키기 위해 로컬 SGD에 가중치 감소, 클리핑, 학습률 스케줄링을 적용한다.
• 재현 가능성을 보장하기 위해 투명한 초모수 및 설정을 포함한 훈련 레시피를 구현한다.
• 수렴 상태를 모니터링하고 과적합 여부를 탐지하기 위해 서버 측 모델 검증 및 퍼플렉서티 추적을 수행한다.
• 직접 데이터를 공유하지 않도록 하여 개인정보 보호를 설계 단계부터 고려한 모델 집계를 구현한다.

실험 결과

연구 질문

• RQ1 제한된 하드웨어와 이질적인 클라이언트 환경에서도 연합 학습이 중앙 집중식 훈련과 경쟁 가능한 LLM 사전 훈련 성능을 달성할 수 있는가?
• RQ2 기대와는 반대로, 연합 훈련의 확장성은 모델 크기가 증가함에 따라 향상되는가?
• RQ3 연합 최적화 과정에서 클라이언트 및 서버 모델 노름은 어떻게 변화하는가? 이는 수렴 동역학에 대해 어떤 통찰을 제공하는가?
• RQ4 FL에서의 모델 평균화가 훈련을 정규화하고 클라이언트 기울기의 노이즈를 얼마나 줄이는가?
• RQ5 비공개 데이터 소스를 중앙 집중화하지 않고도 LLM 사전 훈련에 효과적으로 활용할 수 있는가?

주요 결과

• 표준 클라우드 제공업체의 1~8개 GPU 노드만을 사용하여 이질적인 연합 환경에서 빌리언 스케일 LLM의 사전 훈련을 최초로 성공적으로 수행하였다.
• 연합 훈련은 검증 퍼플렉서티 측면에서 중앙 집중식 훈련과 동일하거나 이를 초월했으며, 서버 모델은 클라이언트의 피크를 소프트 상한선으로 제공하였다.
• 더 큰 연합 모델(예: 350M 파라미터)은 더 작은 모델보다 클라이언트 간에 공감대를 형성하는 데 더 쉽게 성공했으며, 예상과는 반대로 결과를 보였다.
• FedAvg의 의사 기울기 노름은 로컬 클라이언트 기울기보다 더 빨리 감쇠했으며, 이는 평균화 과정에서 상충되거나 노이즈가 많은 최적화 방향이 상쇄되었음을 시사한다.
• 클라이언트와 서버 모델 간의 상호작용을 분석한 결과, 서버 모델은 초기에 클라이언트 노름을 억제하다가 나중에 안정화되고 동일한 노름으로 수렴함을 확인했으며, 이는 성공적인 공감대 형성의 증거로 볼 수 있다.
• 이 시스템은 데이터가 풍부한 기관들이 거대한 컴퓨팅 인프라 없이도 연합 네트워크에 참여함으로써 LLM 사전 훈련의 핵심 참여자로 활동할 수 있음을 입증하였다.