LLM 사후 학습을 위한 파라미터 서버 재조명: 온디맨드 커뮤니케이션(ODC)

초록

LLM 사후 학습에서는 시퀀스 길이 차이로 인한 작업 부하 불균형이 심해 기존의 집합 통신 기반 데이터 병렬(FSDP) 방식이 병목이 된다. 저자는 파라미터 서버 개념을 현대적인 Fully Sharded Data Parallel(FSDP)과 결합해, all‑gather와 reduce‑scatter를 포인트‑투‑포인트 온디맨드 통신으로 교체한 ODC 방식을 제안한다. 이를 통해 동기화 지점을 레이어당 한 번에서 미니배치당 한 번으로 줄이고, 각 디바이스가 독립적으로 진행하도록 하여 스트래거 현상을 크게 완화한다. 실험 결과, 다양한 LLM 사후 학습 작업에서 최대 36 %의 속도 향상을 달성했으며, 구현 코드는 공개되어 있다.

상세 분석

본 논문은 LLM 사후 학습 단계에서 발생하는 시퀀스 길이 변동성으로 인한 워크로드 불균형이 기존의 집합 통신 기반 데이터 병렬(FSDP) 방식에 치명적인 병목을 만든다는 점을 정확히 지적한다. FSDP는 파라미터와 그래디언트를 디바이스에 샤딩하고, 레이어별로 all‑gather와 reduce‑scatter를 수행해 전체 파라미터를 재구성하고 그래디언트를 집계한다. 이러한 레이어 단위 동기화는 모든 디바이스가 가장 느린 디바이스에 맞춰야 함을 의미하며, 특히 긴 시퀀스가 포함된 마이크로배치가 일부 디바이스에 집중될 경우, 다른 디바이스는 최대 50 %까지 유휴 상태에 빠진다. 기존 연구들은 마이크로배치 수준에서 시퀀스 패킹을 통해 부하를 완화하려 했지만, 메모리 제한으로 인해 패킹 효율이 한계에 부딪히고, 오히려 동기화 포인트가 늘어나면서 병목이 심화된다.

이를 해결하기 위해 저자는 파라미터 서버(PS) 아키텍처의 핵심 아이디어—서버와 워커를 분리해 비동기적으로 파라미터를 조회하고 그래디언트를 전송한다—를 현대적인 샤딩 DP와 결합한다. 구체적으로, all‑gather를 “필요한 파라미터 샤드만 개별 요청(gather)”으로, reduce‑scatter를 “그래디언트 샤드에 직접 scatter‑accumulate” 형태의 포인트‑투‑포인트 전송으로 분해한다. 이렇게 하면 각 디바이스는 자신이 준비된 시점에 바로 파라미터를 받아 연산을 시작하고, 그래디언트를 즉시 전송할 수 있어 레이어 단위 동기화가 사라진다. 동기화는 미니배치 전체가 끝난 뒤 한 번만 수행되므로, 스트래거 디바이스가 전체 파이프라인을 지연시키는 현상이 크게 감소한다.

통신 구현 측면에서 저자는 NCCL이나 MPI와 달리 RDMA 기반 CUDA IPC와 NVSHMEM을 활용한다. RDMA는 수신 측이 별도의 서버 스레드 없이도 메모리 영역에 직접 쓰기를 가능하게 하여, “요청이 오면 바로 처리”하는 온디맨드 특성을 보장한다. 또한 Triton‑Distributed를 이용해 파이썬 레벨에서 RDMA 호출을 래핑함으로써 개발 복잡성을 크게 낮추었다. 이러한 설계는 기존 FSDP와 동일한 메모리 샤딩 구조를 유지하면서도, 통신 오버헤드와 동기화 비용을 최소화한다.

실험에서는 두 가지 대표적인 LLM 사후 학습 작업, 즉 Supervised Fine‑Tuning(SFT)과 Reinforcement Learning(RL) 파이프라인을 대상으로 평가하였다. 다양한 모델 크기(7B70B)와 시퀀스 길이(5124096)에서 ODC는 평균 20 %~36 %의 처리량 향상을 보였으며, 특히 긴 시퀀스와 높은 변동성을 가진 데이터셋에서 효과가 두드러졌다. 장비 활용도는 기존 FSDP 대비 10 %~25 % 상승했으며, 이는 스트래거에 의한 유휴 시간이 크게 감소했음을 의미한다. 또한, ODC는 마이크로배치 레벨의 복잡한 패킹 알고리즘을 단순화시켜, 미니배치 전체를 기준으로 부하를 균등하게 분배하는 전략만으로도 충분히 높은 효율을 달성한다는 점을 강조한다.

결과적으로, 본 논문은 “파라미터 서버는 구식이다”는 기존 인식을 뒤집고, 현대적인 샤딩 DP와 결합했을 때 PS‑스타일 온디맨드 통신이 LLM 사후 학습의 불균형 워크로드에 최적화된 솔루션이 될 수 있음을 입증한다. 코드와 구현은 공개되어 있어, 향후 대규모 LLM 파인‑튜닝 및 RL 연구에 바로 적용 가능하다.