[논문 리뷰] Revisiting Parameter Server in LLM Post-Training

Xinyi Wan, Penghui Qi외 4명|arXiv (Cornell University)|2026. 01. 27.

Cloud Computing and Resource Management인용 수 0

한 줄 요약

논문은 On-Demand Communication (ODC)를 도입한 분산 매개변수 서버와 같은 체계로, Fully Sharded Data Parallel(FSDP)의 레이어별 모음(collectives)을 포인트투포인트 전송으로 대체하여, 불균형한 LLM 사후 학습 워크로드(SFT 및 RL)에 대해 장치 활용도와 처리량을 최대 36%까지 향상시킵니다.

ABSTRACT

Modern data parallel (DP) training favors collective communication over parameter servers (PS) for its simplicity and efficiency under balanced workloads. However, the balanced workload assumption no longer holds in large language model (LLM) post-training due to the high variance in sequence lengths. Under imbalanced workloads, collective communication creates synchronization barriers, leading to under-utilization of devices with smaller workloads. This change in training dynamics calls for a revisit of the PS paradigm for its robustness to such imbalance. We propose extbf{On-Demand Communication (ODC)}, which adapts PS into Fully Sharded Data Parallel (FSDP) by replacing collective all-gather and reduce-scatter with direct point-to-point communication. Compared to FSDP, ODC reduces the synchronization barrier from once per layer to once per minibatch and decouples the workload on each device so that faster workers are not stalled. It also enables simpler and more effective load balancing at the minibatch level. Across diverse LLM post-training tasks, ODC consistently improves device utilization and training throughput, achieving up to a 36\% speedup over standard FSDP. These results demonstrate that ODC is a superior fit for the prevalent imbalanced workloads in LLM post-training. Our implementation of ODC and integration with FSDP is open-sourced at https://github.com/sail-sg/odc.

연구 동기 및 목표

• 시퀀스 길이의 변이가 동기화 장벽을 야기하는 상황에서 LLM 사후 학습의 워크로드 불균형에 대한 강건성의 필요성을 제시한다.
• 메모리 효율성을 해치지 않으면서 PS(매개변수 서버) 개념을 FSDP에 적용하기 위해 On-Demand Communication(ODC)을 제안한다.
• ODC가 SFT 및 RL 태스크에서 장치 활용도와 학습 처리량을 향상시킨다는 것을 입증한다.
• 노드 간 통신 오버헤드를 완화하기 위한 부하 분산 및 구현에 관한 실용적 가이드를 제공한다.

제안 방법

• FSDP에서 레이어별 all-gather 및 reduce-scatter를 직접적인 포인트투포인트 파라미터 가져오기 및 그래디언트 푸시로 대체한다.
• 레이어에서 미니배치 단위로 동기화를 완화하여 디바이스 진행 상황을 분리하되, 동기화 최적화의 의미를 유지한다.
• 서버와 워커 역할을 모든 디바이스에 걸쳐 공동 배치해 FSDP를 분산 매개변수 서버로 렌더링한다.
• 데이터 전송을 위해 RDMA 기반 인터페이스(CUDA IPC intra-node; NVSHMEM inter-node)와 Triton-Distributed 커널을 사용해 ODC를 구현한다.
• 컬렉티브를 대체하고 미니배치 경계에서 누적 그래디언트를 수집하여 FSDP와 ODC를 통합한다.
• 마이크로배치에서 메시지 포장을 이동시키고 미니배치 수준의 균형을 개선해 로드 밸런싱 전략을 제안한다.

Figure 1: Collective communications introduces per-layer synchronization barriers in FSDP.

실험 결과

연구 질문

• RQ1 ODC가 LLM 사후 학습에서 워크로드 불균형으로 인한 동기화 장벽과 유휴 시간을 감소시킬 수 있는가?
• RQ2 PS와 같은 분리된 통신을 FSDP에 통합하더라도 불균형한 워크로드에서 메모리 효율성을 유지하면서 처리량이 향상되는가?
• RQ3 장기 컨텍스트 LLM 학습에서 미니배치 단위로 ODC를 보완하는 로드밸런싱 전략은 무엇인가?
• RQ4 ODC는 1.5B에서 32B 매개변수의 모델 규모에 걸친 감독학습 미세조정(SFT) 및 강화학습(RL) 태스크에서 어떻게 성능을 발휘하는가?
• RQ5 노드 간 ODC 통신의 한계와 잠재적 완화 방법은 무엇인가?

주요 결과

• ODC는 SFT 및 RL 태스크 전반에서 디바이스 활용도와 엔드투엔드 처리량을 일관되게 향상시킨다.
• ODC는 롱-context SFT 시나리오에서 표준 FSDP 대비 최대 36%의 속도 향상을 달성한다.
• 전통적 FSDP에서 긴 시퀀스의 감독학습에서 워크로드 불균형으로 인한 유휴 시간이 최대 50%에 이를 수 있다.
• FSDP를 주문형 포인트투포인트 전송이 가능한 분산 PS로 재구성하면 지연띠(Sstragglers)를 완화하고 마이크로벤치마크의 동기화를 완화한다.
• LB-Mini 및 LB-Micro 로드밸런싱 변형은 미니배치 수준의 균형을 효과적으로 가능하게 하며 RL 및 SFT 설정에서 종종 기준치를 능가한다.
• ODC는 단일 노드 내의 집단 통신 방법과 경쟁력을 유지하지만, 노드 간 통신 오버헤드를 보인다. 이는 하이브리드 샤딩 및 연산과의 겹치기 같은 설계 선택으로 완화될 수 있다.

Figure 2: On-demand communications relaxes the synchronization barriers to minibatch end.