ParaBlock: Communication-Computation Parallel Block Coordinate Federated Learning for Large Language Models

기술 요약: ParaBlock: 대규모 언어 모델을 위한 통신-연산 병렬 블록 좌표 연합 학습

문제 정의
연합 학습(Federated Learning, FL)은 분산된 데이터에 걸쳐 모델을 훈련하기 위한 개인정보 보호 패러다임으로 부상했습니다. 최근에는 대규모 모델을 훈련하기 위해 클라이언트가 전체 모델이 아닌 모델 파라미터의 일부(블록)만을 로컬에서 업데이트하는 연합 블록 좌표 하강(Federated Block Coordinate Descent, BCD) 방식이 주목받고 있습니다. 그러나 대규모 언어 모델(LLM) 시대에 들어서면서, 단일 블록조차 수천만에서 수억 개의 파라미터를 포함할 수 있게 되었습니다. 이러한 규모는 심각한 병목 현상을 야기합니다. 즉, 클라이언트와 서버 간에 이러한 거대 블록을 전송하는 데 필요한 통신 지연 시간이 매우 커지며, 특히 자원이 제한된 엣지 클라이언트에게는 치명적입니다.

전통적인 연합 BCD 방식은 단일 스레드 설계(통신-연산-통신)로 작동하여, 클라이언트는 다음 연산 단계를 시작하기 전에 모델 전송이 완료될 때까지 기다려야 합니다. 이러한 통신 시간은 소규모 모델에서는 무시할 만한 수준이었으나, 이제는 LLM의 연합 미세 조정(Fine-tuning)에 있어 확장성과 효율성을 저해하며 클라이언트의 상당한 유휴 시간을 발생시키는 주요 원인이 되고 있습니다.

방법론: ParaBlock
이러한 지연 문제를 해결하기 위해, 저자들은 통신-연산 병렬 구조를 도입한 새로운 연합 BCD 방식인 ParaBlock을 제안합니다. 표준적인 단일 스레드 접근 방식과 달리, ParaBlock은 로컬 클라이언트에 두 개의 병렬 스레드를 구축합니다:

1. 연산 스레드(Computation Thread): 로컬 데이터를 사용하여 할당된 블록의 로컬 블록 좌표 업데이트(미세 조정)를 수행합니다.
2. 통신 스레드(Communication Thread): 이전 라운드의 로컬 업데이트를 서버로 보내고, 집계된 글로벌 업데이트를 받는 동기화 작업을 수행합니다.

핵심 메커니즘:

• 오버래핑(Overlapping): 클라이언트는 서버와 통신하는 동시에 로컬 연산을 실행합니다. 이를 통해 네트워크 전송을 기다리며 발생하는 일반적인 유휴 시간을 제거합니다.
• 스티일니스(Staleness) 및 보정: 통신 스레드가 한 라운드 뒤처지기 때문에(클라이언트는 $t$ 라운드의 연산을 수행하는 동안 $t-1$ 라운드의 글로벌 업데이트를 받음), 로컬 모델 파라미터가 글로벌 모델과 불일치할 수 있습니다. ParaBlock은 보정 메커니즘을 통해 이를 해결합니다. 클라이언트가 다음 라운드를 위해 로컬 모델을 업데이트할 때, 현재 블록에 새로운 로컬 업데이트를 적용할 뿐만 아니라, 수신된 글로벌 업데이트($\Delta_{t-1}$)를 사용하여 이전 블록의 파라미터를 기존의 오래된(stale) 로컬 업데이트($\Delta^i_{t-1}$) 대신 교체함으로써 보정합니다.
• 재귀적 일관성(Recursive Consistency): 이론적 분석에 따르면, 이 보정 과정을 통해 한 번의 글로벌 라운드가 지나면 이전에 업데이트된 블록에 대한 로컬 모델의 파라미터가 글로벌 모델의 파라미터와 정확히 일치하게 되어, 한 라운드 지연의 영향을 효과적으로 완화함을 입증하였습니다.

주요 기여

1. 새로운 아키텍처: 기존의 단일 스레드 연합 학습 설계를 대체하고, LLM 연합 미세 조정의 통신 지연을 줄이기 위해 특별히 설계된 두 개의 스레드 병렬 체계인 ParaBlock을 제안했습니다.
2. 이론적 보장: 저자들은 비볼록(non-convex) 목적 함수에 대한 엄격한 수렴 분석을 제공합니다. ParaBlock이 $O(1/\sqrt{T})$의 수렴 속도를 달enc함이 증명되었으며, 이는 표준 연합 블록 좌표 하강 방식과 일치합니다. 이는 병렬화를 통한 효율성 향상이 이론적 수렴 성능을 저해하지 않음을 확인시켜 줍니다.
3. 실증적 검증: Llama 3-8B 및 Llama 3.2-3B 모델을 대상으로 일반적인 지시 이행(Alpaca-GPT4) 및 수학적 추론(MathInstruct/GSM8K) 작업에 대해 광범한 실험을 수행했습니다.

결과

• 성능: ParaBlock은 강력한 다운스트림 성능을 유지합니다. 표준 연합 BCD 베이스라인(예: FedBCD)과 대등하거나 더 나은 성능을 보였으며, 지시 이행 및 수학 추론 작업에서 여러 PEFT 기반 연합 학습 방법(예: FedIT, FLoRA)보다 유의미하게 높은 성능을 기록했습니다.
• 효율성: 이 방식은 실제 실행 시간(wall-clock runtime)을 크게 단축합니다. 낮은 네트워크 대역폭(50M/s) 및 작은 배치 크기 시나리오에서, ParaBlock은 FedCyBGD와 같은 블록 좌표 베이스라인 대비 실행 시간을 30% 이상 단축했습니다. LoRA 기반 방식인 FedIT와 비교했을 때, 대부분의 설정에서 약 40%의 실행 시간 단축을 달성했습니다.
• 강건성(Robustness): 절제 연구(Ablation studies)를 통해 이 방법이 이질적인 데이터 분포(non-IID) 및 다양한 모델 아키텍처(Qwen-2.5 포함)에서도 효과적임을 보여주었습니다. 또한 크로스-실로(cross-silo) 환경에서의 부분 참여(partial participation) 설정에 대한 호환성도 입증되었습니다.

의의
본 논문은 ParaBlock이 LLM을 위한 연합 BCD에 내재된 통신 병목 현상에 대한 단순하면서도 효과적인 솔루션을 제공한다고 주장합니다. 통신과 연산을 오버래핑함으로써, ParaBlock은 수렴 보장이나 모델 효용성을 희생하지 않고도 자원이 제한된 클라이언트에서 대규모 모델을 위한 연합 미세 조정을 확장 가능하게 배포할 수 있도록 합니다. 이 연구는 차세대 연합 LLM 훈련을 위해 양자화(quantization)나 가지치기(pruning)와 같은 기존 기술과는 구별되며 상호 보완적인 전략으로서, 아키텍처적 병렬화를 통해 통신 지연을 줄이는 것이 필수적이고 유효한 전략임을 강조합니다.